MVP for Alternating Flow

[Satrat]

Initial implementation for alternating between oneshot and finetuning stages. This branch is based on two active PRs, they should be merged first:

• Support for Stacking Recipes #1897
• Enable One-Shot Launch from Finetuning Script #1907

Testing

test_multi_recipe.yaml

test_oneshot_stage:
  obcq_modifiers:
    SparseGPTModifier:
      sparsity: 0.5
      block_size: 128
      sequential_update: False
      quantize: False
      percdamp: 0.01
      prunen: 0
      prunem: 0
      targets: [
        "re:model.layers.\\d+$"
      ]
      target_ids: ["attention_mask", "position_ids"]  
test_finetune_stage:
  pruning_modifiers:
    ConstantPruningModifier:
      targets: [
        "re:.*self_attn.q_proj",
        "re:.*self_attn.k_proj",
        "re:.*self_attn.v_proj",
        "re:.*self_attn.o_proj",
        "re:.*mlp.gate_proj",
        "re:.*mlp.up_proj"
      ]
      start: 0
test_second_oneshot_stage:
  obcq_modifiers:
    SparseGPTModifier:
      sparsity: 0.7
      block_size: 128
      sequential_update: False
      quantize: False
      percdamp: 0.01
      prunen: 0
      prunem: 0
      targets: [
        "re:model.layers.\\d+$"
      ]
      target_ids: ["attention_mask", "position_ids"]  
test_second_finetune_stage:
  pruning_modifiers:
    ConstantPruningModifier:
      targets: [
        "re:.*self_attn.q_proj",
        "re:.*self_attn.k_proj",
        "re:.*self_attn.v_proj",
        "re:.*self_attn.o_proj",
        "re:.*mlp.gate_proj",
        "re:.*mlp.up_proj"
      ]
      start: 0
test_quantization_oneshot_stage:
  obcq_modifiers:
    QuantizationModifier:
      ignore:
        - LlamaRotaryEmbedding
        - LlamaRMSNorm
        - SiLUActivation
        - model.layers.0.mlp.down_proj
        - model.layers.1.mlp.down_proj
        - model.layers.2.mlp.down_proj
        - model.layers.3.mlp.down_proj
        - model.layers.4.mlp.down_proj
        - model.layers.5.mlp.down_proj
      post_oneshot_calibration: False
      scheme_overrides:
        Embedding:
          input_activations: null
          weights:
            num_bits: 8
            symmetric: False

Test script:

def run():
    from sparseml.transformers.finetune.text_generation import run_general
    
    model = "Xenova/llama2.c-stories15M"
    dataset_name = "open_platypus"
    concatenate_data = False
    run_stages = True
    output_dir = "./output_oneshot"
    overwrite_output_dir = True
    recipe = "test_multi_recipe.yaml"
    splits = {
        "calibration": "train[:50%]",
        "train": "train[50%:]"
    }

    run_general(
        model_name_or_path=model,
        dataset_name=dataset_name,
        run_stages=run_stages,
        output_dir=output_dir,
        overwrite_output_dir=overwrite_output_dir,
        recipe=recipe,
        concatenate_data = concatenate_data,
        splits = splits
    )

if __name__ == "__main__":
    run()

Known Issues/ Shortcomings

• FSDP hasn't been tested yet
• Training checkpoints getting overwritten during subsequent finetuning runs
• No way to specify different numbers of epochs for each finetune stage
• No way to specify different datset splits for different finetuning stages
• Checkpoint loading between stages not implemented
• Output recipe doesn't indicate what stages have been run and what hasn't
• No unit or integration tests!

[rahul-tuli]

LGTM! Good tests

[rahul-tuli]

LGTM! Good tests