src/diffusion_lm/train.py

"""Single-device training loop for the masked discrete diffusion LM."""

from __future__ import annotations

import argparse
import contextlib
import gc
import hashlib
import json
import math
import random
import shutil
import time
from dataclasses import replace
from pathlib import Path
import numpy as np
import torch
from torch import Tensor
from torch.optim import AdamW, Optimizer
from torch.utils.data import DataLoader

from diffusion_lm.config import ExperimentConfig, ModelConfig, TrainingConfig, load_config
from diffusion_lm.data import DeterministicBatchSampler, load_packed_dataset
from diffusion_lm.diffusion import corrupt_tokens, diffusion_cross_entropy
from diffusion_lm.model import DiffusionTransformer, format_parameter_count
from diffusion_lm.tokenizer import load_tokenizer, special_token_id


def seed_everything(seed: int, device: torch.device) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)
    if device.type == "mps" and hasattr(torch.mps, "manual_seed"):
        # Materialize the lazy MPS runtime before setting its RNG. Seeding before
        # the first allocation can otherwise be overwritten during initialization.
        torch.empty((), device="mps")
        torch.mps.synchronize()
        torch.mps.manual_seed(seed)


def capture_rng_state(device: torch.device) -> dict[str, object]:
    state: dict[str, object] = {
        "python": random.getstate(),
        "numpy": np.random.get_state(),
        "torch": torch.get_rng_state(),
    }
    if device.type == "cuda":
        state["cuda"] = torch.cuda.get_rng_state_all()
    if device.type == "mps" and hasattr(torch.mps, "get_rng_state"):
        state["mps"] = torch.mps.get_rng_state()
    return state


def restore_rng_state(state: dict[str, object]) -> None:
    random.setstate(state["python"])  # type: ignore[arg-type]
    np.random.set_state(state["numpy"])  # type: ignore[arg-type]
    torch.set_rng_state(state["torch"].cpu())  # type: ignore[union-attr]
    if torch.cuda.is_available() and "cuda" in state:
        torch.cuda.set_rng_state_all(  # type: ignore[arg-type]
            [rng_state.cpu() for rng_state in state["cuda"]]  # type: ignore[union-attr]
        )
    if (
        torch.backends.mps.is_available()
        and "mps" in state
        and hasattr(torch.mps, "set_rng_state")
    ):
        torch.mps.set_rng_state(state["mps"].cpu())  # type: ignore[union-attr]


def resolve_device(requested: str) -> torch.device:
    if requested != "auto":
        device = torch.device(requested)
        if device.type == "cuda" and not torch.cuda.is_available():
            raise RuntimeError("CUDA was requested but is unavailable")
        if device.type == "mps" and not torch.backends.mps.is_available():
            raise RuntimeError("MPS was requested but is unavailable")
        return device
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def resolve_precision(requested: str, device: torch.device) -> str:
    if requested != "auto":
        if device.type == "cpu" and requested == "float16":
            raise ValueError("float16 training on CPU is unsupported; use float32 or bfloat16")
        return requested
    if device.type == "cuda":
        return "bfloat16" if torch.cuda.is_bf16_supported() else "float16"
    # Float32 is the most reliable default for CPU and Apple Silicon in this MVP.
    return "float32"


def autocast_context(device: torch.device, precision: str):
    if precision == "float32":
        return contextlib.nullcontext()
    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16}[precision]
    return torch.autocast(device_type=device.type, dtype=dtype)


def configure_cuda_backends(device: torch.device, require_fused_attention: bool) -> None:
    """Enable Ampere-friendly kernels and optionally forbid quadratic math attention."""

    if device.type != "cuda":
        if require_fused_attention:
            raise ValueError("fused attention can only be required on CUDA")
        return

    torch.set_float32_matmul_precision("high")
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
    # MultiheadAttention delegates to scaled_dot_product_attention when it is
    # called with need_weights=False (as TransformerEncoderLayer does). When
    # required, disabling math makes an unsupported fused path fail fast rather
    # than silently allocating an L x L attention matrix. Set it in both branches
    # so repeated train()/benchmark calls in one process cannot leak backend state.
    torch.backends.cuda.enable_math_sdp(not require_fused_attention)


