train.py

#!/usr/bin/env python3
"""
LeMath: LeJEPA-style fine-tuning for math reasoning LLMs.

Implements:
  - Dataset: nvidia/OpenMathReasoning (cot split) (problem, generated_solution)
  - Token injection:
      immediately after <think>: <|jeton|> repeated with density 1 per N tokens (covers remainder)
  - Objective:
      total = decoding_loss
            + w_sigreg * LeJEPA SIGReg loss on jeton hidden states (via lejepa if installed)
      + w_cos    * mean(1 - cosine(h_jeton_i, h_target_i)) for each i-th N token pairing

Uses a custom PyTorch training loop (not SFTTrainer) to access hidden states.

References:
  - LLM-JEPA: https://arxiv.org/abs/2509.14252
  - LeJEPA / SIGReg: https://arxiv.org/abs/2511.08544 and https://github.com/galilai-group/lejepa
  - Unsloth Qwen3 guide: https://unsloth.ai/blog/qwen3
  - Small batch training: https://arxiv.org/abs/2507.07101
  - Dataset: https://huggingface.co/datasets/nvidia/OpenMathReasoning
"""
from __future__ import annotations

import argparse
import contextlib
import gc
import os
import random
import time
from typing import Any, Dict, Iterable, List, Optional, Tuple
from unsloth import FastLanguageModel # type: ignore
from datasets import load_dataset
# from unsloth_zoo.tokenizer_utils import add_new_tokens
# from unsloth_extensions import add_new_tokens
import lejepa
import torch
from transformers import Adafactor, get_cosine_schedule_with_warmup  # type: ignore

from torch.utils.data import DataLoader

from data_utils import EncodedIterableDataset, encode_one_example, make_lemath_collate_fn
from loss_utils import compute_cos_loss, compute_sigreg_loss

def _is_oom_error(e: BaseException) -> bool:
    s = str(e).lower()
    return ("out of memory" in s) and ("cuda" in s or "cublas" in s or "hip" in s)


