test_ft.py

import argparse
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
import transformers
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
from transformers.trainer_callback import TrainerCallback
import os
import random
import subprocess
from huggingface_hub import HfApi, hf_hub_download

def generate_mmlu_slurm(model_path, hub_model_id, output_dir, num_gpus=1):
    model_short_name = model_path.split('/')[-1]
    filename = f"run_mmlu_{model_short_name}.sbatch"
    
    port = random.randint(10000, 65535)
    
    content = f"""#!/bin/bash
#SBATCH --nodes=1
#SBATCH --gpus-per-node={num_gpus}
#SBATCH --time=24:00:00
#SBATCH --job-name={port}_mmlu_{model_short_name}
#SBATCH --mail-user=mailto:[email protected]
#SBATCH --mail-type=ALL

source /opt/rh/devtoolset-10/enable
source /data/davis_brown/miniconda3/bin/activate
conda init
conda activate quip

CUDA_VISIBLE_DEVICES=0 lm_eval \\
    --model hf \\
    --model_args pretrained={model_path},parallelize=True,peft={hub_model_id} \\
    --tasks mmlu \\
    --device cuda:0 \\
    --batch_size 8 \\
    --output_path={output_dir}/{hub_model_id} \\
    --num_fewshot 5
"""

    with open(filename, 'w') as f:
        f.write(content)
    
    print(f"Generated MMLU evaluation SLURM script: {filename}")
    return filename

def launch_mmlu_evaluation(model_path, hub_model_id, output_dir):
    slurm_script = generate_mmlu_slurm(model_path, hub_model_id, output_dir)
    try:
        subprocess.run(["sbatch", slurm_script], check=True)
        print(f"Submitted MMLU evaluation job: {slurm_script}")
    except subprocess.CalledProcessError as e:
        print(f"Failed to submit MMLU evaluation job: {e}")

# Custom callback to push to Hub
class PushToHubCallback(TrainerCallback):
    def __init__(self, trainer, push_frequency):
        self.trainer = trainer
        self.push_frequency = push_frequency

    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step % self.push_frequency == 0:
            self.trainer.save_model()
            self.trainer.push_to_hub(
                commit_message=f"Training in progress - Step {state.global_step}"
            )
def main(args):
    if args.only_mmlu:
        launch_mmlu_evaluation(args.model_id, args.hub_model_id, args.output_dir)
        return

    model_id = args.model_id
    output_dir = args.output_dir
    hub_model_id = args.hub_model_id

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto", low_cpu_mem_usage=True)

    target_modules = ['q_proj','k_proj','v_proj','o_proj','gate_proj','down_proj','up_proj', ]# 'lm_head']

    config = LoraConfig(
        r=args.lora_rank,
        lora_alpha=args.lora_rank,
        target_modules=target_modules,
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
        use_rslora=True
    )

    model = get_peft_model(model, config)
    model.print_trainable_parameters()
    model.enable_input_require_grads()

    # data = load_dataset("togethercomputer/RedPajama-Data-1T-Sample")
    data = load_dataset("open-web-math/open-web-math")


    max_seq_length = args.max_seq_length
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.model_max_length = max_seq_length

    def preprocess_function(examples):
        return tokenizer(examples["text"], truncation=True, max_length=max_seq_length, padding="max_length")

    processed_dataset = data["train"].map(preprocess_function, batched=True)

    tokenizer.pad_token = tokenizer.eos_token
    torch.cuda.empty_cache()
    trainer = transformers.Trainer(
        model=model,
        train_dataset=processed_dataset,
        args=TrainingArguments(
            per_device_train_batch_size=args.batch_size,
            gradient_accumulation_steps=args.gradient_accumulation_steps,
            gradient_checkpointing=True,
            warmup_steps=200,
            max_steps=args.max_steps,
            learning_rate=2e-4,
            bf16=True,
            logging_steps=25,
            output_dir=output_dir,
            optim="adamw_bnb_8bit",
            logging_first_step=True,
            push_to_hub=True,
            hub_model_id=hub_model_id,
        ),
        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
    )
    model.config.use_cache = False

    push_frequency = 100
    trainer.add_callback(PushToHubCallback(trainer, push_frequency, hub_model_id))

    trainer.train()

    final_commit_hash = trainer.push_to_hub("Training complete")
    print(f"Training complete. Final commit hash: {final_commit_hash}")

    # MMLU Evaluation
    if args.run_mmlu:
        launch_mmlu_evaluation(model_id, hub_model_id, output_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fine-tune a language model and/or run MMLU evaluation")
    parser.add_argument("--model_id", type=str, default="ISTA-DASLab/Meta-Llama-3-8B-Instruct-AQLM-2Bit-1x16", 
                        help="Model ID to fine-tune or evaluate")
    parser.add_argument("--max_seq_length", type=int, default=2048, help="Maximum sequence length")
    parser.add_argument("--output_dir", type=str, required=True, help="Output directory for checkpoints and results")
    parser.add_argument("--hub_model_id", type=str, 
                        default="davisrbr/ISTA-DASLab-Meta-Llama-3-8B-Instruct-AQLM-2Bit-1x16-hf-100000_r8_cont", 
                        help="Hub model ID for pushing or LoRA weights")
    parser.add_argument("--batch_size", type=int, default=1, help="Per-device batch size")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=8, help="Gradient accumulation steps")
    parser.add_argument("--max_steps", type=int, default=50000, help="Maximum number of training steps")
    parser.add_argument("--run_mmlu", action="store_true", help="Run MMLU evaluation after training")
    parser.add_argument("--lora_rank", type=int, default=8, help="Rank of LoRA adaptation")
    parser.add_argument("--only_mmlu", action="store_true", help="Only run MMLU evaluation without training")
    parser.add_argument("--launch_slurm", action="store_true", help="Launch the entire script as a SLURM job")
    parser.add_argument("--num_gpus", type=int, default=4, help="Number of GPUs to use for training")
    parser.add_argument("--commit_hash", type=str, help="Specific commit hash to evaluate (for MMLU only)")
    
    args = parser.parse_args()
    main(args)