def build_optimizer(model: DiffusionTransformer, config: TrainingConfig) -> Optimizer:
    decay: list[Tensor] = []
    no_decay: list[Tensor] = []
    for parameter in model.parameters():
        if not parameter.requires_grad:
            continue
        (decay if parameter.ndim >= 2 else no_decay).append(parameter)
    parameter_groups = [
        {"params": decay, "weight_decay": config.weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    common = {
        "lr": config.learning_rate,
        "betas": (0.9, 0.95),
        "eps": 1e-8,
    }
    if config.optimizer == "adamw8bit":
        try:
            import bitsandbytes as bnb
        except ImportError as exc:
            raise RuntimeError(
                'optimizer=adamw8bit requires bitsandbytes; install the "gpu" extra'
            ) from exc
        if config.optimizer_embedding_32bit:
            # Bitsandbytes recommends higher-precision optimizer state for NLP
            # embeddings. The LM head is tied to this exact Parameter, so one
            # override protects both roles at a modest memory cost.
            manager = bnb.optim.GlobalOptimManager.get_instance()
            manager.register_module_override(
                model.token_embedding,
                "weight",
                {"optim_bits": 32},
            )
        return bnb.optim.AdamW8bit(
            parameter_groups,
            min_8bit_size=config.optimizer_min_8bit_size,
            **common,
        )
    return AdamW(parameter_groups, foreach=False, **common)


def accumulation_mask_probabilities(
    batch_size: int,
    micro_batch_index: int,
    config: TrainingConfig,
    offset: Tensor,
    device: torch.device,
) -> Tensor:
    """Stratify diffusion times across a complete gradient-accumulation step."""

    slots = config.batch_size * config.gradient_accumulation_steps
    start = micro_batch_index * config.batch_size
    indices = torch.arange(start, start + batch_size, device=device, dtype=torch.float32)
    unit = (offset + indices / max(1, slots)) % 1.0
    return config.mask_eps + (1.0 - config.mask_eps) * unit


def learning_rate(step: int, config: TrainingConfig) -> float:
    if step < config.warmup_steps:
        return config.learning_rate * (step + 1) / max(1, config.warmup_steps)
    progress = (step - config.warmup_steps) / max(1, config.max_steps - config.warmup_steps - 1)
    cosine = 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))
    return config.min_learning_rate + cosine * (
        config.learning_rate - config.min_learning_rate
    )


def create_grad_scaler(enabled: bool):
    """Use the unified API when available and retain PyTorch 2.2 support."""

    unified_scaler = getattr(torch.amp, "GradScaler", None)
    if unified_scaler is not None:
        return unified_scaler("cuda", enabled=enabled)
    return torch.cuda.amp.GradScaler(enabled=enabled)


def validate_inputs(config: ExperimentConfig) -> None:
    tokenizer = load_tokenizer(config.training.tokenizer)
    actual_vocab = tokenizer.get_vocab_size(with_added_tokens=True)
    actual_mask = special_token_id(tokenizer, "mask")
    if actual_vocab != config.model.vocab_size:
        raise ValueError(
            f"config vocab_size is {config.model.vocab_size}, tokenizer has {actual_vocab}; "
            "train the requested tokenizer or update the model budget"
        )
    if actual_mask != config.model.mask_token_id:
        raise ValueError(
            f"config mask_token_id is {config.model.mask_token_id}, tokenizer uses {actual_mask}"
        )
    tokenizer_hash = hashlib.sha256(Path(config.training.tokenizer).read_bytes()).hexdigest()
    for data_path in (config.training.train_data, config.training.val_data):
        if data_path is None:
            continue
        dataset = load_packed_dataset(data_path, config.model.max_seq_len)
        metadata = dataset.metadata
        if int(metadata["vocab_size"]) != config.model.vocab_size:
            raise ValueError(f"{data_path} was encoded with a different vocabulary size")
        if metadata["tokenizer_sha256"] != tokenizer_hash:
            raise ValueError(f"{data_path} was encoded with a different tokenizer file")


