Fix for OOM Errors during Ultrachat200k Finetuning

[Satrat]

The following SparseZoo recipes were causing OOM errors even with FSDP:

• "zoo:llama2-7b-ultrachat200k_llama2_pretrain-pruned40"
• "zoo:llama2-7b-ultrachat200k_llama2_pretrain-pruned40_quantized"

The fix was to offload weights to CPU when gathering FSDP params during modifier initialization and finalization

Test

Tested on 6 48GB GPUs. Example only runs a few training samples for ease of testing.

from sparseml.transformers import compress, SparseAutoModelForCausalLM, SparseAutoTokenizer

model = SparseAutoModelForCausalLM.from_pretrained("zoo:llama2-7b-ultrachat200k_llama2_pretrain-base")
teacher = SparseAutoModelForCausalLM.from_pretrained("zoo:llama2-7b-ultrachat200k_llama2_pretrain-base")
tokenizer = SparseAutoTokenizer.from_pretrained("zoo:llama2-7b-ultrachat200k_llama2_pretrain-base")
dataset="open_platypus"
MODEL_STUB="zoo:llama2-7b-ultrachat200k_llama2_pretrain-pruned40"

compress(
    model=model,
    distill_teacher=teacher,
    tokenizer=tokenizer,
    dataset=dataset,
    recipe=MODEL_STUB,
    output_dir="./output",
    gradient_checkpointing = True,
    num_train_epochs=0.02
)

Run with FSDP: accelerate launch --config_file fsdp_config.yaml test.py