def _set_seed(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def _now_ts() -> str:
    return time.strftime("%Y-%m-%d %H:%M:%S")

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="LeMath training (LeJEPA-style finetune).")
    # Model
    p.add_argument("--model_name_or_path", type=str, default="Qwen/Qwen3-14B")
    p.add_argument("--max_seq_len", type=int, default=10240)
    p.add_argument(
        "--load_in_4bit",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Load model in 4-bit. Recommended for Qwen3-14B; training will use LoRA adapters.",
    )
    p.add_argument(
        "--use_lora",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Use LoRA adapters for training (recommended; required for 4-bit).",
    )
    p.add_argument("--lora_r", type=int, default=64)
    p.add_argument("--lora_alpha", type=int, default=128)
    p.add_argument("--lora_dropout", type=float, default=0.0)
    p.add_argument(
        "--lora_target_modules",
        type=str,
        default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
        help="Comma-separated module names to apply LoRA to.",
    )
    p.add_argument("--trust_remote_code", action="store_true")
    p.add_argument("--gradient_checkpointing", action="store_true")

    # Batching
    p.add_argument("--microbatch", type=int, default=2, help="Per-step microbatch size (before grad accumulation).")

    # Dataset
    p.add_argument("--dataset_name", type=str, default="nvidia/OpenMathReasoning")
    p.add_argument("--dataset_split", type=str, default="cot")
    p.add_argument("--streaming", action=argparse.BooleanOptionalAction, default=True)
    p.add_argument("--shuffle_buffer", type=int, default=10_000)
    p.add_argument("--num_workers", type=int, default=0)

    # Formatting
    p.add_argument("--chat_format", action=argparse.BooleanOptionalAction, default=True, help="Use tokenizer chat template formatting.")
    p.add_argument(
        "--system_prompt",
        type=str,
        default=(
            "You are a helpful math reasoning assistant.\n"
            "Solve the user's problem.\n"
            "Put the final answer in \\\\boxed{...}."
        ),
    )
    # The user message should be *only* the problem statement by default.
    p.add_argument("--prompt_template", type=str, default="{problem}")
    p.add_argument("--jeton_token", type=str, default="<|jeton|>")
    p.add_argument(
        "--tokens_per_jeton",
        type=int,
        default=8,
        help="Inject 1 jeton per this many tokens inside the <think> region (covers remainder chunk).",
    )

    # Loss weights
    p.add_argument("--w_sigreg", type=float, default=0.02)
    p.add_argument("--w_cos", type=float, default=1.0)
    p.add_argument("--w_decoding", type=float, default=1.0)

    # SIGReg config (lejepa)
    p.add_argument("--sigreg_num_points", type=int, default=17)
    p.add_argument("--sigreg_num_slices", type=int, default=1024)

    # Optim / schedule
    p.add_argument("--optimizer", type=str, default="adafactor", choices=["adafactor"])
    p.add_argument("--learning_rate", type=float, default=5e-5)
    p.add_argument("--weight_decay", type=float, default=0.0)
    p.add_argument("--warmup_steps", type=int, default=200)
    p.add_argument("--max_steps", type=int, default=1000)
    p.add_argument("--grad_accum", type=int, default=1)
    p.add_argument("--max_grad_norm", type=float, default=1.0)

    # Precision
    p.add_argument("--bf16", action=argparse.BooleanOptionalAction, default=True)
    # Kept only for backwards-compat; float16 is intentionally unsupported.
    p.add_argument(
        "--fp16",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="(deprecated) float16 is not supported; use --bf16 (or --no-bf16 for fp32).",
    )

    # IO / logging
    p.add_argument("--output_dir", type=str, default="./out")
    p.add_argument("--save_every", type=int, default=1000)
    p.add_argument("--log_every", type=int, default=10)
    p.add_argument(
        "--wandb_project",
        type=str,
        default="",
        help="If set, log training metrics to Weights & Biases under this project name.",
    )
    p.add_argument(
        "--wandb-entity",
        type=str,
        default="",
        help="Optional W&B entity/team. If empty, uses your default entity.",
    )
    p.add_argument(
        "--wandb-run-id",
        type=str,
        default="",
        help="Optional W&B run id to resume/attach to (e.g. 2ii557i5). Uses resume='allow'.",
    )
    p.add_argument(
        "--wandb-name",
        type=str,
        default="",
        help="Optional W&B run display name.",
    )
    p.add_argument(
        "--wandb-log-checkpoints",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="If true, upload checkpoints as W&B Artifacts (only when --wandb-project is set).",
    )
    p.add_argument(
        "--wandb-artifact-name",
        type=str,
        default="checkpoint",
        help="Artifact base name (actual artifact name is '<base>-<run_id>' to avoid collisions).",
    )
    p.add_argument("--seed", type=int, default=1234)
    p.add_argument("--train_on_last_pair", action=argparse.BooleanOptionalAction, default=True)
    return p.parse_args()