@torch.no_grad()
def evaluate(
    model: DiffusionTransformer,
    loader: DataLoader[Tensor],
    device: torch.device,
    precision: str,
    mask_eps: float,
    max_batches: int,
) -> dict[str, float]:
    training_rng_state = capture_rng_state(device)
    was_training = model.training
    try:
        # Fixed corruption masks make validation checkpoints directly comparable.
        # The complete caller RNG state is restored below, so evaluation remains
        # invisible to the subsequent training trajectory.
        seed_everything(0, device)
        model.eval()
        losses: list[float] = []
        correct_weighted = 0.0
        masked_total = 0
        for batch_index, clean_tokens in enumerate(loader):
            if batch_index >= max_batches:
                break
            clean_tokens = clean_tokens.to(device, non_blocking=True)
            # Cover the full noise range deterministically even when evaluation uses
            # microbatch one. Random batch-1 evaluation can otherwise miss the hard
            # near-fully-masked regime for many consecutive checkpoints.
            level = mask_eps + (1.0 - mask_eps) * (batch_index + 0.5) / max_batches
            mask_probability = torch.full(
                (clean_tokens.shape[0],), level, device=device, dtype=torch.float32
            )
            corruption = corrupt_tokens(
                clean_tokens,
                model.config.mask_token_id,
                mask_probability=mask_probability,
                eps=mask_eps,
            )
            with autocast_context(device, precision):
                logits = model(corruption.noisy_tokens, output_positions=corruption.mask)
                output = diffusion_cross_entropy(logits, clean_tokens, corruption)
            losses.append(float(output.loss))
            correct_weighted += float(output.masked_accuracy) * output.masked_tokens
            masked_total += output.masked_tokens
        return {
            "loss": sum(losses) / max(1, len(losses)),
            "masked_accuracy": correct_weighted / max(1, masked_total),
        }
    finally:
        restore_rng_state(training_rng_state)
        model.train(was_training)


def save_checkpoint(
    output_dir: Path,
    model: DiffusionTransformer,
    optimizer: Optimizer,
    scaler,
    experiment: ExperimentConfig,
    step: int,
    tokens_seen: int,
    keep_last_checkpoints: int,
    data_generator: torch.Generator,
    micro_batches_seen: int,
) -> Path:
    output_dir.mkdir(parents=True, exist_ok=True)
    checkpoint = {
        "format": "mini-diffusion-lm-checkpoint-v1",
        "step": step,
        "tokens_seen": tokens_seen,
        "config": experiment.to_dict(),
        "tokenizer_sha256": hashlib.sha256(
            Path(experiment.training.tokenizer).read_bytes()
        ).hexdigest(),
        "rng_state": capture_rng_state(next(model.parameters()).device),
        "data_generator_state": data_generator.get_state(),
        "micro_batches_seen": micro_batches_seen,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scaler": scaler.state_dict(),
    }
    numbered_path = output_dir / f"step-{step:08d}.pt"
    temporary_path = output_dir / ".checkpoint.tmp"
    torch.save(checkpoint, temporary_path)
    temporary_path.replace(numbered_path)

    # Keep `latest.pt` as a hard link when possible so a large optimizer state is
    # not stored twice. The temporary name makes replacement atomic.
    latest_path = output_dir / "latest.pt"
    latest_temporary = output_dir / ".latest.tmp"
    latest_temporary.unlink(missing_ok=True)
    try:
        latest_temporary.hardlink_to(numbered_path)
    except OSError:
        shutil.copyfile(numbered_path, latest_temporary)
    latest_temporary.replace(latest_path)

    if keep_last_checkpoints:
        numbered_checkpoints = sorted(output_dir.glob("step-*.pt"))
        for old_checkpoint in numbered_checkpoints[:-keep_last_checkpoints]:
            old_checkpoint.unlink()
    if experiment.training.save_inference_checkpoint:
        save_inference_checkpoint(output_dir, model, experiment, step, tokens_seen)
    return numbered_path


