Add training script: QLoRA SFT on Qwen2.5-Coder-7B with SOTA code datasets

train.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+"""
+Train a Claude Code-level coding model via QLoRA SFT on Qwen2.5-Coder-7B-Instruct.
+
+Recipe (based on published SOTA results):
+- Base: Qwen2.5-Coder-7B-Instruct (88.4% HumanEval baseline)
+- Data: KodCode-V1-SFT-R1 (verified competitive programming with R1-style CoT)
+       + Code-Feedback (multi-turn code dialogue)
+       + Magicoder-OSS-Instruct (diverse code generation)
+       + Magicoder-Evol-Instruct (evolved code instructions)
+- Method: QLoRA (4-bit NF4 + LoRA r=64, all-linear)
+- Target: Push past 90%+ HumanEval, maximize LiveCodeBench with CoT reasoning
+
+References:
+- rStar-Coder (arxiv:2505.21297): Qwen2.5-Coder-7B → 57.3% LiveCodeBench
+- KodCode (arxiv:2503.02951): Verified coding dataset with R1-style reasoning
+- Qwen2.5-Coder (arxiv:2409.12186): Base model technical report
+"""
+
+import os
+import torch
+from datasets import load_dataset, concatenate_datasets
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import LoraConfig, prepare_model_for_kbit_training
+from trl import SFTConfig, SFTTrainer
+
+# ─── Configuration ────────────────────────────────────────────────────────────
+MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"
+OUTPUT_DIR = "./qwen25-coder-7b-mythos"
+HUB_MODEL_ID = "ashhhhhh26/qwen25-coder-32b-mythos"
+
+SYSTEM_PROMPT = """You are an elite software engineer with deep expertise across all programming languages, frameworks, and paradigms. You write clean, efficient, well-documented code. You think step-by-step through complex problems, consider edge cases, and provide production-quality solutions. When debugging, you methodically trace through the code to identify root causes. You explain your reasoning clearly and concisely."""
+
+# ─── 1. Load datasets ────────────────────────────────────────────────────────
+print("=" * 60)
+print("Loading datasets...")
+print("=" * 60)
+
+# Dataset 1: KodCode-V1-SFT-R1 (verified competitive programming + R1-style reasoning)
+ds_kodcode = load_dataset("KodCode/KodCode-V1-SFT-R1", split="train")
+print(f"KodCode-V1-SFT-R1 (raw): {len(ds_kodcode)} samples")
+
+# Filter for R1-verified correct solutions only
+ds_kodcode = ds_kodcode.filter(lambda x: x["r1_correctness"] is True, num_proc=4)
+print(f"KodCode-V1-SFT-R1 (r1_correctness=True): {len(ds_kodcode)} samples")
+
+# Dataset 2: Code-Feedback (~66K) - already in messages format
+ds_feedback = load_dataset("m-a-p/Code-Feedback", split="train")
+print(f"Code-Feedback: {len(ds_feedback)} samples")
+
+# Dataset 3: Magicoder-OSS-Instruct-75K
+ds_magicoder_oss = load_dataset("ise-uiuc/Magicoder-OSS-Instruct-75K", split="train")
+print(f"Magicoder-OSS-Instruct: {len(ds_magicoder_oss)} samples")
+
+# Dataset 4: Magicoder-Evol-Instruct-110K
+ds_magicoder_evol = load_dataset("ise-uiuc/Magicoder-Evol-Instruct-110K", split="train")
+print(f"Magicoder-Evol-Instruct: {len(ds_magicoder_evol)} samples")
+
+
+# ─── 2. Convert all datasets to ChatML messages format ───────────────────────
+print("\\nConverting datasets to ChatML format...")
+
+def convert_kodcode(example):
+    """Convert KodCode conversations (human/gpt) to standard ChatML messages."""
+    role_map = {"human": "user", "gpt": "assistant"}
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for msg in example["conversations"]:
+        role = role_map.get(msg["from"], msg["from"])
+        messages.append({"role": role, "content": msg["value"]})
+    return {"messages": messages}
+
+def convert_feedback(example):
+    """Code-Feedback already has messages, just add system prompt."""
+    messages = example["messages"]
+    if messages and messages[0]["role"] != "system":
+        messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
+    return {"messages": messages}
+
+def convert_magicoder_oss(example):
+    """Convert problem/solution to messages format."""
+    return {
+        "messages": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["problem"]},
+            {"role": "assistant", "content": example["solution"]},
+        ]
+    }
+
+def convert_magicoder_evol(example):
+    """Convert instruction/response to messages format."""
+    return {
+        "messages": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["instruction"]},
+            {"role": "assistant", "content": example["response"]},
+        ]
+    }
+
+# Apply conversions
+ds_kodcode = ds_kodcode.map(convert_kodcode, num_proc=4, remove_columns=ds_kodcode.column_names)
+ds_feedback = ds_feedback.map(convert_feedback, num_proc=4, remove_columns=[c for c in ds_feedback.column_names if c != "messages"])
+ds_magicoder_oss = ds_magicoder_oss.map(convert_magicoder_oss, num_proc=4, remove_columns=ds_magicoder_oss.column_names)
+ds_magicoder_evol = ds_magicoder_evol.map(convert_magicoder_evol, num_proc=4, remove_columns=ds_magicoder_evol.column_names)
+
+# Combine all datasets
+combined_dataset = concatenate_datasets([ds_kodcode, ds_feedback, ds_magicoder_oss, ds_magicoder_evol])
+combined_dataset = combined_dataset.shuffle(seed=42)
+print(f"\\nTotal combined dataset: {len(combined_dataset)} samples")
+
+# Quality filter
+def filter_quality(example):
+    """Remove examples with very short responses."""
+    msgs = example["messages"]
+    assistant_msgs = [m for m in msgs if m["role"] == "assistant"]
+    if not assistant_msgs:
+        return False
+    total_assistant_len = sum(len(m["content"]) for m in assistant_msgs)
+    return total_assistant_len >= 50
+
+combined_dataset = combined_dataset.filter(filter_quality, num_proc=4)
+print(f"After quality filter: {len(combined_dataset)} samples")
+
+
+# ─── 3. Load model with QLoRA (4-bit quantization) ───────────────────────────
+print("\\n" + "=" * 60)
+print(f"Loading {MODEL_ID} with 4-bit quantization...")
+print("=" * 60)
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    quantization_config=bnb_config,
+    attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+model = prepare_model_for_kbit_training(model)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+print(f"Model loaded. Parameters: {model.num_parameters():,}")
+
+
+# ─── 4. LoRA config ──────────────────────────────────────────────────────────
+peft_config = LoraConfig(
+    r=64,
+    lora_alpha=128,
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+    target_modules="all-linear",
+)
+
+
+# ─── 5. Training config ──────────────────────────────────────────────────────
+training_args = SFTConfig(
+    output_dir=OUTPUT_DIR,
+
+    # Data
+    max_length=4096,
+    packing=True,
+    dataset_num_proc=8,
+
+    # Training hyperparams
+    num_train_epochs=2,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=16,
+    learning_rate=2e-4,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.05,
+    weight_decay=0.01,
+    max_grad_norm=1.0,
+    optim="paged_adamw_8bit",
+
+    # Memory optimization
+    gradient_checkpointing=True,
+    bf16=True,
+    tf32=True,
+
+    # Logging
+    logging_steps=5,
+    logging_first_step=True,
+    disable_tqdm=True,
+
+    # Saving & Hub
+    save_strategy="steps",
+    save_steps=1000,
+    save_total_limit=3,
+    push_to_hub=True,
+    hub_model_id=HUB_MODEL_ID,
+    hub_strategy="every_save",
+
+    # Monitoring
+    report_to="trackio",
+    run_name="qwen25-coder-7b-mythos-sft",
+    project="code-mythos",
+
+    # Misc
+    seed=42,
+    dataloader_num_workers=4,
+    remove_unused_columns=False,
+)
+
+
+# ─── 6. Create trainer and train ─────────────────────────────────────────────
+print("\\n" + "=" * 60)
+print("Initializing SFTTrainer...")
+print("=" * 60)
+
+trainer = SFTTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=combined_dataset,
+    processing_class=tokenizer,
+    peft_config=peft_config,
+)
+
+trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total_params = model.num_parameters()
+print(f"Trainable: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)")
+
+print("\\n" + "=" * 60)
+print("Starting training...")
+print("=" * 60)
+
+trainer.train()
+
+# ─── 7. Save and push final model ────────────────────────────────────────────
+print("\\n" + "=" * 60)
+print("Saving final model...")
+print("=" * 60)
+
+trainer.save_model(OUTPUT_DIR)
+trainer.push_to_hub()
+
+print("\\n" + "=" * 60)
+print(f"✅ Training complete! Model pushed to: https://huggingface.co/{HUB_MODEL_ID}")
+print("=" * 60)