def main() -> None:
    args = parse_args()
    _set_seed(args.seed)

    os.makedirs(args.output_dir, exist_ok=True)
    if args.fp16:
        raise ValueError("--fp16 is no longer supported. Use --bf16 (recommended) or --no-bf16 for fp32.")

    if args.load_in_4bit and not args.use_lora:
        print(f"[{_now_ts()}] WARNING: --load_in_4bit requires adapters. Forcing --use_lora.", flush=True)
        args.use_lora = True

    # Tokens we want to inject / track during training.
    special_tokens = [args.jeton_token]

    # Environment hint (requested venv: /venv/main)
    expected_venv = os.path.realpath("/venv/main")
    active_venv = os.environ.get("VIRTUAL_ENV")
    if active_venv is None:
        print(f"[{_now_ts()}] WARNING: VIRTUAL_ENV is not set. Recommended to run with `{expected_venv}/bin/python`.", flush=True)
    else:
        if os.path.realpath(active_venv) != expected_venv:
            print(
                f"[{_now_ts()}] WARNING: active venv is `{os.path.realpath(active_venv)}`; "
                f"recommended venv is `{expected_venv}`.",
                flush=True,
            )

    # ---- Optional W&B logging (only if --wandb_project is provided)
    wandb = None  # type: ignore
    wandb_run = None
    if args.wandb_project:
        try:
            import wandb as _wandb  # type: ignore
        except Exception as e:
            raise RuntimeError(
                "Failed to import wandb. Install it (e.g. `pip install wandb`) or run without --wandb_project."
            ) from e
        wandb = _wandb
        wandb_init_kwargs: Dict[str, Any] = {"project": args.wandb_project, "config": vars(args)}
        if args.wandb_entity:
            wandb_init_kwargs["entity"] = args.wandb_entity
        if args.wandb_run_id:
            wandb_init_kwargs["id"] = args.wandb_run_id
            wandb_init_kwargs["resume"] = "allow"
        if args.wandb_name:
            wandb_init_kwargs["name"] = args.wandb_name
        wandb_run = wandb.init(**wandb_init_kwargs)

    # ---- Load model/tokenizer via Unsloth (preferred), fallback to HF otherwise.

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model_name_or_path,
        max_seq_length=args.max_seq_len,
        dtype=torch.bfloat16 if args.bf16 else None,
        load_in_4bit=args.load_in_4bit,
        trust_remote_code=args.trust_remote_code,
    )
    print(f"[{_now_ts()}] Adding new tokens: {special_tokens}")
    # add_new_tokens(model, tokenizer, new_tokens=special_tokens)
    tokenizer.add_tokens(special_tokens)

    if args.use_lora:
        target_modules = [m.strip() for m in args.lora_target_modules.split(",") if m.strip()]
        model = FastLanguageModel.get_peft_model(
            model,
            r=args.lora_r,
            target_modules=target_modules,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            bias="none",
            use_gradient_checkpointing=args.gradient_checkpointing,
            random_state=args.seed,
        )

    # Make sure we can train:
    FastLanguageModel.for_training(model, use_gradient_checkpointing=args.gradient_checkpointing)

    if tokenizer.pad_token is None:
        # Qwen-style tokenizers often don't have a pad token.
        tokenizer.pad_token = tokenizer.eos_token

    jeton_id = tokenizer.convert_tokens_to_ids(args.jeton_token)
    if jeton_id == tokenizer.unk_token_id:
        raise RuntimeError("Failed to register jeton token with tokenizer.")
    think_open_ids = tokenizer("<think>", add_special_tokens=False).input_ids
    think_close_ids = tokenizer("</think>", add_special_tokens=False).input_ids
    if not think_open_ids or not think_close_ids:
        raise RuntimeError("Failed to tokenize <think> / </think> markers.")
    if len(think_open_ids) != 1 or len(think_close_ids) != 1:
        raise RuntimeError(
            "Expected <think> and </think> to be single tokens. "
            "Please ensure they are registered as special tokens in the tokenizer."
        )
    think_open_id = int(think_open_ids[0])
    think_close_id = int(think_close_ids[0])

    # Determine device from model (important for quantized / device_map='auto').
    model_param_device = None
    for p in model.parameters():
        model_param_device = p.device
        break
    device = model_param_device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # Quantized / device-mapped models should not be moved manually.
    if not args.load_in_4bit:
        model.to(device)

    if args.load_in_4bit and not args.use_lora:
        raise RuntimeError("--load_in_4bit requires --use_lora in this training script.")

    # Autocast setup (used by both preflight + training).
    use_cuda = device.type == "cuda"
    use_autocast = use_cuda and args.bf16
    amp_dtype = torch.bfloat16

    # ---- SIGReg loss function (LeJEPA) if installed, else fallback.
    # LeJEPA's EppsPulley stores some tensors as plain attributes (not registered buffers),
    # so a bare `.to(device)` may not move everything. We patch tensor attributes manually.
    univariate_test = lejepa.univariate.EppsPulley(n_points=args.sigreg_num_points).to(device)
    sigreg_fn = lejepa.multivariate.SlicingUnivariateTest(
            univariate_test=univariate_test,
            num_slices=args.sigreg_num_slices,
        )

    ds = load_dataset(args.dataset_name, split=args.dataset_split, streaming=args.streaming)
    if args.streaming and args.shuffle_buffer > 0:
        ds = ds.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed)

    def _encode(ex: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        return encode_one_example(
            ex=ex,
            prompt_template=args.prompt_template,
            jeton_token=args.jeton_token,
            jeton_id=jeton_id,
            tokens_per_jeton=args.tokens_per_jeton,
            max_seq_len=args.max_seq_len,
            think_open_id=think_open_id,
            think_close_id=think_close_id,
            tokenizer=tokenizer,
            system_prompt=args.system_prompt,
        )

    # We train any trainable model params (LoRA if enabled).
    trainable_params = [p for p in model.parameters() if p.requires_grad]
    total_param_count = sum(p.numel() for p in model.parameters())
    trainable_param_count = sum(p.numel() for p in trainable_params)
    pct_trainable = (100.0 * trainable_param_count / total_param_count) if total_param_count > 0 else 0.0
    print(
        f"[{_now_ts()}] params: trainable={trainable_param_count:,} / total={total_param_count:,} ({pct_trainable:.4f}%)",
        flush=True,
    )
    if len(trainable_params) == 0:
        raise RuntimeError(
            "Model has 0 trainable parameters. If you loaded quantized weights, you likely need adapters; "
            "this script is written for full fine-tuning."
        )
    opt = Adafactor(
        trainable_params,
        lr=args.learning_rate,
        scale_parameter=False,
        relative_step=False,
        warmup_init=False,
        weight_decay=args.weight_decay,
        decay_rate=-1.2
    )
    sched = get_cosine_schedule_with_warmup(opt, num_warmup_steps=args.warmup_steps, num_training_steps=args.max_steps)

    # Autocast:
    # - bf16 autocast only (no float16)
    model.train()

    # ---- Training loop
    step = 0
    running = {"loss": 0.0, "decoding_loss": 0.0, "sigreg": 0.0, "cos": 0.0}
    last_log = time.time()

    def _wandb_log_checkpoint(ckpt_dir: str, tag: str, step_for_log: Optional[int]) -> None:
        if wandb_run is None or wandb is None or (not args.wandb_log_checkpoints):
            return
        try:
            # Keep the artifact name unique per-run, but versioned across checkpoints.
            artifact_name = f"{args.wandb_artifact_name}-{wandb_run.id}"
            artifact = wandb.Artifact(
                name=artifact_name,
                type="model",
                metadata={
                    "tag": tag,
                    "step": int(step_for_log) if step_for_log is not None else None,
                    "output_dir": args.output_dir,
                },
            )
            artifact.add_dir(ckpt_dir)
            aliases = ["latest", tag]
            wandb_run.log_artifact(artifact, aliases=aliases)
        except Exception as e:
            print(f"[{_now_ts()}] WARNING: failed to log checkpoint to W&B: {e}", flush=True)

    def save_checkpoint(tag: str, *, step_for_log: Optional[int] = None) -> None:
        print(f"[{_now_ts()}] Saving checkpoint to {args.output_dir}/{tag}")
        out = os.path.join(args.output_dir, tag)
        os.makedirs(out, exist_ok=True)
        # Save model + tokenizer in HF format. (Unsloth model is still HF-compatible.)
        model.save_pretrained(out)
        tokenizer.save_pretrained(out)
        _wandb_log_checkpoint(out, tag=tag, step_for_log=step_for_log)

    microbatch = int(args.microbatch)
    # Torch DataLoader over encoded streaming examples.
    torch_ds = EncodedIterableDataset(ds, encode_fn=_encode)
    dl = DataLoader(
        torch_ds,
        batch_size=microbatch,
        collate_fn=make_lemath_collate_fn(pad_id=int(tokenizer.pad_token_id)),
        num_workers=int(args.num_workers),
        pin_memory=(device.type == "cuda"),
        persistent_workers=(int(args.num_workers) > 0),
    )
    data_it = iter(dl)

    while step < args.max_steps:
        opt.zero_grad(set_to_none=True)

        for ga in range(args.grad_accum):
            try:
                batch = next(data_it)
            except StopIteration:
                data_it = iter(dl)
                batch = next(data_it)

            base_input_ids = batch["base_input_ids"].to(device)
            base_attention_mask = batch["base_attention_mask"].to(device)
            input_ids = batch["input_ids"].to(device)
            grad_attention_mask = batch["grad_attention_mask"].to(device)
            labels = batch["labels"].to(device)

            # Two-pass forward:
            #  - base pass: no grad, up to </think>, with original thinking tokens
            #  - grad pass: with grad, think content replaced by jetons; NLL ignores jetons
            ctx = torch.autocast(device_type="cuda", dtype=amp_dtype) if use_autocast else contextlib.nullcontext()
            with ctx:
                if args.w_cos > 0:
                    with torch.no_grad():
                        base_outputs = model(
                            input_ids=base_input_ids,
                            attention_mask=base_attention_mask,
                            output_hidden_states=True,
                            use_cache=False,
                        )
                        base_last_hidden = base_outputs.hidden_states[-1]  # [B, Tb, H]
                grad_outputs = model(
                    input_ids=input_ids,
                    attention_mask=grad_attention_mask,
                    labels=labels,
                    output_hidden_states=True,
                    use_cache=False,
                )

            decoding_loss = grad_outputs.loss
            last_hidden = grad_outputs.hidden_states[-1]  # [B, Tg, H]

            if args.w_cos > 0:
                cos_loss = compute_cos_loss(
                    base_last_hidden=base_last_hidden,
                    last_hidden=last_hidden,
                    jeton_pairs=batch["cross_pairs"],
                    train_on_last_pair=args.train_on_last_pair,
                )
            else:
                cos_loss = torch.tensor(0.0, device=device)
            if args.w_sigreg > 0:
                sigreg_loss = compute_sigreg_loss(
                    last_hidden=last_hidden,
                    jeton_pairs=batch["cross_pairs"],
                    sigreg_fn=sigreg_fn,
                )
            else:
                sigreg_loss = torch.tensor(0.0, device=device)

            loss = args.w_decoding * decoding_loss + args.w_sigreg * sigreg_loss + args.w_cos * cos_loss
            loss = loss / args.grad_accum
            loss.backward()

            running["loss"] += loss.detach()
            running["decoding_loss"] += decoding_loss.detach()
            running["sigreg"] += sigreg_loss.detach()
            running["cos"] += cos_loss.detach()

        # Step
        torch.nn.utils.clip_grad_norm_(trainable_params, args.max_grad_norm)
        opt.step()
        sched.step()

        step += 1

        if step % args.log_every == 0:
            dt = time.time() - last_log
            last_log = time.time()
            denom = args.log_every
            metrics = {
                "loss": float(running["loss"].cpu()) / denom,
                "decoding_loss": float(running["decoding_loss"].cpu()) / denom,
                "sigreg": float(running["sigreg"].cpu()) / denom,
                "cos": float(running["cos"].cpu()) / denom,
                "lr": sched.get_last_lr()[0],
                "secs_per_step": (dt / denom) if denom > 0 else float("nan"),
                "steps_per_sec": (denom / dt) if dt > 0 else 0.0,
            }
            msg = (
                f"[{_now_ts()}] step {step}/{args.max_steps} "
                f"loss={metrics['loss']:.4f} "
                f"decoding_loss={metrics['decoding_loss']:.4f} "
                f"sigreg={metrics['sigreg']:.4f} "
                f"cos={metrics['cos']:.4f} "
                f"num_jetons={sum(batch['num_jetons'])} "
                f"lr={metrics['lr']:.2e} "
                f"({dt:.1f}s/{args.log_every} steps)"
            )
            print(msg, flush=True)
            if wandb_run is not None:
                wandb.log(metrics, step=step)
            for k in running:
                running[k] = 0.0

        if args.save_every > 0 and step % args.save_every == 0:
            save_checkpoint(f"step-{step}", step_for_log=step)

    save_checkpoint("final", step_for_log=step)
    print(f"[{_now_ts()}] done. saved to {args.output_dir}")
    if wandb_run is not None:
        wandb.finish()


if __name__ == "__main__":
    main()