tinymind-native-8b-remote-handoff/bundle/model/tinymind-12b/train_12b_qlora.py

from __future__ import annotations

import argparse
import json
import logging
import math
import sys
import warnings
from pathlib import Path

logging.getLogger("torch.utils.flop_counter").setLevel(logging.ERROR)
warnings.filterwarnings(
    "ignore",
    message=r"_check_is_size will be removed in a future PyTorch release.*",
    category=FutureWarning,
    module=r"bitsandbytes\.backends\.cuda\.ops",
)

import torch
import torch.nn.functional as F

PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))
from train.lora_patterns import parse_lora_pattern


ROOT = Path(r"D:\ad\tinymind\model\tinymind-12b")


class TinyListDataset:
    """Minimal dataset for parser/unit tests without importing pyarrow/datasets."""

    def __init__(self, rows: list[dict[str, object]]):
        self.rows = rows
        self.column_names = list(rows[0].keys()) if rows else []

    @classmethod
    def from_list(cls, rows: list[dict[str, object]]) -> "TinyListDataset":
        return cls(rows)

    def __len__(self) -> int:
        return len(self.rows)

    def __getitem__(self, idx):
        if isinstance(idx, str):
            return [row[idx] for row in self.rows]
        return self.rows[idx]

    def shuffle(self, seed: int) -> "TinyListDataset":
        import random

        rng = random.Random(seed)
        rows = list(self.rows)
        rng.shuffle(rows)
        return TinyListDataset(rows)

    def select(self, indices) -> "TinyListDataset":
        return TinyListDataset([self.rows[i] for i in indices])


def make_dataset(rows: list[dict[str, object]], *, backend: str = "auto"):
    if backend == "tiny":
        return TinyListDataset.from_list(rows)
    from datasets import Dataset

    return Dataset.from_list(rows)


def format_messages(tokenizer, messages: list[dict]) -> str:
    if getattr(tokenizer, "chat_template", None):
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
    chunks: list[str] = []
    for message in messages:
        role = str(message.get("role", "user")).strip().upper()
        content = str(message.get("content", "")).strip()
        if content:
            chunks.append(f"{role}: {content}")
    chunks.append("ASSISTANT:")
    return "\n\n".join(chunks)


def _loss_weight(item: dict) -> float:
    qg = item.get("quality_governor") or {}
    meta = item.get("metadata") or {}
    for src in (qg, meta):
        if src.get("loss_weight") is not None:
            try:
                return float(src["loss_weight"])
            except (TypeError, ValueError):
                pass
    source = str(item.get("source", "")).lower()
    if "coverage_100k" in source:
        return 0.15
    if "logic_agent_code" in source:
        return 1.25
    return 1.0


def load_chat_dataset(
    path: str | Path,
    tokenizer,
    limit: int | None = None,
    *,
    dataset_format: str = "text",
    dataset_backend: str = "auto",
):
    rows: list[dict[str, object]] = []
    with Path(path).open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            item = json.JSONDecoder(strict=False).decode(line)
            messages = item.get("messages", [])
            if dataset_format == "prompt_completion" and messages:
                normalized_messages = []
                for message in messages:
                    role = str(message.get("role", "")).strip()
                    content = str(message.get("content", "")).strip()
                    if role in {"system", "user", "assistant"} and content:
                        normalized_messages.append({"role": role, "content": content})
                assistant_indices = [idx for idx, message in enumerate(normalized_messages) if message["role"] == "assistant"]
                if assistant_indices:
                    last_assistant = assistant_indices[-1]
                    prompt_messages = normalized_messages[:last_assistant]
                    completion = normalized_messages[last_assistant]["content"]
                    if prompt_messages and completion:
                        if getattr(tokenizer, "chat_template", None):
                            prompt = tokenizer.apply_chat_template(
                                prompt_messages,
                                tokenize=False,
                                add_generation_prompt=True,
                            )
                        else:
                            prompt = format_messages(tokenizer, prompt_messages)
                        rows.append({"prompt": prompt, "completion": completion, "loss_weight": _loss_weight(item)})
            elif dataset_format == "conversational" and messages:
                normalized_messages = []
                for message in messages:
                    role = str(message.get("role", "")).strip()
                    content = str(message.get("content", "")).strip()
                    if role in {"system", "user", "assistant"} and content:
                        normalized_messages.append({"role": role, "content": content})
                if normalized_messages and any(m["role"] == "assistant" for m in normalized_messages):
                    rows.append({"messages": normalized_messages, "loss_weight": _loss_weight(item)})
            elif messages:
                text = format_messages(tokenizer, messages)
                if text:
                    rows.append({"text": text, "loss_weight": _loss_weight(item)})
            else:
                text = item.get("text", "")
                if text:
                    rows.append({"text": text, "loss_weight": _loss_weight(item)})
            if limit is not None and len(rows) >= limit:
                break
    return make_dataset(rows, backend=dataset_backend)


