scripts/train_dinodev_m4_qwen3_coder_30b.py

import os
import json
import hashlib
import inspect
import torch

from datasets import load_dataset, Dataset, concatenate_datasets
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer, SFTConfig

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-Coder-30B-A3B-Instruct")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/workspace/outputs/dinodev-m4-qwen3-coder-30b-code-sft-adapter")

MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "2048"))
MAX_STEPS = int(os.environ.get("MAX_STEPS", "1200"))
MAX_SAMPLES = int(os.environ.get("MAX_SAMPLES", "120000"))

print("Base model:", BASE_MODEL)
print("Output:", OUTPUT_DIR)
print("Max length:", MAX_LENGTH)
print("Max steps:", MAX_STEPS)
print("Max samples:", MAX_SAMPLES)

DATASETS = [
    "m-a-p/CodeFeedback-Filtered-Instruction",
    "theblackcat102/evol-codealpaca-v1",
    "iamtarun/code_instructions_120k_alpaca",
    "iamtarun/python_code_instructions_18k_alpaca",
]

SYSTEM_PROMPT = (
    "You are DinoDev, a senior full-stack coding assistant. "
    "Write clean, production-ready code. Explain important decisions briefly. "
    "Prefer secure, maintainable, testable solutions."
)

def first_value(row, keys):
    for k in keys:
        if k in row and row[k] is not None and str(row[k]).strip():
            return str(row[k]).strip()
    return ""

def row_to_messages(row):
    # Already chat formatted
    if "messages" in row and isinstance(row["messages"], list):
        return row["messages"]

    instruction = first_value(row, [
        "instruction", "query", "question", "prompt", "input", "problem", "task"
    ])

    extra_input = first_value(row, [
        "context", "additional_input", "given", "description"
    ])

    output = first_value(row, [
        "output", "response", "answer", "completion", "solution", "code"
    ])

    # Some alpaca rows have prompt containing full instruction
    if not instruction and "text" in row:
        instruction = str(row["text"]).strip()

    if extra_input and extra_input != instruction:
        user_content = instruction + "\n\nInput / Context:\n" + extra_input
    else:
        user_content = instruction

    if not user_content or not output:
        return None

    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
        {"role": "assistant", "content": output},
    ]

def load_and_format(tokenizer):
    formatted = []
    seen = set()

    for ds_name in DATASETS:
        print(f"\nLoading dataset: {ds_name}")
        try:
            ds = load_dataset(ds_name, split="train")
        except Exception as e:
            print(f"Skipping {ds_name}: {e}")
            continue

        print(ds)

        # Shuffle before taking rows
        ds = ds.shuffle(seed=42)

        for row in ds:
            messages = row_to_messages(row)
            if not messages:
                continue

            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=False,
            )

            # dedupe
            h = hashlib.sha256(text.encode("utf-8")).hexdigest()
            if h in seen:
                continue
            seen.add(h)

            # length safety before tokenization-heavy trainer
            if len(text) < 100:
                continue

            formatted.append({"text": text})

            if len(formatted) >= MAX_SAMPLES:
                break

        print(f"Collected so far: {len(formatted)}")

        if len(formatted) >= MAX_SAMPLES:
            break

    if not formatted:
        raise RuntimeError("No training samples created. Check dataset schemas.")

    return Dataset.from_list(formatted)

tokenizer = AutoTokenizer.from_pretrained(
    BASE_MODEL,
    trust_remote_code=True,
    use_fast=True,
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

dataset = load_and_format(tokenizer)
dataset = dataset.train_test_split(test_size=0.02, seed=42)

train_ds = dataset["train"]
eval_ds = dataset["test"]

print("Train rows:", len(train_ds))
print("Eval rows:", len(eval_ds))
print("Sample text:\n", train_ds[0]["text"][:1000])

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
)

model.config.use_cache = False

# Low-VRAM QLoRA prep:
# Avoid PEFT prepare_model_for_kbit_training() because it can cast large params to fp32
# and cause CUDA OOM on 30B MoE models.
for param in model.parameters():
    param.requires_grad = False

if hasattr(model, "gradient_checkpointing_enable"):
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})

if hasattr(model, "enable_input_require_grads"):
    model.enable_input_require_grads()
else:
    def make_inputs_require_grad(module, input, output):
        output.requires_grad_(True)
    model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    # Attention-only LoRA first. Safer for 30B MoE / 40GB-80GB GPU.
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
    ],
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Make SFTConfig compatible with older/newer TRL
cfg_params = inspect.signature(SFTConfig.__init__).parameters

kwargs = dict(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    learning_rate=2e-4,
    max_steps=MAX_STEPS,
    warmup_ratio=0.03,
    logging_steps=10,
    save_steps=100,
    eval_steps=100,
    bf16=True,
    optim="paged_adamw_8bit",
    lr_scheduler_type="cosine",
    gradient_checkpointing=True,
    report_to="none",
    save_total_limit=3,
)

if "eval_strategy" in cfg_params:
    kwargs["eval_strategy"] = "steps"
elif "evaluation_strategy" in cfg_params:
    kwargs["evaluation_strategy"] = "steps"

if "dataset_text_field" in cfg_params:
    kwargs["dataset_text_field"] = "text"

if "packing" in cfg_params:
    kwargs["packing"] = True

if "max_seq_length" in cfg_params:
    kwargs["max_seq_length"] = MAX_LENGTH
elif "max_length" in cfg_params:
    kwargs["max_length"] = MAX_LENGTH

training_args = SFTConfig(**kwargs)

try:
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        processing_class=tokenizer,
    )
except TypeError:
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        tokenizer=tokenizer,
    )

trainer.train()

trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

readme = f"""---
base_model: {BASE_MODEL}
library_name: peft
tags:
- qwen3
- coder
- qlora
- lora
- peft
- dinodev
- code-sft
---

# DinoDev M4 Qwen3 Coder 30B Code SFT Adapter

Base model: `{BASE_MODEL}`

Training type: QLoRA / PEFT LoRA adapter

Datasets:
- m-a-p/CodeFeedback-Filtered-Instruction
- theblackcat102/evol-codealpaca-v1
- iamtarun/code_instructions_120k_alpaca
- iamtarun/python_code_instructions_18k_alpaca

Settings:
- max_length: {MAX_LENGTH}
- max_steps: {MAX_STEPS}
- max_samples: {MAX_SAMPLES}
- LoRA r: 32
- LoRA alpha: 64
"""
with open(os.path.join(OUTPUT_DIR, "README.md"), "w") as f:
    f.write(readme)

print("Training complete. Saved to:", OUTPUT_DIR)