"""Training loop for the reasoning variants: ar, diffusion, and hybrid objectives.""" from __future__ import annotations import argparse import hashlib import json import time from dataclasses import dataclass, asdict from pathlib import Path from typing import Any import torch import yaml from torch.utils.data import DataLoader import torch.nn.functional as F from diffusion_lm.config import ModelConfig, TrainingConfig from diffusion_lm.data import DeterministicBatchSampler, load_packed_dataset from diffusion_lm.diffusion import diffusion_cross_entropy from diffusion_lm.hybrid import ( adaptive_hybrid_objective, ar_objective, block_diffusion_objective, block_size_curriculum, diffusion_objective, hybrid_objective, ) from diffusion_lm.model import DiffusionTransformer, build_denoiser, format_parameter_count from diffusion_lm.reasoning import ReasoningTokenDataset from diffusion_lm.tokenizer import load_tokenizer from diffusion_lm.train import ( _inference_state_dict, autocast_context, build_optimizer, capture_rng_state, configure_cuda_backends, create_grad_scaler, learning_rate, resolve_device, resolve_precision, restore_rng_state, seed_everything, ) REASONING_CHECKPOINT_FORMAT = 'mini-diffusion-lm-reasoning-checkpoint-v1' OBJECTIVES = ('ar', 'diffusion', 'hybrid', 'lm', 'block_diffusion') # Objectives trained on continuous packed text (manifest datasets, no region annotations). PACKED_OBJECTIVES = ('lm', 'block_diffusion') @dataclass(frozen=True) class ReasoningConfig: """Objective selection and hybrid-mode hyperparameters.""" objective: str think_probability: float = 0.65 adaptive: bool = False sizes: tuple[int, ...] = () causal_prefix: bool = False curriculum_steps: int = 0 ar_probability: float = 0.0 # Fraction of think tokens corrupted in the causal samples, so the controller learns its # decisions from damaged context instead of only from pristine traces. control_context_noise: float = 0.0 def __post_init__(self) -> None: if self.objective not in OBJECTIVES: raise ValueError(f'objective must be one of {OBJECTIVES}') if not 0.0 < self.think_probability < 1.0: raise ValueError('think_probability must be in (0, 1)') if self.adaptive and self.objective != 'hybrid': raise ValueError('adaptive block sizing is only defined for the hybrid objective') if not isinstance(self.causal_prefix, bool): raise ValueError('causal_prefix must be a boolean') if self.curriculum_steps < 0: raise ValueError('curriculum_steps must be non-negative') if not 0.0 <= self.ar_probability < 1.0: raise ValueError('ar_probability must be in [0, 1)') if not 0.0 <= self.control_context_noise < 1.0: raise ValueError('control_context_noise must be in [0, 1)') sizes = tuple(self.sizes) if self.objective == 'block_diffusion': if not sizes: raise ValueError('block_diffusion requires a non-empty sizes menu') if list(sizes) != sorted(set(sizes)) or sizes[0] <= 0: raise ValueError('sizes must be unique, ascending, and positive') object.__setattr__(self, 'sizes', sizes) @dataclass(frozen=True) class ReasoningExperiment: model: ModelConfig training: TrainingConfig reasoning: ReasoningConfig def to_dict(self) -> dict[str, Any]: return asdict(self) def load_reasoning_config(path: str | Path) -> ReasoningExperiment: with Path(path).open('r', encoding='utf-8') as handle: raw = yaml.safe_load(handle) for section in ('model', 'training', 'reasoning'): if section not in raw: raise ValueError(f'config must contain a top-level {section} mapping') return ReasoningExperiment( model=ModelConfig(**raw['model']), training=TrainingConfig(**raw['training']), reasoning=ReasoningConfig(**raw['reasoning']), ) def _validate_inputs(experiment: ReasoningExperiment, datasets: list) -> None: tokenizer = load_tokenizer(experiment.training.tokenizer) actual_vocab = tokenizer.get_vocab_size(with_added_tokens=True) if experiment.model.backbone == 'project': if actual_vocab != experiment.model.vocab_size: raise ValueError( f'config vocab_size is {experiment.model.vocab_size}, ' f'tokenizer has {actual_vocab}' ) elif actual_vocab > experiment.model.vocab_size: # Pretrained embedding matrices may carry unused tail rows beyond the tokenizer. raise ValueError( f'tokenizer has {actual_vocab} tokens, beyond the {experiment.model.vocab_size} ' f'embedding rows of the pretrained backbone' ) tokenizer_hash = hashlib.sha256( Path(experiment.training.tokenizer).read_bytes() ).hexdigest() for dataset in datasets: if dataset.metadata['tokenizer_sha256'] != tokenizer_hash: raise ValueError(f'{dataset.path} was encoded with a different tokenizer file') if int(dataset.metadata['vocab_size']) != actual_vocab: raise ValueError(f'{dataset.path} was encoded with a different vocabulary size') if 'layout' not in dataset.metadata: continue if experiment.reasoning.objective == 'hybrid': expected_layout = 'adaptive' if experiment.reasoning.adaptive else 'slotted' else: expected_layout = 'flat' if dataset.metadata['layout'] != expected_layout: raise ValueError( f'{dataset.path} has layout {dataset.metadata["layout"]!r}; the ' f'{experiment.reasoning.objective} objective requires {expected_layout!r}' ) if dataset.seq_len != experiment.model.max_seq_len: raise ValueError(f'{dataset.path} sequence length does not match max_seq_len') def _lm_objective( model: DiffusionTransformer, tokens: torch.Tensor ) -> tuple[torch.Tensor, dict[str, float]]: """Plain causal-LM pretraining over continuous packed text.""" seq_len = tokens.shape[1] causal_blocked = torch.triu( torch.ones(seq_len, seq_len, dtype=torch.bool, device=tokens.device), diagonal=1 ) output_positions = torch.ones_like(tokens, dtype=torch.bool) output_positions[:, -1] = False logits = model(tokens, output_positions=output_positions, attn_mask=causal_blocked) targets = tokens[:, 1:].reshape(-1) loss = F.cross_entropy(logits.float(), targets) accuracy = float((logits.argmax(dim=-1) == targets).float().mean()) return loss, {'accuracy': accuracy} def _objective_step( model: DiffusionTransformer, tokens: torch.Tensor, regions: torch.Tensor | None, experiment: ReasoningExperiment, *, eval_mask_level: float | None = None, step: int | None = None, ) -> tuple[torch.Tensor, dict[str, float]]: objective = experiment.reasoning.objective if objective == 'lm': return _lm_objective(model, tokens) if objective == 'block_diffusion': size_weights = block_size_curriculum( step, n_sizes=len(experiment.reasoning.sizes), curriculum_steps=experiment.reasoning.curriculum_steps, ) output = block_diffusion_objective( model, tokens, sizes=experiment.reasoning.sizes, size_weights=size_weights, mask_eps=experiment.training.mask_eps, ar_probability=experiment.reasoning.ar_probability, ) return output.loss, { 'think_loss': output.think_loss, 'answer_loss': output.answer_loss, 'think_accuracy': output.think_accuracy, 'answer_accuracy': output.answer_accuracy, } if objective == 'hybrid': if experiment.reasoning.adaptive: size_ids = torch.tensor(model.adaptive_size_ids, device=tokens.device) output = adaptive_hybrid_objective( model, tokens, regions, size_ids=size_ids, end_think_id=model.adaptive_end_think_id, think_probability=experiment.reasoning.think_probability, mask_eps=experiment.training.mask_eps, causal_prefix=experiment.reasoning.causal_prefix, control_context_noise=experiment.reasoning.control_context_noise, ) else: output = hybrid_objective( model, tokens, regions, block=int(model.reasoning_block), think_probability=experiment.reasoning.think_probability, mask_eps=experiment.training.mask_eps, ) metrics = { 'think_loss': output.think_loss, 'answer_loss': output.answer_loss, 'think_accuracy': output.think_accuracy, 'answer_accuracy': output.answer_accuracy, } if experiment.reasoning.adaptive: metrics['control_accuracy'] = output.control_accuracy metrics['stop_accuracy'] = output.stop_accuracy return output.loss, metrics if objective == 'ar': output = ar_objective(model, tokens, regions) return output.loss, {'accuracy': output.accuracy} mask_probability = None if eval_mask_level is not None: mask_probability = torch.full( (tokens.shape[0],), eval_mask_level, device=tokens.device, dtype=torch.float32 ) logits, corruption, _ = diffusion_objective( model, tokens, regions, mask_eps=experiment.training.mask_eps, mask_probability=mask_probability, ) output = diffusion_cross_entropy(logits, tokens, corruption) return output.loss, {'masked_accuracy': float(output.masked_accuracy)} @torch.no_grad() def _evaluate( model: DiffusionTransformer, loader: DataLoader, experiment: ReasoningExperiment, device: torch.device, precision: str, ) -> dict[str, float]: state = capture_rng_state(device) was_training = model.training try: seed_everything(0, device) model.eval() totals: dict[str, float] = {} batches = 0 max_batches = experiment.training.eval_batches for batch_index, batch in enumerate(loader): if batch_index >= max_batches: break if experiment.reasoning.objective in PACKED_OBJECTIVES: tokens, regions = batch, None else: tokens, regions = batch regions = regions.to(device, non_blocking=True) tokens = tokens.to(device, non_blocking=True) level = experiment.training.mask_eps + (1.0 - experiment.training.mask_eps) * ( (batch_index + 0.5) / max_batches ) with autocast_context(device, precision): loss, metrics = _objective_step( model, tokens, regions, experiment, eval_mask_level=level ) totals['loss'] = totals.get('loss', 0.0) + float(loss) for key, value in metrics.items(): totals[key] = totals.get(key, 0.0) + value batches += 1 return {key: value / max(1, batches) for key, value in totals.items()} finally: restore_rng_state(state) model.train(was_training) def _save_checkpoint( output_dir: Path, model: DiffusionTransformer, optimizer: torch.optim.Optimizer, scaler: Any, experiment: ReasoningExperiment, step: int, micro_batches_seen: int, data_generator: torch.Generator, ) -> Path: output_dir.mkdir(parents=True, exist_ok=True) payload = { 'format': REASONING_CHECKPOINT_FORMAT, 'step': step, 'micro_batches_seen': micro_batches_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(), 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scaler': scaler.state_dict(), } path = output_dir / 'latest.pt' temporary = output_dir / '.checkpoint.tmp' torch.save(payload, temporary) temporary.replace(path) inference_payload = { 'format': 'mini-diffusion-lm-reasoning-inference-v1', 'step': step, 'config': experiment.to_dict(), 'tokenizer_sha256': payload['tokenizer_sha256'], 'model': _inference_state_dict(model), } inference_temporary = output_dir / '.inference.tmp' torch.save(inference_payload, inference_temporary) inference_temporary.replace(output_dir / 'inference-latest.pt') return path def _model_architecture(config: ModelConfig) -> dict[str, Any]: """Config payload compared on resume; pretrained_path is machine-local and may move.""" fields = asdict(config) fields.pop('pretrained_path', None) return fields def _load_init_weights(model: DiffusionTransformer, path: str | Path) -> None: """Initialize model weights from any project checkpoint, ignoring optimizer state. Enables the pretrain-AR-then-adapt recipe: architectures must match tensor-wise, while objective-level settings (forbidden outputs, objective) may differ. """ checkpoint = torch.load(path, map_location='cpu', weights_only=False) state = checkpoint.get('model') if state is None: raise ValueError(f'{path} does not contain model weights') model.load_state_dict(state) def train( experiment: ReasoningExperiment, resume: str | Path | None = None, max_run_steps: int | None = None, init_weights: str | Path | None = None, ) -> Path: config = experiment.training device = resolve_device(config.device) seed_everything(config.seed, device) configure_cuda_backends(device, config.require_fused_attention) precision = resolve_precision(config.precision, device) is_packed = experiment.reasoning.objective in PACKED_OBJECTIVES def _open_dataset(path: str): if is_packed: return load_packed_dataset(path, experiment.model.max_seq_len) return ReasoningTokenDataset(path) train_dataset = _open_dataset(config.train_data) datasets = [train_dataset] val_loader = None if config.val_data is not None: val_dataset = _open_dataset(config.val_data) datasets.append(val_dataset) val_loader = DataLoader( val_dataset, batch_size=config.batch_size, shuffle=False, num_workers=config.num_workers, pin_memory=device.type == 'cuda', ) _validate_inputs(experiment, datasets) data_generator = torch.Generator().manual_seed(config.seed) batch_sampler = DeterministicBatchSampler( len(train_dataset), config.batch_size, seed=config.seed ) train_loader = DataLoader( train_dataset, batch_sampler=batch_sampler, num_workers=config.num_workers, pin_memory=device.type == 'cuda', generator=data_generator, ) load_pretrained = resume is None and init_weights is None # fp32 master weights regardless of the checkpoint's serialized dtype (transformers 5 # defaults to 'auto'/bf16); compute precision comes from autocast like every project run. model = build_denoiser( experiment.model, load_pretrained=load_pretrained, dtype=torch.float32 ).to(device) if not is_packed: model.reasoning_block = int(train_dataset.metadata['block']) if experiment.reasoning.adaptive: metadata = train_dataset.metadata if 'size_token_ids' not in metadata: raise ValueError(f'{config.train_data} lacks size_token_ids; rebuild it adaptively') model.adaptive_size_ids = tuple(sorted(int(v) for v in metadata['size_token_ids'].values())) model.adaptive_end_think_id = int(metadata['reasoning_token_ids']['end_think']) if init_weights is not None: if resume is not None: raise ValueError('init_weights and resume are mutually exclusive') _load_init_weights(model, init_weights) print(json.dumps({'event': 'init_weights', 'path': str(init_weights)})) optimizer = build_optimizer(model, config) scaler = create_grad_scaler(device.type == 'cuda' and precision == 'float16') start_step = 0 micro_batches_seen = 0 if resume is not None: checkpoint = torch.load(resume, map_location='cpu', weights_only=False) if checkpoint.get('format') != REASONING_CHECKPOINT_FORMAT: raise ValueError('unsupported checkpoint format') if _model_architecture( ModelConfig(**checkpoint['config']['model']) ) != _model_architecture(experiment.model): raise ValueError('checkpoint model configuration does not match the config') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) scaler.load_state_dict(checkpoint.get('scaler', {})) data_generator.set_state(checkpoint['data_generator_state'].cpu()) restore_rng_state(checkpoint['rng_state']) start_step = int(checkpoint['step']) + 1 micro_batches_seen = int(checkpoint['micro_batches_seen']) del checkpoint 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', 'objective': experiment.reasoning.objective, 'device': str(device), 'precision': precision, 'parameters': model.num_parameters, 'parameters_human': format_parameter_count(model.num_parameters), 'training_examples': len(train_dataset), 'start_step': start_step, } ) ) model.train() log_started = time.perf_counter() log_loss = 0.0 log_metrics: dict[str, float] = {} log_count = 0 end_step = config.max_steps if max_run_steps is not None: end_step = min(end_step, start_step + max_run_steps) last_checkpoint = output_dir / 'latest.pt' 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) for _ in range(config.gradient_accumulation_steps): batch = next(train_iterator) if is_packed: tokens, regions = batch, None else: tokens, regions = batch regions = regions.to(device, non_blocking=True) tokens = tokens.to(device, non_blocking=True) micro_batches_seen += 1 with autocast_context(device, precision): loss, metrics = _objective_step(model, tokens, regions, experiment, step=step) scaled = loss / config.gradient_accumulation_steps scaler.scale(scaled).backward() log_loss += float(loss) / config.gradient_accumulation_steps for key, value in metrics.items(): log_metrics[key] = ( log_metrics.get(key, 0.0) + value / config.gradient_accumulation_steps ) scaler.unscale_(optimizer) grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip) scaler.step(optimizer) scaler.update() log_count += 1 if (step + 1) % config.log_interval == 0: elapsed = time.perf_counter() - log_started payload = { 'event': 'train', 'step': step + 1, 'loss': log_loss / max(1, log_count), 'learning_rate': lr, 'grad_norm': float(grad_norm), 'steps_per_second': log_count / max(elapsed, 1e-9), } payload.update( {key: value / max(1, log_count) for key, value in log_metrics.items()} ) print(json.dumps(payload)) log_started = time.perf_counter() log_loss = 0.0 log_metrics = {} log_count = 0 if val_loader is not None and (step + 1) % config.eval_interval == 0: metrics = _evaluate(model, val_loader, experiment, device, precision) 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, micro_batches_seen, data_generator, ) print(json.dumps({'event': 'checkpoint', 'path': str(last_checkpoint)})) final_step = end_step - 1 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, micro_batches_seen, data_generator, ) 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('--init-weights', type=Path) parser.add_argument('--max-run-steps', type=int) parser.add_argument('--device') parser.add_argument( '--precision', choices=('auto', 'float32', 'bfloat16', 'float16') ) args = parser.parse_args() experiment = load_reasoning_config(args.config) if args.device or args.precision: from dataclasses import replace training = replace( experiment.training, device=args.device or experiment.training.device, precision=args.precision or experiment.training.precision, ) experiment = ReasoningExperiment( model=experiment.model, training=training, reasoning=experiment.reasoning ) train( experiment, resume=args.resume, max_run_steps=args.max_run_steps, init_weights=args.init_weights, ) if __name__ == '__main__': main()