trainer/train.py

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
from datasets import Dataset
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint

from .common import ARTIFACT_ROOT, CHECKPOINT_ROOT, DATASET_ROOT, atomic_json, ensure_dirs, log_event, update_status


BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-1.7B")
MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "2048"))
OUTPUT_NAME = os.environ.get("OUTPUT_NAME", "adapter")
TRAIN_LIMIT = int(os.environ.get("TRAIN_LIMIT", "0"))
MAX_STEPS = int(os.environ.get("MAX_STEPS", "-1"))
SEED = int(os.environ.get("TRAIN_SEED", "42"))


def load_rows(path: Path, limit: int = 0) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            rows.append(json.loads(line))
            if limit and len(rows) >= limit:
                break
    return rows


def common_prefix_length(left: list[int], right: list[int]) -> int:
    length = 0
    for a, b in zip(left, right):
        if a != b:
            break
        length += 1
    return length


def as_token_ids(value: Any) -> list[int]:
    if isinstance(value, dict) or hasattr(value, "keys"):
        value = value["input_ids"]
    if isinstance(value, torch.Tensor):
        value = value.tolist()
    if value and isinstance(value[0], list):
        value = value[0]
    return list(value)


def tokenized_dataset(tokenizer: Any, rows: list[dict[str, Any]]) -> Dataset:
    dataset = Dataset.from_list([{"messages": row["messages"], "id": row["id"]} for row in rows])

    def encode(example: dict[str, Any]) -> dict[str, Any]:
        messages = example["messages"]
        prefix = messages[:-1]
        prefix_ids = as_token_ids(
            tokenizer.apply_chat_template(
                prefix,
                tokenize=True,
                add_generation_prompt=True,
                enable_thinking=False,
            )
        )
        full_ids = as_token_ids(
            tokenizer.apply_chat_template(
                messages,
                tokenize=True,
                add_generation_prompt=False,
                enable_thinking=False,
            )
        )
        if tokenizer.eos_token_id is not None and (not full_ids or full_ids[-1] != tokenizer.eos_token_id):
            full_ids.append(tokenizer.eos_token_id)
        mask_to = common_prefix_length(prefix_ids, full_ids)
        input_ids = full_ids[:MAX_LENGTH]
        labels = [-100] * min(mask_to, len(input_ids)) + input_ids[min(mask_to, len(input_ids)) :]
        if not any(label != -100 for label in labels):
            raise ValueError(f"Assistant response was fully truncated for example {example['id']}")
        return {"input_ids": input_ids, "attention_mask": [1] * len(input_ids), "labels": labels}

    encoded = dataset.map(encode, remove_columns=dataset.column_names, desc="Tokenizing conversations")
    packed: list[dict[str, list[int]]] = []
    current = {"input_ids": [], "attention_mask": [], "labels": []}
    for item in encoded:
        item_length = len(item["input_ids"])
        if item_length > MAX_LENGTH:
            raise ValueError("Tokenized example exceeded MAX_LENGTH after truncation")
        if current["input_ids"] and len(current["input_ids"]) + item_length > MAX_LENGTH:
            packed.append(current)
            current = {"input_ids": [], "attention_mask": [], "labels": []}
        current["input_ids"].extend(item["input_ids"])
        current["attention_mask"].extend(item["attention_mask"])
        current["labels"].extend(item["labels"])
    if current["input_ids"]:
        packed.append(current)
    return Dataset.from_list(packed)


@dataclass
class AssistantOnlyCollator:
    pad_token_id: int

    def __call__(self, features: list[dict[str, list[int]]]) -> dict[str, torch.Tensor]:
        max_len = max(len(feature["input_ids"]) for feature in features)
        input_ids = []
        attention_mask = []
        labels = []
        for feature in features:
            padding = max_len - len(feature["input_ids"])
            input_ids.append(feature["input_ids"] + [self.pad_token_id] * padding)
            attention_mask.append(feature["attention_mask"] + [0] * padding)
            labels.append(feature["labels"] + [-100] * padding)
        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }


class StatusCallback(TrainerCallback):
    def on_log(self, args: TrainingArguments, state: Any, control: Any, logs: dict[str, Any] | None = None, **_: Any) -> None:
        if not logs:
            return
        numeric = {key: round(float(value), 6) for key, value in logs.items() if isinstance(value, (int, float))}
        update_status(
            "smoke_train" if OUTPUT_NAME == "smoke_adapter" else "full_train",
            f"Training step {state.global_step}",
            global_step=state.global_step,
            max_steps=state.max_steps,
            metrics=numeric,
        )


