train_qlora.py

import os
import sys
import json
import torch
import argparse
from datetime import datetime
from pathlib import Path
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

from datasets import load_dataset, Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq,
)
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    TaskType,
)
from tqdm import tqdm

class Config:
    base_model = "Qwen/Qwen2.5-3B-Instruct"
    train_file = "data/final/train.jsonl"
    val_file = "data/final/validation.jsonl"
    output_dir = "outputs/ml-slm-qwen3b"
    lora_r = 64
    lora_alpha = 128
    lora_dropout = 0.05
    lora_target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ]
    load_in_4bit = True
    bnb_4bit_compute_dtype = torch.bfloat16
    bnb_4bit_quant_type = "nf4"
    bnb_4bit_use_double_quant = True
    num_epochs = 3
    per_device_train_batch_size = 1
    gradient_accumulation_steps = 16
    learning_rate = 2e-4
    lr_scheduler_type = "cosine"
    warmup_ratio = 0.03
    weight_decay = 0.01
    max_grad_norm = 0.3
    max_seq_length = 1024
    gradient_checkpointing = True
    optim = "paged_adamw_8bit"
    save_strategy = "steps"
    save_steps = 500
    save_total_limit = 3
    logging_steps = 50
    report_to = "none"
    eval_strategy = "steps"
    eval_steps = 500
    seed = 42
    bf16 = True
    tf32 = True

SYSTEM_PROMPT = """You are an expert ML/DL teaching assistant. Your role is to:

1. Explain machine learning and deep learning concepts clearly and accurately
2. Use intuitive analogies and simple language when possible
3. Provide the mathematical intuition behind algorithms
4. Compare and contrast different methods when relevant
5. Focus on conceptual understanding, not code implementation

IMPORTANT RULES:
- Do NOT write code or code examples
- Do NOT make up terms or concepts that don't exist
- If you're unsure about something, say so
- Keep explanations grounded in established ML/DL theory
- Be concise but thorough"""


def get_prompt_length(sample: dict, tokenizer) -> int:
    """Get the token length of the prompt (everything before the assistant response)."""
    instruction = sample.get('instruction', '')
    input_text = sample.get('input', '')
    if input_text:
        user_content = f"{instruction}\n\n{input_text}"
    else:
        user_content = instruction

    prompt_messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]
    prompt_text = tokenizer.apply_chat_template(
        prompt_messages,
        tokenize=False,
        add_generation_prompt=True
    )
    prompt_ids = tokenizer(prompt_text, add_special_tokens=False)["input_ids"]
    return len(prompt_ids)


def format_instruction(sample: dict, tokenizer) -> str:
    instruction = sample.get('instruction', '')
    input_text = sample.get('input', '')
    output = sample.get('output', '')
    if input_text:
        user_content = f"{instruction}\n\n{input_text}"
    else:
        user_content = instruction

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
        {"role": "assistant", "content": output}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False
    )
    return text


def tokenize_function(examples, tokenizer, max_length):
    all_input_ids = []
    all_attention_mask = []
    all_labels = []

    for i in range(len(examples['instruction'])):
        sample = {
            'instruction': examples['instruction'][i],
            'input': examples['input'][i] if examples['input'][i] else '',
            'output': examples['output'][i],
        }
        text = format_instruction(sample, tokenizer)
        prompt_len = get_prompt_length(sample, tokenizer)

        tokenized = tokenizer(
            text,
            truncation=True,
            max_length=max_length,
            padding=False,
            return_tensors=None,
        )

        input_ids = tokenized["input_ids"]
        attention_mask = tokenized["attention_mask"]

        # Mask prompt tokens in labels with -100 so loss is only on the response
        labels = list(input_ids)
        for j in range(min(prompt_len, len(labels))):
            labels[j] = -100

        # Skip samples where the prompt fills the entire sequence after truncation,
        # leaving no response tokens to compute loss on (would produce NaN loss)
        if all(l == -100 for l in labels):
            continue

        all_input_ids.append(input_ids)
        all_attention_mask.append(attention_mask)
        all_labels.append(labels)

    return {
        "input_ids": all_input_ids,
        "attention_mask": all_attention_mask,
        "labels": all_labels,
    }

def load_and_prepare_data(config, tokenizer):
    print("\nLoading dataset...")
    train_dataset = load_dataset('json', data_files=config.train_file, split='train')
    val_dataset = load_dataset('json', data_files=config.val_file, split='train')
    print(f"   Train samples: {len(train_dataset):,}")
    print(f"   Val samples: {len(val_dataset):,}")
    print("\nTokenizing dataset...")
    
    train_dataset = train_dataset.map(
        lambda x: tokenize_function(x, tokenizer, config.max_seq_length),
        batched=True,
        batch_size=1000,
        remove_columns=train_dataset.column_names,
        desc="Tokenizing train",
    )
    
    val_dataset = val_dataset.map(
        lambda x: tokenize_function(x, tokenizer, config.max_seq_length),
        batched=True,
        batch_size=1000,
        remove_columns=val_dataset.column_names,
        desc="Tokenizing val",
    )
    return train_dataset, val_dataset

