Add trainer/train.py

trainer/train.py

@@ -0,0 +1,265 @@
+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}
+
+    return dataset.map(encode, remove_columns=dataset.column_names, desc="Tokenizing conversations")
+
+
+@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,
+            "thinking": False,
+            "max_length": MAX_LENGTH,
+            "seed": SEED,
+            "train_examples": len(train_dataset),
+            "validation_examples": 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()