Add validated training script

train_coding_agent.py

@@ -0,0 +1,267 @@
+"""
+Qwen3-8B Coding & Agentic Reasoning Expert — Multi-Dataset SFT Training
+========================================================================
+Base: Qwen/Qwen3-8B (Apache 2.0, 8.2B params, 32K context)
+Method: QLoRA SFT with assistant-only loss masking
+Datasets:
+  - TIGER-Lab/VisCode-200K (visualization/chart generation) — ChatML ready
+  - m-a-p/CodeFeedback-Filtered-Instruction (code instruction tuning)
+  - nvidia/OpenCodeReasoning (reasoning with <think> blocks)
+  - glaiveai/glaive-function-calling-v2 (tool calling)
+  - ise-uiuc/Magicoder-OSS-Instruct-75K (code generation)
+
+Recipe: Based on Qwen3-Coder-Next + LoRA Without Regret papers
+Target: Coding + agentic reasoning + visualization + tool-use expert
+
+Usage:
+  pip install transformers>=4.51.0 trl>=1.3.0 peft>=0.15.0 datasets accelerate bitsandbytes torch trackio
+  HUB_MODEL_ID=your-username/model-name python train_coding_agent.py
+"""
+
+import os
+import re
+import json
+import torch
+import trackio
+from datasets import load_dataset, concatenate_datasets, Dataset
+from transformers import AutoTokenizer, BitsAndBytesConfig, TrainerCallback
+from trl import SFTTrainer, SFTConfig
+from peft import LoraConfig, TaskType
+
+# ============================================================
+# Configuration
+# ============================================================
+MODEL_ID = "Qwen/Qwen3-8B"
+OUTPUT_DIR = "./qwen3-8b-coding-agent"
+HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "sukritvemula/Qwen3-8B-CodeAgent")
+
+# Training hyperparameters (from Qwen3 + LoRA Without Regret papers)
+LEARNING_RATE = 2e-4
+NUM_EPOCHS = 2
+BATCH_SIZE = 2
+GRAD_ACCUM = 8
+MAX_LENGTH = 4096
+LORA_R = 64
+LORA_ALPHA = 16
+WARMUP_RATIO = 0.05
+
+# Dataset proportions (~50K samples)
+MAX_VISCODE = 12000
+MAX_CODEFEEDBACK = 10000
+MAX_OPENCODE = 10000
+MAX_GLAIVE = 8000
+MAX_MAGICODER = 10000
+
+SYSTEM_PROMPT = """You are an expert AI assistant specialized in coding, agentic reasoning, data visualization, and tool use. You can:
+1. Write, debug, and explain code in any programming language
+2. Reason step-by-step through complex problems using <think>...</think> blocks
+3. Generate charts, graphs, and data visualizations using matplotlib, plotly, seaborn
+4. Call functions and tools when needed, returning structured JSON for tool invocations
+5. Search the web and read research papers to provide accurate, up-to-date information
+6. Replicate images and diagrams programmatically
+
+Always think carefully before responding. Be precise, avoid hallucination, and cite sources when possible."""
+
+
+class AlertCallback(TrainerCallback):
+    def __init__(self):
+        self.best_loss = float('inf')
+        self.initial_loss = None
+        self.steps_since_improvement = 0
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs is None:
+            return
+        loss = logs.get("loss")
+        if loss is None:
+            return
+        step = state.global_step
+        if self.initial_loss is None:
+            self.initial_loss = loss
+            trackio.alert(title="Training Started", text=f"Initial loss={loss:.4f} at step {step}. Model: {MODEL_ID}, lr={LEARNING_RATE}, batch={BATCH_SIZE}x{GRAD_ACCUM}={BATCH_SIZE*GRAD_ACCUM}", level="INFO")
+        if loss != loss or loss > 20.0:
+            trackio.alert(title="DIVERGENCE DETECTED", text=f"loss={loss} at step {step} — training has diverged. lr likely too high, try lr={LEARNING_RATE*0.1:.1e}", level="ERROR")
+            return
+        if loss < self.best_loss:
+            self.best_loss = loss
+            self.steps_since_improvement = 0
+        else:
+            self.steps_since_improvement += 1
+        if step > 100 and loss > self.initial_loss * 0.9:
+            trackio.alert(title="Slow Convergence", text=f"loss={loss:.4f} at step {step}, only {((self.initial_loss - loss) / self.initial_loss * 100):.1f}% reduction from initial {self.initial_loss:.4f}. Consider lr={LEARNING_RATE*2:.1e}", level="WARN")
+        if self.steps_since_improvement > 200:
+            trackio.alert(title="Loss Stagnation", text=f"No improvement for {self.steps_since_improvement} steps. Best loss={self.best_loss:.4f}, current={loss:.4f}.", level="WARN")
+        if step > 0 and step % 500 == 0:
+            trackio.alert(title="Training Milestone", text=f"Step {step}: loss={loss:.4f}, best_loss={self.best_loss:.4f}, lr={logs.get('learning_rate', 'N/A')}", level="INFO")
+
+
+def process_viscode(max_samples):
+    print(f"Loading VisCode-200K (max {max_samples})...")
+    ds = load_dataset("TIGER-Lab/VisCode-200K", split=f"train[:{max_samples}]")
+    def add_system(example):
+        messages = example["messages"]
+        if messages and messages[0]["role"] != "system":
+            messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
+        return {"messages": messages}
+    ds = ds.map(add_system, num_proc=4)
+    print(f"  VisCode: {len(ds)} samples loaded")
+    return ds
+
+
+def process_codefeedback(max_samples):
+    print(f"Loading CodeFeedback (max {max_samples})...")
+    ds = load_dataset("m-a-p/CodeFeedback-Filtered-Instruction", split=f"train[:{max_samples}]")
+    def to_messages(example):
+        return {"messages": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["query"]},
+            {"role": "assistant", "content": example["answer"]}
+        ]}
+    ds = ds.map(to_messages, remove_columns=ds.column_names, num_proc=4)
+    print(f"  CodeFeedback: {len(ds)} samples loaded")
+    return ds
+
+
+def process_opencode_reasoning(max_samples):
+    print(f"Loading OpenCodeReasoning (max {max_samples})...")
+    ds = load_dataset("nvidia/OpenCodeReasoning", "split_0", split=f"split_0[:{max_samples}]")
+    def to_messages(example):
+        return {"messages": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["input"]},
+            {"role": "assistant", "content": example["output"]}
+        ]}
+    ds = ds.map(to_messages, remove_columns=ds.column_names, num_proc=4)
+    print(f"  OpenCodeReasoning: {len(ds)} samples loaded")
+    return ds
+
+
+def process_glaive_function_calling(max_samples):
+    print(f"Loading Glaive Function Calling (max {max_samples})...")
+    ds = load_dataset("glaiveai/glaive-function-calling-v2", split=f"train[:{max_samples}]")
+    def to_messages(example):
+        system_content = re.sub(r'^SYSTEM:\s*', '', example["system"])
+        chat = example["chat"]
+        messages = [{"role": "system", "content": system_content}]
+        parts = re.split(r'\n*(USER:|ASSISTANT:|FUNCTION RESPONSE:)', chat)
+        current_role, current_content = None, ""
+        for part in parts:
+            part = part.strip()
+            if not part:
+                continue
+            if part == "USER:":
+                if current_role and current_content.strip():
+                    messages.append({"role": current_role, "content": current_content.strip()})
+                current_role, current_content = "user", ""
+            elif part == "ASSISTANT:":
+                if current_role and current_content.strip():
+                    messages.append({"role": current_role, "content": current_content.strip()})
+                current_role, current_content = "assistant", ""
+            elif part == "FUNCTION RESPONSE:":
+                if current_role and current_content.strip():
+                    messages.append({"role": current_role, "content": current_content.strip()})
+                current_role, current_content = "user", "[Function Response] "
+            else:
+                current_content += part
+        if current_role and current_content.strip():
+            messages.append({"role": current_role, "content": current_content.strip()})
+        merged = [messages[0]]
+        for msg in messages[1:]:
+            if merged and msg["role"] == merged[-1]["role"]:
+                merged[-1]["content"] += "\n" + msg["content"]
+            else:
+                merged.append(msg)
+        messages = merged
+        if len(messages) < 3 or messages[-1]["role"] != "assistant":
+            return {"messages": []}
+        return {"messages": messages}
+    ds = ds.map(to_messages, remove_columns=ds.column_names, num_proc=4)
+    ds = ds.filter(lambda x: len(x["messages"]) >= 3 and any(m["role"] == "assistant" for m in x["messages"]))
+    print(f"  Glaive Function Calling: {len(ds)} samples loaded")
+    return ds
+
+
+def process_magicoder(max_samples):
+    print(f"Loading Magicoder (max {max_samples})...")
+    ds = load_dataset("ise-uiuc/Magicoder-OSS-Instruct-75K", split=f"train[:{max_samples}]")
+    def to_messages(example):
+        return {"messages": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["problem"]},
+            {"role": "assistant", "content": example["solution"]}
+        ]}
+    ds = ds.map(to_messages, remove_columns=ds.column_names, num_proc=4)
+    print(f"  Magicoder: {len(ds)} samples loaded")
+    return ds
+
+
+def main():
+    print("=" * 60)
+    print("Qwen3-8B Coding & Agentic Reasoning Expert Training")
+    print("=" * 60)
+
+    datasets_list = []
+    for loader in [
+        lambda: process_viscode(MAX_VISCODE),
+        lambda: process_codefeedback(MAX_CODEFEEDBACK),
+        lambda: process_opencode_reasoning(MAX_OPENCODE),
+        lambda: process_glaive_function_calling(MAX_GLAIVE),
+        lambda: process_magicoder(MAX_MAGICODER),
+    ]:
+        try:
+            datasets_list.append(loader())
+        except Exception as e:
+            print(f"  ⚠️ Failed: {e}")
+
+    if not datasets_list:
+        raise ValueError("No datasets loaded!")
+
+    combined = concatenate_datasets(datasets_list).shuffle(seed=42)
+    print(f"✅ Total training samples: {len(combined)}")
+
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True, bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True,
+    )
+    peft_config = LoraConfig(
+        r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=0.05, bias="none",
+        task_type=TaskType.CAUSAL_LM, target_modules="all-linear", use_rslora=True,
+    )
+    training_args = SFTConfig(
+        output_dir=OUTPUT_DIR, push_to_hub=True, hub_model_id=HUB_MODEL_ID,
+        hub_strategy="every_save", max_length=MAX_LENGTH, packing=False,
+        assistant_only_loss=True, num_train_epochs=NUM_EPOCHS,
+        per_device_train_batch_size=BATCH_SIZE, gradient_accumulation_steps=GRAD_ACCUM,
+        learning_rate=LEARNING_RATE, lr_scheduler_type="cosine", warmup_ratio=WARMUP_RATIO,
+        weight_decay=0.01, max_grad_norm=1.0, bf16=True, tf32=True,
+        gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False},
+        logging_steps=10, logging_first_step=True, disable_tqdm=True,
+        save_strategy="steps", save_steps=500, save_total_limit=3, eval_strategy="no",
+        report_to="trackio", run_name="sft-qwen3-8b-coding-agent-v1",
+        model_init_kwargs={
+            "quantization_config": bnb_config, "device_map": "auto",
+            "use_cache": False, "torch_dtype": torch.bfloat16,
+        },
+        seed=42, dataloader_num_workers=4, dataloader_pin_memory=True,
+    )
+
+    trainer = SFTTrainer(
+        model=MODEL_ID, args=training_args, train_dataset=combined,
+        peft_config=peft_config, callbacks=[AlertCallback()],
+    )
+
+    total_params = sum(p.numel() for p in trainer.model.parameters())
+    trainable_params = sum(p.numel() for p in trainer.model.parameters() if p.requires_grad)
+    print(f"Total: {total_params:,} | Trainable: {trainable_params:,} ({100 * trainable_params / total_params:.2f}%)")
+
+    train_result = trainer.train()
+    trainer.save_model(OUTPUT_DIR)
+    trainer.push_to_hub(commit_message="Training complete: Qwen3-8B Coding Agent v1")
+
+    metrics = train_result.metrics
+    trackio.alert(title="Training Complete", text=f"Final loss={metrics.get('train_loss', 'N/A')}, hub_model={HUB_MODEL_ID}", level="INFO")
+    print(f"✅ DONE! Model: https://huggingface.co/{HUB_MODEL_ID}")
+
+
+if __name__ == "__main__":
+    main()