train_sft.py

#!/usr/bin/env python3
"""
Reference SFT training script for agentic coding.
Loads a 60/30/10 mix of SWE + tool-use + code-act datasets,
normalizes to unified message format with multi-template tool formats.

Usage:
    python train_sft.py \
        --model nvidia/Nemotron-Terminal-8B \
        --output_dir ./nexus-coder-sft
"""

import argparse
import random
import json
from datasets import load_dataset, concatenate_datasets, Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig

# ---------------------------------------------------------------------------
# Tool template formatters (multi-template trick for generalization)
# ---------------------------------------------------------------------------

def format_openai_json(tool_calls: list) -> str:
    out = ""
    for tc in tool_calls:
        out += f'<tool_call>{{"type": "function", "function": {{"name": "{tc.get("name","")}", "arguments": {json.dumps(tc.get("arguments",""))}}}}}</tool_call>\n'
    return out.strip()

def format_xml(tool_calls: list) -> str:
    out = ""
    for tc in tool_calls:
        out += f"<tool_call><name>{tc.get('name','')}</name><arguments>{tc.get('arguments','')}</arguments></tool_call>\n"
    return out.strip()

def format_python(tool_calls: list) -> str:
    out = ""
    for tc in tool_calls:
        out += f"{tc.get('name','')}({tc.get('arguments','')})\n"
    return out.strip()

def format_typescript(tool_calls: list) -> str:
    out = ""
    for tc in tool_calls:
        out += f"{{ tool: '{tc.get('name','')}', args: {tc.get('arguments','')} }}\n"
    return out.strip()

def format_qwen3_xml(tool_calls: list) -> str:
    out = ""
    for tc in tool_calls:
        out += f"<qwen3_coder><tool>{tc.get('name','')}</tool><params>{tc.get('arguments','')}</params></qwen3_coder>\n"
    return out.strip()

FORMAT_CHOICES = [format_openai_json, format_xml, format_python, format_typescript, format_qwen3_xml]

# ---------------------------------------------------------------------------
# Dataset loaders
# ---------------------------------------------------------------------------

def load_swe_smith(tokenizer) -> Dataset:
    """Load SWE-smith trajectories (tool split, resolved only)."""
    ds = load_dataset("SWE-bench/SWE-smith-trajectories", split="tool")
    ds = ds.filter(lambda x: x.get("resolved", False) is True)
    def normalize(example):
        msgs = example.get("messages", [])
        if isinstance(msgs, str):
            msgs = json.loads(msgs)
        text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False)
        return {"text": text}
    return ds.map(normalize, remove_columns=ds.column_names)

def load_nemotron_agentic(tokenizer) -> Dataset:
    """Load Nemotron-Agentic-v1 interactive_agent + tool_calling."""
    ds_ia = load_dataset("nvidia/Nemotron-Agentic-v1", split="interactive_agent")
    ds_tc = load_dataset("nvidia/Nemotron-Agentic-v1", split="tool_calling")
    ds = concatenate_datasets([ds_ia, ds_tc])
    def normalize(example):
        msgs = example.get("messages", [])
        if isinstance(msgs, str):
            msgs = json.loads(msgs)
        # Apply random template to any assistant tool_calls
        for m in msgs:
            if m.get("role") == "assistant" and m.get("tool_calls"):
                fmt = random.choice(FORMAT_CHOICES)
                m["content"] = fmt(m["tool_calls"])
        text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False)
        return {"text": text}
    return ds.map(normalize, remove_columns=ds.column_names)

def load_code_act(tokenizer) -> Dataset:
    """Load xingyaoww/code-act codeact split."""
    ds = load_dataset("xingyaoww/code-act", split="codeact")
    def normalize(example):
        conv = example.get("conversations", [])
        if isinstance(conv, str):
            conv = json.loads(conv)
        msgs = []
        for c in conv:
            role = "user" if c.get("from") in ("human", "user") else "assistant"
            if c.get("from") == "system":
                role = "system"
            msgs.append({"role": role, "content": c.get("value", "")})
        text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False)
        return {"text": text}
    return ds.map(normalize, remove_columns=ds.column_names)

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="nvidia/Nemotron-Terminal-8B")
    parser.add_argument("--output_dir", default="./nexus-coder-sft")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--grad_accum", type=int, default=8)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--max_seq_length", type=int, default=16384)
    parser.add_argument("--hub_model_id", default=None)
    parser.add_argument("--lora", action="store_true", help="Use LoRA if VRAM-constrained")
    parser.add_argument("--lora_r", type=int, default=64)
    parser.add_argument("--lora_alpha", type=int, default=128)
    args = parser.parse_args()

    print("[1/5] Loading model and tokenizer...")
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        torch_dtype="bfloat16",
        device_map="auto",
        trust_remote_code=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # LoRA setup if requested
    peft_config = None
    if args.lora:
        from peft import LoraConfig, TaskType
        peft_config = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
            task_type=TaskType.CAUSAL_LM,
            lora_dropout=0.05,
            bias="none",
        )
        print(f"    Using LoRA (r={args.lora_r}, alpha={args.lora_alpha})")

    print("[2/5] Loading and mixing datasets...")
    ds_swe = load_swe_smith(tokenizer)
    ds_agentic = load_nemotron_agentic(tokenizer)
    ds_code = load_code_act(tokenizer)

    # Shuffle and sample to approximate 60/30/10 by token count
    # Simple heuristic: sample proportional to raw example counts
    n_swe = min(len(ds_swe), 10000)
    n_agentic = min(len(ds_agentic), 5000)
    n_code = min(len(ds_code), 2000)
    ds_swe = ds_swe.shuffle(seed=42).select(range(n_swe))
    ds_agentic = ds_agentic.shuffle(seed=42).select(range(n_agentic))
    ds_code = ds_code.shuffle(seed=42).select(range(n_code))

    mixed = concatenate_datasets([ds_swe, ds_agentic, ds_code])
    mixed = mixed.shuffle(seed=42)
    print(f"    Mixed dataset: {len(mixed)} examples")

    print("[3/5] Applying multi-template normalization...")
    def ensure_text(example):
        return {"text": example.get("text", "")}
    mixed = mixed.map(ensure_text).filter(lambda x: len(x.get("text", "")) > 200)

    print("[4/5] Configuring SFT trainer...")
    sft_config = SFTConfig(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        max_seq_length=args.max_seq_length,
        logging_strategy="steps",
        logging_steps=10,
        logging_first_step=True,
        save_strategy="epoch",
        bf16=True,
        gradient_checkpointing=True,
        disable_tqdm=True,
        push_to_hub=args.hub_model_id is not None,
        hub_model_id=args.hub_model_id,
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=mixed,
        args=sft_config,
        peft_config=peft_config,
    )

    print("[5/5] Starting SFT training...")
    trainer.train()
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)
    print(f"Done. Model saved to {args.output_dir}")


if __name__ == "__main__":
    main()