src/diffusion_lm/reasoning_train.py

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