def setup_model_and_tokenizer(config):
    print("\nLoading model and tokenizer...")
    print(f"   Base model: {config.base_model}")
    
    tokenizer = AutoTokenizer.from_pretrained(
        config.base_model,
        trust_remote_code=True,
        padding_side="right",
    )
    
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=config.load_in_4bit,
        bnb_4bit_compute_dtype=config.bnb_4bit_compute_dtype,
        bnb_4bit_quant_type=config.bnb_4bit_quant_type,
        bnb_4bit_use_double_quant=config.bnb_4bit_use_double_quant,
    )
    
    print("   Loading model in 4-bit...")
    model = AutoModelForCausalLM.from_pretrained(
        config.base_model,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        dtype=torch.bfloat16,
        attn_implementation="sdpa",
    )
    
    model = prepare_model_for_kbit_training(
        model,
        use_gradient_checkpointing=config.gradient_checkpointing
    )
    
    lora_config = LoraConfig(
        r=config.lora_r,
        lora_alpha=config.lora_alpha,
        lora_dropout=config.lora_dropout,
        target_modules=config.lora_target_modules,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )
    print("   Applying LoRA...")
    model = get_peft_model(model, lora_config)
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"   Trainable parameters: {trainable_params:,} ({100 * trainable_params / total_params:.2f}%)")
    return model, tokenizer

def train(config, resume_from_checkpoint=False):
    print("=" * 70)
    print("ML-SLM QLoRA TRAINING")
    print("=" * 70)
    print(f"\nStarted at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Output directory: {config.output_dir}")
    
    if not torch.cuda.is_available():
        print("CUDA not available!")
        sys.exit(1)
    
    print(f"\nGPU: {torch.cuda.get_device_name(0)}")
    print(f"   VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    
    model, tokenizer = setup_model_and_tokenizer(config)
    train_dataset, val_dataset = load_and_prepare_data(config, tokenizer)
    data_collator = DataCollatorForSeq2Seq(
        tokenizer=tokenizer,
        model=model,
        padding=True,
        pad_to_multiple_of=8,
    )
    training_args = TrainingArguments(
        output_dir=config.output_dir,
        num_train_epochs=config.num_epochs,
        per_device_train_batch_size=config.per_device_train_batch_size,
        per_device_eval_batch_size=config.per_device_train_batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        learning_rate=config.learning_rate,
        lr_scheduler_type=config.lr_scheduler_type,
        warmup_ratio=config.warmup_ratio,
        weight_decay=config.weight_decay,
        max_grad_norm=config.max_grad_norm,
        optim=config.optim,
        bf16=config.bf16,
        tf32=config.tf32,
        gradient_checkpointing=config.gradient_checkpointing,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        save_strategy=config.save_strategy,
        save_steps=config.save_steps,
        save_total_limit=config.save_total_limit,
        eval_strategy=config.eval_strategy,
        eval_steps=config.eval_steps,
        logging_steps=config.logging_steps,
        report_to=config.report_to,
        seed=config.seed,
        dataloader_num_workers=2,
        remove_unused_columns=False,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
    )
    total_steps = len(train_dataset) // (config.per_device_train_batch_size * config.gradient_accumulation_steps) * config.num_epochs
    
    print(f"\nTraining Configuration:")
    print(f"   Epochs: {config.num_epochs}")
    print(f"   Batch size: {config.per_device_train_batch_size}")
    print(f"   Gradient accumulation: {config.gradient_accumulation_steps}")
    print(f"   Effective batch size: {config.per_device_train_batch_size * config.gradient_accumulation_steps}")
    print(f"   Total steps: ~{total_steps:,}")
    print(f"   Learning rate: {config.learning_rate}")
    print(f"   Max sequence length: {config.max_seq_length}")
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=data_collator,
        processing_class=tokenizer,
    )
    print("\nStarting training...")
    print("-" * 70)
    
    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
    print("\n💾 Saving final model...")
    final_path = f"{config.output_dir}/final"
    trainer.save_model(final_path)
    tokenizer.save_pretrained(final_path)

    config_dict = {}
    for key in dir(config):
        if not key.startswith('__') and not callable(getattr(config, key)):
            config_dict[key] = getattr(config, key)
            
    if 'bnb_4bit_compute_dtype' in config_dict:
        config_dict['bnb_4bit_compute_dtype'] = str(config_dict['bnb_4bit_compute_dtype'])
        
    with open(f"{final_path}/training_config.json", 'w') as f:
        json.dump(config_dict, f, indent=2, default=str)
    
    print(f"\nTraining complete!")
    print(f"   Model saved to: {final_path}")
    print(f"   Finished at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    return trainer

def test_run(config):
    print("\nTEST RUN MODE")
    print("   Using 100 samples, 1 epoch\n")
    config.num_epochs = 1
    config.save_steps = 50
    config.eval_steps = 50
    config.logging_steps = 10
    config.output_dir = "outputs/test-run"
    train_data = []
    with open(config.train_file, 'r') as f:
        for i, line in enumerate(f):
            if i >= 100:
                break
            train_data.append(json.loads(line))
    
    val_data = train_data[:20]
    os.makedirs("data/test", exist_ok=True)
    with open("data/test/train.jsonl", 'w') as f:
        for item in train_data:
            f.write(json.dumps(item) + '\n')
    with open("data/test/val.jsonl", 'w') as f:
        for item in val_data:
            f.write(json.dumps(item) + '\n')
    
    config.train_file = "data/test/train.jsonl"
    config.val_file = "data/test/val.jsonl"
    train(config)

def main():
    parser = argparse.ArgumentParser(description='QLoRA Fine-tuning for ML-SLM')
    parser.add_argument('--resume', action='store_true', help='Resume from checkpoint')
    parser.add_argument('--test', action='store_true', help='Quick test run')
    parser.add_argument('--epochs', type=int, help='Override number of epochs')
    parser.add_argument('--lr', type=float, help='Override learning rate')
    parser.add_argument('--output', type=str, help='Override output directory')
    
    args = parser.parse_args()
    
    config = Config()
    if args.epochs:
        config.num_epochs = args.epochs
    if args.lr:
        config.learning_rate = args.lr
    if args.output:
        config.output_dir = args.output
    
    if args.test:
        test_run(config)
    else:
        train(config, resume_from_checkpoint=args.resume)

if __name__ == '__main__':
    main()