def main() -> None:
    ensure_dirs()
    if not torch.cuda.is_available():
        raise RuntimeError("QLoRA training requires the requested CUDA GPU")
    set_seed(SEED)
    output_dir = CHECKPOINT_ROOT / OUTPUT_NAME
    adapter_dir = ARTIFACT_ROOT / OUTPUT_NAME
    output_dir.mkdir(parents=True, exist_ok=True)
    adapter_dir.mkdir(parents=True, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True, trust_remote_code=False)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    train_rows = load_rows(DATASET_ROOT / "train.jsonl", TRAIN_LIMIT)
    eval_limit = min(200, max(30, len(train_rows) // 10)) if TRAIN_LIMIT else 500
    eval_rows = load_rows(DATASET_ROOT / "validation.jsonl", eval_limit)
    train_dataset = tokenized_dataset(tokenizer, train_rows)
    eval_dataset = tokenized_dataset(tokenizer, eval_rows)

    quantization = 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(
        BASE_MODEL,
        quantization_config=quantization,
        torch_dtype=torch.bfloat16,
        device_map={"": 0},
        trust_remote_code=False,
        attn_implementation="sdpa",
    )
    model.config.use_cache = False
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
    lora = LoraConfig(
        r=int(os.environ.get("LORA_R", "16")),
        lora_alpha=int(os.environ.get("LORA_ALPHA", "32")),
        lora_dropout=float(os.environ.get("LORA_DROPOUT", "0.05")),
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    )
    model = get_peft_model(model, lora)
    model.print_trainable_parameters()

    save_strategy = os.environ.get("SAVE_STRATEGY", "steps")
    smoke = OUTPUT_NAME == "smoke_adapter"
    args = TrainingArguments(
        output_dir=str(output_dir),
        overwrite_output_dir=False,
        num_train_epochs=float(os.environ.get("EPOCHS", "2")),
        max_steps=MAX_STEPS,
        per_device_train_batch_size=int(os.environ.get("BATCH_SIZE", "4")),
        per_device_eval_batch_size=int(os.environ.get("EVAL_BATCH_SIZE", "4")),
        gradient_accumulation_steps=int(os.environ.get("GRAD_ACCUM", "4")),
        learning_rate=float(os.environ.get("LEARNING_RATE", "1e-4")),
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        weight_decay=0.0,
        max_grad_norm=1.0,
        optim="paged_adamw_8bit",
        bf16=True,
        fp16=False,
        tf32=True,
        logging_steps=1 if smoke else 10,
        eval_strategy="no" if smoke else "steps",
        eval_steps=100,
        save_strategy=save_strategy,
        save_steps=100,
        save_total_limit=3,
        load_best_model_at_end=not smoke,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        report_to="none",
        remove_unused_columns=False,
        dataloader_num_workers=2,
        dataloader_pin_memory=True,
        gradient_checkpointing=True,
        seed=SEED,
        data_seed=SEED,
    )
    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=None if smoke else eval_dataset,
        data_collator=AssistantOnlyCollator(tokenizer.pad_token_id),
        callbacks=[StatusCallback()],
    )
    last_checkpoint = get_last_checkpoint(str(output_dir)) if output_dir.exists() and not smoke else None
    update_status(
        "smoke_train" if smoke else "full_train",
        "Model loaded; optimizer starting",
        base_model=BASE_MODEL,
        train_examples=len(train_dataset),
        validation_examples=len(eval_dataset),
        resume_from=last_checkpoint,
        gpu=torch.cuda.get_device_name(0),
    )
    result = trainer.train(resume_from_checkpoint=last_checkpoint)
    trainer.save_metrics("train", result.metrics)
    model.save_pretrained(adapter_dir, safe_serialization=True)
    tokenizer.save_pretrained(adapter_dir)
    metrics = {key: float(value) for key, value in result.metrics.items() if isinstance(value, (int, float))}
    atomic_json(adapter_dir / "run_metrics.json", metrics)
    atomic_json(
        adapter_dir / "training_config.json",
        {
            "base_model": BASE_MODEL,
            "method": "QLoRA NF4 double-quantization with BF16 compute",
            "assistant_only_loss": True,
            "packing": True,
            "thinking": False,
            "max_length": MAX_LENGTH,
            "seed": SEED,
            "source_train_examples": len(train_rows),
            "source_validation_examples": len(eval_rows),
            "packed_train_sequences": len(train_dataset),
            "packed_validation_sequences": len(eval_dataset),
            "lora_r": lora.r,
            "lora_alpha": lora.lora_alpha,
        },
    )
    log_event("training_complete", output=OUTPUT_NAME, metrics=metrics)
    update_status(
        "smoke_complete" if smoke else "training_complete",
        f"{OUTPUT_NAME} training completed",
        metrics=metrics,
        adapter=str(adapter_dir),
    )


if __name__ == "__main__":
    main()