def effective_record_limit(
    *,
    limit_records: int | None,
    eval_records: int,
    eval_only: bool,
    eval_only_load_multiplier: int,
) -> int | None:
    if not eval_only:
        return limit_records
    cap = max(eval_records, 1) * max(eval_only_load_multiplier, 1)
    return min(limit_records, cap) if limit_records is not None else cap


def train_eval_split(
    dataset,
    eval_records: int,
    *,
    seed: int = 20260524,
    min_eval_records: int = 2,
) -> tuple[Dataset, Dataset | None]:
    if eval_records <= 0 or len(dataset) <= 1:
        return dataset, None
    n_eval = min(max(eval_records, min_eval_records), len(dataset) - 1)
    if n_eval < min_eval_records:
        raise ValueError(f"eval split too small: {n_eval} < min_eval_records={min_eval_records}")
    shuffled = dataset.shuffle(seed=seed)
    train = shuffled.select(range(0, len(shuffled) - n_eval))
    eval_ds = shuffled.select(range(len(shuffled) - n_eval, len(shuffled)))
    return train, eval_ds


class WeightedSFTTrainer:
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        weights = inputs.pop("loss_weight", None)
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        if labels is None:
            loss = outputs.loss
        else:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            token_loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
                reduction="none",
            ).view_as(shift_labels)
            mask = shift_labels.ne(-100).to(token_loss.dtype)
            per_sample = (token_loss * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0)
            if weights is not None:
                weights = weights.to(per_sample.device, dtype=per_sample.dtype).view(-1)
                active = weights.gt(0)
                if active.any():
                    loss = (per_sample[active] * weights[active]).sum() / weights[active].sum().clamp_min(1.0)
                else:
                    loss = per_sample.mean()
            else:
                loss = per_sample.mean()
        return (loss, outputs) if return_outputs else loss


def make_weighted_sft_trainer_class(sft_trainer_cls):
    class RuntimeWeightedSFTTrainer(WeightedSFTTrainer, sft_trainer_cls):
        pass

    return RuntimeWeightedSFTTrainer


def align_special_token_ids(model, tokenizer) -> dict[str, int | None]:
    token_ids = {
        "pad_token_id": tokenizer.pad_token_id,
        "eos_token_id": tokenizer.eos_token_id,
        "bos_token_id": tokenizer.bos_token_id,
    }
    for key, value in token_ids.items():
        if value is None:
            continue
        setattr(model.config, key, value)
        if getattr(model, "generation_config", None) is not None:
            setattr(model.generation_config, key, value)
    return token_ids