def _inference_state_dict(model: DiffusionTransformer) -> dict[str, Tensor]:
    """Copy weights to CPU BF16 while preserving tied tensor storage."""

    converted: dict[str, Tensor] = {}
    shared: dict[tuple[object, ...], Tensor] = {}
    for name, tensor in model.state_dict().items():
        key = (
            tensor.untyped_storage().data_ptr(),
            tensor.storage_offset(),
            tuple(tensor.shape),
            tuple(tensor.stride()),
        )
        value = shared.get(key)
        if value is None:
            dtype = torch.bfloat16 if tensor.is_floating_point() else tensor.dtype
            value = tensor.detach().to(device="cpu", dtype=dtype)
            shared[key] = value
        converted[name] = value
    return converted


def save_inference_checkpoint(
    output_dir: Path,
    model: DiffusionTransformer,
    experiment: ExperimentConfig,
    step: int,
    tokens_seen: int,
) -> Path:
    """Write a compact weights-only checkpoint for sampling and the playground."""

    path = output_dir / "inference-latest.pt"
    temporary = output_dir / ".inference.tmp"
    state = _inference_state_dict(model)
    payload = {
        "format": "mini-diffusion-lm-inference-v1",
        "step": step,
        "tokens_seen": tokens_seen,
        "config": experiment.to_dict(),
        "tokenizer_sha256": hashlib.sha256(
            Path(experiment.training.tokenizer).read_bytes()
        ).hexdigest(),
        "model": state,
    }
    torch.save(payload, temporary)
    temporary.replace(path)
    del payload, state
    gc.collect()
    return path


