[BUG] LISA: same loss regardless of lisa_activated_layers

Describe the bug
I think there might be something wrong with the current LISA implementation. There is no difference in training loss, no matter how many layers are active.

Not using LMFlow but HF Trainer with DynamicLayerActivationCallback from https://github.com/OptimalScale/LMFlow/blob/main/src/lmflow/pipeline/finetuner.py

To Reproduce

class DynamicLayerActivationCallback(TrainerCallback):
    def __init__(self, n_layers, interval_steps, model):
        super().__init__()
        self.n_layers = n_layers
        self.interval_steps = interval_steps
        self.model = model
        # Determine the way to access layers based on the model type
        if self.model.__class__.__name__ == 'LlamaForCausalLM':
            self.layers_attribute = 'model.model.layers'  # Layer access path for LlamaForCausalLM
        else:
            self.layers_attribute = 'model.transformer.h'  # General access path
        self.total_layers = len(eval('self.' + self.layers_attribute))  # Dynamically execute to get the number of layers

        # Freeze all layers upon initialization
        self.freeze_all_layers()
        self.active_layers_indices = []

    def freeze_all_layers(self):
        layers = eval('self.' + self.layers_attribute)  # Dynamically execute to get layers
        for layer in layers:
            for param in layer.parameters():
                param.requires_grad = False

    def on_step_begin(self, args, state, control, **kwargs):
        # Check if it's time to switch active layers, including at step 0
        if state.global_step % self.interval_steps == 0 or state.global_step == 1:
            self.switch_active_layers()

    def switch_active_layers(self):
        # First, disable gradients for all layers
        self.freeze_all_layers()

        # Randomly select n_layers to activate
        layers = eval('self.' + self.layers_attribute)  # Re-fetch layer references
        self.active_layers_indices = np.random.choice(range(self.total_layers), self.n_layers, replace=False)
        print(f"Activating layers at indices: {self.active_layers_indices} for the next steps.")

        # Enable gradients only for the selected layers
        for idx in self.active_layers_indices:
            for param in layers[idx].parameters():
                param.requires_grad = True

# Instantiate the callback
dynamic_layer_activation_callback = DynamicLayerActivationCallback(
    n_layers = lisa_activated_layers,                     # Number of layers to activate
    interval_steps = lisa_interval_steps,               # Step interval to update active layers
    model = model
)

trainer.add_callback(dynamic_layer_activation_callback)

model llama2-7b

Expected behavior

• different loss for different lisa_activated_layers
• same loss (and VRAM usage) for lisa_activated_layers==32 and full finetune (without LISA) - loss curves are different, they diverge after a few steps

Screenshots

Setup
2x 3090

torch==2.2.1
transformers==4.39.2
Python 3.10.12

Start with DynamicLayerActivationCallback in src/lmflow/pipeline/finetuner.py, using the provided Hugging Face Trainer reproduction with Llama2-7b. Compare training loss and VRAM across different lisa_activated_layers values, including 32 and full fine-tuning. Done means the active-layer count changes the observed behavior and the 32-layer and full-fine-tuning results match the expected distinction.