def main() -> int:
    from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from trl import SFTConfig, SFTTrainer

    RuntimeWeightedSFTTrainer = make_weighted_sft_trainer_class(SFTTrainer)

    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", default="mistralai/Mistral-Nemo-Instruct-2407")
    parser.add_argument("--dataset", default=str(ROOT / "data" / "tinymind_sft.jsonl"))
    parser.add_argument("--output", default=str(ROOT / "adapters" / "tinymind-12b-lora"))
    parser.add_argument("--max-seq-length", type=int, default=2048)
    parser.add_argument("--epochs", type=float, default=1.0)
    parser.add_argument("--max-steps", type=int, default=800)
    parser.add_argument("--learning-rate", type=float, default=1.5e-4)
    parser.add_argument("--limit-records", type=int, default=None)
    parser.add_argument("--eval-records", type=int, default=128)
    parser.add_argument("--min-train-records", type=int, default=8)
    parser.add_argument("--min-eval-records", type=int, default=2)
    parser.add_argument("--seed", type=int, default=20260524)
    parser.add_argument("--target-modules", nargs="*", default=None)
    parser.add_argument("--resume-adapter", default=None)
    parser.add_argument("--lora-rank", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--lora-dropout", type=float, default=0.05)
    parser.add_argument("--lora-rank-pattern", default=None, help="Comma list such as q_proj:32,gate_proj:128")
    parser.add_argument("--lora-alpha-pattern", default=None, help="Comma list such as q_proj:64,gate_proj:256")
    parser.add_argument("--per-device-train-batch-size", type=int, default=1)
    parser.add_argument("--per-device-eval-batch-size", type=int, default=4)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=16)
    parser.add_argument("--eval-only-load-multiplier", type=int, default=2)
    parser.add_argument("--dataset-format", choices=["text", "conversational", "prompt_completion"], default="text")
    parser.add_argument("--dataset-backend", choices=["auto", "tiny"], default="auto")
    parser.add_argument("--completion-only-loss", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--assistant-only-loss", action=argparse.BooleanOptionalAction, default=False)
    parser.add_argument("--eval-before-train", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--eval-only", action="store_true")
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model_id, trust_remote_code=True, fix_mistral_regex=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    record_limit = effective_record_limit(
        limit_records=args.limit_records,
        eval_records=args.eval_records,
        eval_only=args.eval_only,
        eval_only_load_multiplier=args.eval_only_load_multiplier,
    )
    dataset = load_chat_dataset(
        args.dataset,
        tokenizer,
        limit=record_limit,
        dataset_format=args.dataset_format,
        dataset_backend=args.dataset_backend,
    )
    if len(dataset) < args.min_train_records + max(0, args.min_eval_records):
        raise ValueError(
            f"dataset too small for credible training/eval: records={len(dataset)} "
            f"required>={args.min_train_records + max(0, args.min_eval_records)}"
        )
    train_dataset, eval_dataset = train_eval_split(
        dataset,
        args.eval_records,
        seed=args.seed,
        min_eval_records=args.min_eval_records,
    )
    if len(train_dataset) < args.min_train_records:
        raise ValueError(f"train split too small: {len(train_dataset)} < min_train_records={args.min_train_records}")

    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        bnb_4bit_use_double_quant=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        quantization_config=bnb,
        device_map="auto",
        trust_remote_code=True,
        dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    )
    special_token_ids = align_special_token_ids(model, tokenizer)
    model = prepare_model_for_kbit_training(model)

    target_modules = args.target_modules
    if target_modules is None:
        model_type = getattr(model.config, "model_type", "")
        if model_type == "gpt2":
            target_modules = ["c_attn", "c_proj", "c_fc"]
        else:
            target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

    peft_config = None
    rank_pattern = parse_lora_pattern(args.lora_rank_pattern)
    alpha_pattern = parse_lora_pattern(args.lora_alpha_pattern)
    if args.resume_adapter:
        if rank_pattern or alpha_pattern:
            raise ValueError(
                "--resume-adapter cannot change rank_pattern/alpha_pattern. "
                "Start a new adapter from the base model for asymmetric-rank runs."
            )
        adapter_path = Path(args.resume_adapter)
        if not (adapter_path / "adapter_config.json").exists():
            raise FileNotFoundError(f"resume adapter is missing adapter_config.json: {adapter_path}")
        model = PeftModel.from_pretrained(model, str(adapter_path), is_trainable=True)
    else:
        peft_config = LoraConfig(
            r=args.lora_rank,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=target_modules,
            rank_pattern=rank_pattern,
            alpha_pattern=alpha_pattern,
        )

    warmup_steps = max(1, int(args.max_steps * 0.03)) if args.max_steps and args.max_steps > 0 else 0
    training_kwargs = {
        "output_dir": args.output,
        "per_device_train_batch_size": args.per_device_train_batch_size,
        "per_device_eval_batch_size": args.per_device_eval_batch_size,
        "gradient_accumulation_steps": args.gradient_accumulation_steps,
        "learning_rate": args.learning_rate,
        "num_train_epochs": args.epochs,
        "max_steps": args.max_steps,
        "max_length": args.max_seq_length,
        "warmup_steps": warmup_steps,
        "lr_scheduler_type": "cosine",
        "logging_steps": 5,
        "save_steps": 100,
        "save_total_limit": 3,
        "bf16": torch.cuda.is_available(),
        "fp16": False,
        "optim": "paged_adamw_8bit",
        "gradient_checkpointing": True,
        "remove_unused_columns": False,
        "report_to": [],
        "packing": False,
    }
    if args.dataset_format == "text":
        training_kwargs["dataset_text_field"] = "text"
    import inspect

    sft_params = inspect.signature(SFTConfig.__init__).parameters
    if "completion_only_loss" in sft_params:
        training_kwargs["completion_only_loss"] = (
            args.completion_only_loss if args.dataset_format in {"text", "prompt_completion"} else None
        )
    if "assistant_only_loss" in sft_params:
        training_kwargs["assistant_only_loss"] = bool(args.assistant_only_loss and args.dataset_format == "conversational")
    training_args = SFTConfig(**training_kwargs)

    trainer = RuntimeWeightedSFTTrainer(
        model=model,
        processing_class=tokenizer,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        peft_config=peft_config,
        args=training_args,
    )
    pre_train_eval_metrics = {}
    if args.eval_before_train and not args.eval_only and eval_dataset is not None:
        pre_train_eval_metrics = trainer.evaluate(metric_key_prefix="pre_train_eval")
    train_result = None if args.eval_only else trainer.train()
    eval_metrics = trainer.evaluate() if eval_dataset is not None else {}
    if not args.eval_only:
        trainer.save_model(args.output)
        tokenizer.save_pretrained(args.output)

    train_metrics = train_result.metrics if train_result is not None else {}
    train_loss = train_metrics.get("train_loss")
    train_loss_per_microbatch_estimate = None
    if train_loss is not None and args.gradient_accumulation_steps:
        train_loss_per_microbatch_estimate = float(train_loss) / float(args.gradient_accumulation_steps)

    manifest = {
        "base_model": args.model_id,
        "adapter": args.output,
        "dataset": args.dataset,
        "max_seq_length": args.max_seq_length,
        "max_steps": args.max_steps,
        "learning_rate": args.learning_rate,
        "eval_only": args.eval_only,
        "effective_record_limit": record_limit,
        "records_loaded": len(dataset),
        "train_records": len(train_dataset),
        "eval_records": len(eval_dataset) if eval_dataset is not None else 0,
        "seed": args.seed,
        "min_train_records": args.min_train_records,
        "min_eval_records": args.min_eval_records,
        "eval_split": "deterministic_shuffle_holdout",
        "eval_credible": eval_dataset is not None and len(eval_dataset) >= args.min_eval_records,
        "qlora": True,
        "lora_rank": args.lora_rank,
        "lora_alpha": args.lora_alpha,
        "lora_dropout": args.lora_dropout,
        "lora_rank_pattern": rank_pattern,
        "lora_alpha_pattern": alpha_pattern,
        "per_device_train_batch_size": args.per_device_train_batch_size,
        "per_device_eval_batch_size": args.per_device_eval_batch_size,
        "gradient_accumulation_steps": args.gradient_accumulation_steps,
        "eval_only_load_multiplier": args.eval_only_load_multiplier,
        "dataset_format": args.dataset_format,
        "dataset_backend": args.dataset_backend,
        "target_modules": target_modules,
        "resume_adapter": args.resume_adapter,
        "weighted_loss": True,
        "completion_only_loss_requested": args.completion_only_loss,
        "assistant_only_loss_requested": args.assistant_only_loss,
        "eval_before_train": args.eval_before_train,
        "special_token_ids": special_token_ids,
        "loss_weight_normalization": "weighted_mean_by_sum_of_weights",
        "pre_train_eval_metrics": pre_train_eval_metrics,
        "train_metrics": train_metrics,
        "train_loss_per_microbatch_estimate": train_loss_per_microbatch_estimate,
        "train_loss_note": (
            "HF Trainer train_loss is logged over gradient-accumulated training microbatches; "
            "compare eval_loss/pre_train_eval_loss for heldout quality."
        ),
        "eval_metrics": eval_metrics,
    }
    if eval_metrics.get("eval_loss") is not None:
        manifest["eval_loss"] = float(eval_metrics["eval_loss"])
        manifest["perplexity"] = float(math.exp(min(float(eval_metrics["eval_loss"]), 20.0)))
    Path(args.output).mkdir(parents=True, exist_ok=True)
    (Path(args.output) / "tinymind_12b_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    print(json.dumps(manifest, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())