"""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()