def train(
    experiment: ExperimentConfig,
    resume: str | Path | None = None,
    max_run_steps: int | None = None,
) -> Path:
    config = experiment.training
    if max_run_steps is not None and max_run_steps <= 0:
        raise ValueError("max_run_steps must be positive")
    device = resolve_device(config.device)
    seed_everything(config.seed, device)
    configure_cuda_backends(device, config.require_fused_attention)
    validate_inputs(experiment)

    precision = resolve_precision(config.precision, device)
    train_dataset = load_packed_dataset(config.train_data, experiment.model.max_seq_len)
    data_generator = torch.Generator().manual_seed(config.seed)
    train_batch_sampler = DeterministicBatchSampler(
        len(train_dataset), config.batch_size, seed=config.seed
    )
    train_loader = DataLoader(
        train_dataset,
        batch_sampler=train_batch_sampler,
        num_workers=config.num_workers,
        pin_memory=device.type == "cuda",
        generator=data_generator,
    )

    val_loader = None
    if config.val_data is not None:
        val_dataset = load_packed_dataset(config.val_data, experiment.model.max_seq_len)
        val_loader = DataLoader(
            val_dataset,
            batch_size=config.batch_size,
            shuffle=False,
            num_workers=config.num_workers,
            pin_memory=device.type == "cuda",
        )

    model = DiffusionTransformer(experiment.model).to(device)
    optimizer = build_optimizer(model, config)
    scaler = create_grad_scaler(device.type == "cuda" and precision == "float16")
    start_step = 0
    tokens_seen = 0
    micro_batches_seen = 0
    if resume is not None:
        checkpoint = torch.load(resume, map_location="cpu", weights_only=False)
        if checkpoint.get("format") != "mini-diffusion-lm-checkpoint-v1":
            raise ValueError("unsupported checkpoint format")
        if ModelConfig(**checkpoint["config"]["model"]) != experiment.model:
            raise ValueError("checkpoint model configuration does not match the requested config")
        checkpoint_training = checkpoint["config"].get("training", {})
        checkpoint_optimizer = checkpoint_training.get("optimizer", "adamw")
        if checkpoint_optimizer != config.optimizer:
            raise ValueError(
                f"checkpoint optimizer is {checkpoint_optimizer}, requested {config.optimizer}"
            )
        if config.optimizer == "adamw8bit":
            checkpoint_min_size = int(
                checkpoint_training.get("optimizer_min_8bit_size", 4096)
            )
            checkpoint_embedding_32bit = bool(
                # An absent legacy field means no explicit 32-bit override was
                # guaranteed. Never silently reinterpret it as the safer setting.
                checkpoint_training.get("optimizer_embedding_32bit", False)
            )
            if checkpoint_min_size != config.optimizer_min_8bit_size:
                raise ValueError("checkpoint 8-bit optimizer minimum tensor size does not match")
            if checkpoint_embedding_32bit != config.optimizer_embedding_32bit:
                raise ValueError("checkpoint embedding optimizer precision does not match")
        current_tokenizer_hash = hashlib.sha256(Path(config.tokenizer).read_bytes()).hexdigest()
        if checkpoint.get("tokenizer_sha256") != current_tokenizer_hash:
            raise ValueError("checkpoint was trained with a different tokenizer")
        model.load_state_dict(checkpoint["model"])
        optimizer.load_state_dict(checkpoint["optimizer"])
        optimizer.param_groups[0]["weight_decay"] = config.weight_decay
        optimizer.param_groups[1]["weight_decay"] = 0.0
        for group in optimizer.param_groups:
            group["betas"] = (0.9, 0.95)
            group["eps"] = 1e-8
        scaler.load_state_dict(checkpoint.get("scaler", {}))
        if "data_generator_state" in checkpoint:
            data_generator.set_state(checkpoint["data_generator_state"].cpu())
        if "rng_state" in checkpoint:
            restore_rng_state(checkpoint["rng_state"])
        start_step = int(checkpoint["step"]) + 1
        tokens_seen = int(checkpoint.get("tokens_seen", 0))
        micro_batches_seen = int(
            checkpoint.get(
                "micro_batches_seen",
                start_step * int(checkpoint["config"]["training"]["gradient_accumulation_steps"]),
            )
        )
        del checkpoint
        gc.collect()

    train_batch_sampler.start_batch = micro_batches_seen
    train_iterator = iter(train_loader)

    output_dir = Path(config.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    with (output_dir / "config.json").open("w", encoding="utf-8") as handle:
        json.dump(experiment.to_dict(), handle, indent=2)
        handle.write("\n")

    print(
        json.dumps(
            {
                "event": "start",
                "device": str(device),
                "precision": precision,
                "parameters": model.num_parameters,
                "parameters_human": format_parameter_count(model.num_parameters),
                "optimizer": config.optimizer,
                "activation_checkpointing": experiment.model.activation_checkpointing,
                "fused_attention_required": config.require_fused_attention,
                "training_blocks": len(train_dataset),
                "start_step": start_step,
            }
        )
    )

    last_checkpoint = output_dir / "latest.pt"
    model.train()
    log_started = time.perf_counter()
    log_loss = torch.zeros((), device=device)
    log_accuracy = torch.zeros((), device=device)
    log_tokens = 0
    end_step = config.max_steps
    if max_run_steps is not None:
        end_step = min(end_step, start_step + max_run_steps)
    for step in range(start_step, end_step):
        lr = learning_rate(step, config)
        for group in optimizer.param_groups:
            group["lr"] = lr
        optimizer.zero_grad(set_to_none=True)
        step_loss = torch.zeros((), device=device)
        step_accuracy = torch.zeros((), device=device)
        noise_offset = torch.rand((), device=device)

        for micro_batch_index in range(config.gradient_accumulation_steps):
            clean_tokens = next(train_iterator).to(device, non_blocking=True)
            micro_batches_seen += 1
            mask_probability = accumulation_mask_probabilities(
                clean_tokens.shape[0],
                micro_batch_index,
                config,
                noise_offset,
                device,
            )
            corruption = corrupt_tokens(
                clean_tokens,
                experiment.model.mask_token_id,
                mask_probability=mask_probability,
                eps=config.mask_eps,
            )
            with autocast_context(device, precision):
                logits = model(corruption.noisy_tokens, output_positions=corruption.mask)
                output = diffusion_cross_entropy(logits, clean_tokens, corruption)
                scaled_loss = output.loss / config.gradient_accumulation_steps
            scaler.scale(scaled_loss).backward()
            step_loss += output.loss.detach() / config.gradient_accumulation_steps
            step_accuracy += output.masked_accuracy.detach() / config.gradient_accumulation_steps
            tokens_seen += clean_tokens.numel()
            log_tokens += clean_tokens.numel()

        scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
        scaler.step(optimizer)
        scaler.update()
        log_loss += step_loss
        log_accuracy += step_accuracy

        if (step + 1) % config.log_interval == 0:
            elapsed = time.perf_counter() - log_started
            print(
                json.dumps(
                    {
                        "event": "train",
                        "step": step + 1,
                        "loss": float(log_loss / config.log_interval),
                        "masked_accuracy": float(log_accuracy / config.log_interval),
                        "learning_rate": lr,
                        "grad_norm": float(grad_norm),
                        "tokens_seen": tokens_seen,
                        "tokens_per_second": log_tokens / max(elapsed, 1e-9),
                        "memory_allocated_gib": (
                            torch.cuda.memory_allocated(device) / 1024**3
                            if device.type == "cuda"
                            else 0.0
                        ),
                        "memory_reserved_gib": (
                            torch.cuda.memory_reserved(device) / 1024**3
                            if device.type == "cuda"
                            else 0.0
                        ),
                        "peak_memory_allocated_gib": (
                            torch.cuda.max_memory_allocated(device) / 1024**3
                            if device.type == "cuda"
                            else 0.0
                        ),
                    }
                )
            )
            log_started = time.perf_counter()
            log_loss.zero_()
            log_accuracy.zero_()
            log_tokens = 0

        if val_loader is not None and (step + 1) % config.eval_interval == 0:
            metrics = evaluate(
                model,
                val_loader,
                device,
                precision,
                config.mask_eps,
                config.eval_batches,
            )
            print(json.dumps({"event": "validation", "step": step + 1, **metrics}))

        if (step + 1) % config.save_interval == 0:
            last_checkpoint = save_checkpoint(
                output_dir,
                model,
                optimizer,
                scaler,
                experiment,
                step,
                tokens_seen,
                config.keep_last_checkpoints,
                data_generator,
                micro_batches_seen,
            )
            print(json.dumps({"event": "checkpoint", "path": str(last_checkpoint)}))

    final_step = end_step - 1
    if final_step < start_step:
        raise ValueError("checkpoint step is already at or beyond max_steps")
    if not last_checkpoint.exists() or (final_step + 1) % config.save_interval != 0:
        last_checkpoint = save_checkpoint(
            output_dir,
            model,
            optimizer,
            scaler,
            experiment,
            final_step,
            tokens_seen,
            config.keep_last_checkpoints,
            data_generator,
            micro_batches_seen,
        )
    event = "complete" if end_step == config.max_steps else "paused"
    print(json.dumps({"event": event, "checkpoint": str(last_checkpoint)}))
    return last_checkpoint


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--resume", type=Path)
    parser.add_argument(
        "--max-run-steps",
        type=int,
        help="stop safely after this many optimizer steps (useful for scheduled jobs/tests)",
    )
    parser.add_argument("--device", help="override config device, e.g. cpu, mps, cuda")
    parser.add_argument(
        "--precision",
        choices=("auto", "float32", "bfloat16", "float16"),
        help="override config precision",
    )
    args = parser.parse_args()

    experiment = load_config(args.config)
    if args.device or args.precision:
        training = replace(
            experiment.training,
            device=args.device or experiment.training.device,
            precision=args.precision or experiment.training.precision,
        )
        experiment = replace(experiment, training=training)
    train(experiment, resume=args.resume, max_run_steps=args.max_run_steps)


if __name__ == "__main__":
    main()