Add lisa support for gpt2

- add preliminary of support for mixtral
- add --lisa_layers_attribute for customizable models

src/lmflow/args.py

@@ -591,6 +591,12 @@ class FinetunerArguments(TrainingArguments):
             "help": "the number of steps in each freezing interval of LISA, i.e. the selected unfreezed layers are randomly switched every {lisa_interval_steps} steps."
         }
     )
+    lisa_layers_attribute: int = field(
+        default="model.model.layers",
+        metadata={
+            "help": "where the layer attribute stores, e.g. model.model.layers"
+        }
+    )
 
 
 @dataclass

src/lmflow/pipeline/finetuner.py

@@ -311,10 +311,14 @@ def __init__(self, n_layers, interval_steps, model):
                         self.layers_attribute = 'model.model.layers'  # Layer access path for Qwen model
                     elif self.model.__class__.__name__ == 'MistralForCausalLM':
                         self.layers_attribute = 'model.model.layers'
+                    elif self.model.__class__.__name__ == 'MixtralForCausalLM':
+                        self.layers_attribute = 'model.model.layers'
                     elif self.model.__class__.__name__ == 'GemmaForCausalLM':
                         self.layers_attribute = 'model.model.layers'
+                    elif self.model.__class__.__name__ == 'GPT2LMHeadModel':
+                        self.layers_attribute = 'model.transformer.h'
                     else:
-                        self.layers_attribute = 'model.transformer.h'  # General access path
+                        self.layers_attribute = training_args.lisa_layers_attribute
                     self.total_layers = len(eval('self.' + self.layers_attribute))  # Dynamically execute to get the number of layers
 
                     self.active_layers_indices = []