source/src/strata/training/loop.py

"""Pure-DDP causal-LM training loop for STRATA.

Constraints honoured here:

* Plain PyTorch + ``DistributedDataParallel`` only. No DeepSpeed, no ZeRO, no
  CPU/parameter offload.
* Reproducible: seeds, resolved configs, and per-step metrics are recorded; a
  run directory is never silently overwritten.
* Resumable: checkpoints carry model + optimizer + step + RNG state.

The loop is deliberately framework-light so it also runs single-process on CPU
(used by the unit tests) and single-/multi-GPU under ``torchrun``.
"""

from __future__ import annotations

import json
import logging
import math
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import torch
from torch import nn

from strata.modeling.config import StrataConfig
from strata.modeling.graph_object import GRAPH_OBJECT_INTERVENTIONS
from strata.modeling.model import StrataForCausalLM
from strata.training.distributed import DistInfo
from strata.training.lm_data import PackedLMDataset

LOGGER = logging.getLogger("strata.train")

_PRECISIONS = {"bf16", "fp16", "fp32"}


@dataclass(frozen=True, slots=True)
class TrainConfig:
    """Optimisation and runtime hyperparameters (no machine-specific paths)."""

    seed: int = 1234
    seq_len: int = 512
    micro_batch_size: int = 16
    grad_accum: int = 1
    max_steps: int = 100
    lr: float = 6e-4
    min_lr: float = 6e-5
    warmup_steps: int = 10
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0
    precision: str = "bf16"
    log_every: int = 10
    ckpt_every: int = 0  # 0 disables intra-run checkpoints (only final is saved)
    num_workers: int = 2
    compile: bool = False
    # STRATA's auxiliary predicate/graph heads are computed and returned by the
    # forward pass but do not contribute to the pure-LM loss. DDP therefore sees
    # parameters that are reachable from a forward output yet never receive a
    # gradient. static_graph=True records the real backward graph (the LM path,
    # incl. the tied embedding used twice) on the first step and treats the rest
    # as consistently unused, which is both correct and efficient here.
    ddp_static_graph: bool = True
    # Only consulted when ddp_static_graph is False.
    find_unused_parameters: bool = True
    # UD graph-supervision loss weights (used by the graph training loop). LM loss
    # always stays at weight 1.0; these are conservative by default.
    lambda_ud_arc: float = 0.1
    lambda_ud_rel: float = 0.1
    lambda_node_type: float = 0.05
    ud_max_len: int = 256
    # SRL (predicate-argument) supervision weights, used by the SRL training loop.
    lambda_srl_predicate: float = 0.1
    lambda_srl_role: float = 0.1
    # Down-weight the dominant NONE (no-edge) class in the SRL role loss.
    srl_none_weight: float = 0.2
    # Optional targeted CE over gold ARG0/ARG1 cells. Default is zero so legacy
    # runs/configs are unchanged; enable only after semantic ablations indicate
    # weak core-role causality.
    lambda_srl_core_role: float = 0.0
    # Chart supervision (per-token chart-item type) weight.
    lambda_chart_type: float = 0.05
    # Explicit GraphObjectMemory auxiliary objective over edge slots. This is
    # separate from UD/SRL/chart prediction heads and is only used by graph-object
    # replacement experiments.
    lambda_graph_object_aux: float = 0.0
    # GraphObjectMemory v6 same-topology ARG0/ARG1 counterfactual objective. This
    # scores RoleEdit examples with learned predicate-memory readback disabled,
    # so only the explicit graph object can distinguish original vs edited role
    # labels. Stability keeps unrelated probes from flipping under the edit.
    lambda_graph_object_role_edit: float = 0.0
    lambda_graph_object_role_edit_stability: float = 0.0
    lambda_graph_object_role_edit_untyped_null: float = 0.0
    lambda_graph_object_role_edit_random_null: float = 0.0
    # Counterfactual ARG0/ARG1 edit-control objective. Defaults are inert so all
    # legacy LM/mixed configs keep identical behavior unless explicitly enabled.
    lambda_role_edit: float = 0.0
    lambda_role_edit_specificity: float = 0.0
    lambda_role_edit_scope: float = 0.0
    role_edit_batch_size: int = 1
    role_edit_margin: float = 0.2
    role_edit_scope_tolerance: float = 0.1
    role_edit_bias_value: float = 4.0
    role_edit_override_scope: str = "legacy"
    # Explicit graph-object memory source used by graph-object-enabled mixed
    # training arms. In legacy configs this stays "none"; when the model config
    # has use_graph_object_memory=true it selects typed/untyped/random object
    # content at matched tensor/edge-slot budget.
    graph_object_intervention: str = "none"
    # GraphObjectMemory v2 forced-use knob. During mixed training this suppresses
    # the learned predicate-memory residual on graph examples in the full-graph
    # pass only. Pure LM examples and the causal LM pass keep the learned memory
    # path intact so broad LM preservation remains measured.
    graph_object_predicate_memory_dropout: float = 0.0
    # GraphObjectMemory v4 annealed readback mixture. At 1.0 the graph-object
    # path uses legacy residual strength (no deterministic replacement). Values
    # below 1.0 linearly anneal the full-graph readback on graph rows from
    # learned=1/object=0 to learned=value/object=(1-value).
    graph_object_anneal_final_learned_weight: float = 1.0
    # Fixed readback mixture for graph-object replacement experiments. When set,
    # it overrides the anneal schedule and uses this learned-memory weight on
    # graph rows for the full run.
    graph_object_readback_learned_weight: float | None = None
    # v9 intervention-controller mode. When true, mixed training freezes every
    # non-GraphObject parameter after warm-start so quartet control cannot move
    # the canonical observational LM/graph path.
    freeze_non_graph_object_parameters: bool = False

    def validate(self) -> None:
        if self.precision not in _PRECISIONS:
            raise ValueError(f"precision must be one of {sorted(_PRECISIONS)}")
        for name in ("seq_len", "micro_batch_size", "grad_accum", "max_steps"):
            if getattr(self, name) <= 0:
                raise ValueError(f"{name} must be positive")
        if self.warmup_steps < 0 or self.warmup_steps > self.max_steps:
            raise ValueError("warmup_steps must be in [0, max_steps]")
        if self.lr <= 0 or self.min_lr < 0 or self.min_lr > self.lr:
            raise ValueError("require 0 < min_lr <= lr")
        if self.srl_none_weight < 0:
            raise ValueError("srl_none_weight must be non-negative")
        if self.lambda_srl_core_role < 0:
            raise ValueError("lambda_srl_core_role must be non-negative")
        if self.lambda_role_edit < 0 or self.lambda_role_edit_specificity < 0:
            raise ValueError("role-edit loss weights must be non-negative")
        if self.lambda_role_edit_scope < 0:
            raise ValueError("lambda_role_edit_scope must be non-negative")
        if self.lambda_graph_object_aux < 0:
            raise ValueError("lambda_graph_object_aux must be non-negative")
        if (
            self.lambda_graph_object_role_edit < 0
            or self.lambda_graph_object_role_edit_stability < 0
            or self.lambda_graph_object_role_edit_untyped_null < 0
            or self.lambda_graph_object_role_edit_random_null < 0
        ):
            raise ValueError("graph-object role-edit loss weights must be non-negative")
        if self.role_edit_batch_size <= 0:
            raise ValueError("role_edit_batch_size must be positive")
        if self.role_edit_margin < 0:
            raise ValueError("role_edit_margin must be non-negative")
        if self.role_edit_scope_tolerance < 0:
            raise ValueError("role_edit_scope_tolerance must be non-negative")
        if self.role_edit_override_scope not in {"legacy", "role_only"}:
            raise ValueError("role_edit_override_scope must be 'legacy' or 'role_only'")
        if self.graph_object_intervention not in GRAPH_OBJECT_INTERVENTIONS:
            raise ValueError(f"unknown graph_object_intervention {self.graph_object_intervention!r}")
        if not 0.0 <= self.graph_object_predicate_memory_dropout < 1.0:
            raise ValueError("graph_object_predicate_memory_dropout must be in [0, 1)")
        if not 0.0 <= self.graph_object_anneal_final_learned_weight <= 1.0:
            raise ValueError("graph_object_anneal_final_learned_weight must be in [0, 1]")
        if self.graph_object_readback_learned_weight is not None and not 0.0 <= self.graph_object_readback_learned_weight <= 1.0:
            raise ValueError("graph_object_readback_learned_weight must be in [0, 1]")

    def to_dict(self) -> dict[str, object]:
        return asdict(self)


def lr_at_step(step: int, cfg: TrainConfig) -> float:
    """Linear warmup then cosine decay from ``lr`` to ``min_lr``."""

    if step < cfg.warmup_steps:
        return cfg.lr * (step + 1) / max(1, cfg.warmup_steps)
    if step >= cfg.max_steps:
        return cfg.min_lr
    span = max(1, cfg.max_steps - cfg.warmup_steps)
    progress = (step - cfg.warmup_steps) / span
    return cfg.min_lr + 0.5 * (cfg.lr - cfg.min_lr) * (1.0 + math.cos(math.pi * progress))


def build_param_groups(model: nn.Module, weight_decay: float) -> list[dict]:
    """Weight-decay matmul/embedding tensors (ndim>=2); exclude biases/norms."""

    decay, no_decay = [], []
    for _, param in model.named_parameters():
        if not param.requires_grad:
            continue
        (decay if param.ndim >= 2 else no_decay).append(param)
    return [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]


def _autocast_context(precision: str, device_type: str):
    if precision == "fp32" or device_type == "cpu":
        return torch.autocast(device_type=device_type, enabled=False)
    dtype = torch.bfloat16 if precision == "bf16" else torch.float16
    return torch.autocast(device_type=device_type, dtype=dtype)


def _infinite_batches(loader, sampler):
    epoch = 0
    while True:
        if sampler is not None and hasattr(sampler, "set_epoch"):
            sampler.set_epoch(epoch)
        for batch in loader:
            yield batch
        epoch += 1


def _all_reduce_mean(value: float, info: DistInfo, device: torch.device) -> float:
    if not info.is_distributed:
        return value
    import torch.distributed as dist

    tensor = torch.tensor([value], device=device, dtype=torch.float32)
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    return (tensor / info.world_size).item()


@dataclass(slots=True)
class TrainResult:
    steps: int
    first_loss: float
    final_loss: float
    best_loss: float
    tokens_seen: int
    checkpoint_dir: str | None


def train_lm(
    *,
    train_cfg: TrainConfig,
    model_cfg: StrataConfig,
    corpus_dir: Path,
    out_dir: Path,
    dist_info: DistInfo,
    resume_from: Path | None = None,
    device: torch.device | None = None,
    logger: logging.Logger | None = None,
) -> TrainResult:
    """Train ``StrataForCausalLM`` on a packed corpus with pure DDP."""

    log = logger or LOGGER
    train_cfg.validate()
    model_cfg.validate()

    if device is None:
        device = torch.device(f"cuda:{dist_info.local_rank}" if torch.cuda.is_available() else "cpu")
    device_type = device.type

    # ---- data ----
    dataset = PackedLMDataset(corpus_dir, seq_len=train_cfg.seq_len)
    if dataset.meta.vocab_size != model_cfg.vocab_size:
        raise ValueError(
            f"corpus vocab_size {dataset.meta.vocab_size} != model vocab_size "
            f"{model_cfg.vocab_size}; regenerate the corpus or fix the model config"
        )
    sampler = None
    if dist_info.is_distributed:
        from torch.utils.data.distributed import DistributedSampler

        sampler = DistributedSampler(
            dataset,
            num_replicas=dist_info.world_size,
            rank=dist_info.rank,
            shuffle=True,
            seed=train_cfg.seed,
            drop_last=True,
        )
    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=train_cfg.micro_batch_size,
        sampler=sampler,
        shuffle=(sampler is None),
        num_workers=train_cfg.num_workers,
        pin_memory=(device_type == "cuda"),
        drop_last=True,
    )
    batches = _infinite_batches(loader, sampler)

    # ---- model ----
    model = StrataForCausalLM(model_cfg).to(device)
    raw_model = model
    if train_cfg.compile:
        model = torch.compile(model)  # type: ignore[assignment]
    if dist_info.is_distributed:
        from torch.nn.parallel import DistributedDataParallel as DDP

        ddp_kwargs: dict = {
            "device_ids": [dist_info.local_rank] if device_type == "cuda" else None,
            "output_device": dist_info.local_rank if device_type == "cuda" else None,
            "gradient_as_bucket_view": True,
        }
        if train_cfg.ddp_static_graph:
            ddp_kwargs["static_graph"] = True
        else:
            ddp_kwargs["find_unused_parameters"] = train_cfg.find_unused_parameters
        model = DDP(model, **ddp_kwargs)

    optimizer = torch.optim.AdamW(
        build_param_groups(raw_model, train_cfg.weight_decay),
        lr=train_cfg.lr,
        betas=(train_cfg.beta1, train_cfg.beta2),
    )
    scaler = torch.amp.GradScaler(device_type, enabled=(train_cfg.precision == "fp16" and device_type == "cuda"))

    start_step = 0
    if resume_from is not None:
        start_step = _load_checkpoint(resume_from, raw_model, optimizer, device, log)

    if dist_info.is_main:
        out_dir.mkdir(parents=True, exist_ok=True)
        metrics_path = out_dir / "metrics.jsonl"
        metrics_handle = metrics_path.open("a", encoding="utf-8")
        num_params = sum(p.numel() for p in raw_model.parameters())
        log.info(
            "model params=%.2fM windows=%d vocab=%d device=%s world=%d precision=%s",
            num_params / 1e6, len(dataset), model_cfg.vocab_size, device, dist_info.world_size, train_cfg.precision,
        )
    else:
        metrics_handle = None

    model.train()
    first_loss = math.nan
    final_loss = math.nan
    best_loss = math.inf
    tokens_per_step = (
        train_cfg.micro_batch_size * train_cfg.grad_accum * train_cfg.seq_len * dist_info.world_size
    )
    tokens_seen = 0
    tokens_at_last_log = 0
    time_at_last_log = time.perf_counter()

    for step in range(start_step, train_cfg.max_steps):
        lr = lr_at_step(step, train_cfg)
        for group in optimizer.param_groups:
            group["lr"] = lr

        optimizer.zero_grad(set_to_none=True)
        accum_loss = 0.0
        for micro in range(train_cfg.grad_accum):
            batch = next(batches).to(device, non_blocking=True)
            is_last_micro = micro == train_cfg.grad_accum - 1
            # Skip the all-reduce on non-final micro-batches to speed up gradient
            # accumulation. Not compatible with static_graph, which manages
            # reduction itself, so fall back to a plain backward there.
            use_no_sync = (
                dist_info.is_distributed
                and not train_cfg.ddp_static_graph
                and not is_last_micro
            )
            sync_context = model.no_sync() if use_no_sync else _nullcontext()
            with sync_context:
                with _autocast_context(train_cfg.precision, device_type):
                    output = model(batch, labels=batch)
                    loss = output.loss / train_cfg.grad_accum
                scaler.scale(loss).backward()
            accum_loss += loss.item()

        if scaler.is_enabled():
            scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(raw_model.parameters(), train_cfg.grad_clip)
        scaler.step(optimizer)
        scaler.update()

        tokens_seen += tokens_per_step
        mean_loss = _all_reduce_mean(accum_loss, dist_info, device)
        if step == start_step:
            first_loss = mean_loss
        final_loss = mean_loss
        best_loss = min(best_loss, mean_loss)

        if dist_info.is_main and (step % train_cfg.log_every == 0 or step == train_cfg.max_steps - 1):
            now = time.perf_counter()
            dt = now - time_at_last_log
            tok_per_s = (tokens_seen - tokens_at_last_log) / dt if dt > 0 else 0.0
            time_at_last_log = now
            tokens_at_last_log = tokens_seen
            record = {
                "step": step,
                "loss": round(mean_loss, 5),
                "lr": lr,
                "grad_norm": round(float(grad_norm), 4),
                "tokens_seen": tokens_seen,
                "tokens_per_s": round(tok_per_s, 1),
            }
            log.info(
                "step %d loss %.4f lr %.2e grad_norm %.3f tok/s %.0f",
                step, mean_loss, lr, float(grad_norm), tok_per_s,
            )
            if metrics_handle is not None:
                metrics_handle.write(json.dumps(record) + "\n")
                metrics_handle.flush()

        if (
            dist_info.is_main
            and train_cfg.ckpt_every > 0
            and step > start_step
            and step % train_cfg.ckpt_every == 0
        ):
            _save_checkpoint(out_dir, step, raw_model, optimizer, train_cfg, log)

    checkpoint_dir = None
    if dist_info.is_main:
        checkpoint_dir = str(_save_checkpoint(out_dir, train_cfg.max_steps, raw_model, optimizer, train_cfg, log))
        if metrics_handle is not None:
            metrics_handle.close()

    return TrainResult(
        steps=train_cfg.max_steps - start_step,
        first_loss=first_loss,
        final_loss=final_loss,
        best_loss=best_loss,
        tokens_seen=tokens_seen,
        checkpoint_dir=checkpoint_dir,
    )


class _nullcontext:
    def __enter__(self):
        return None

    def __exit__(self, *exc):
        return False


def _save_checkpoint(
    out_dir: Path, step: int, raw_model: StrataForCausalLM, optimizer, train_cfg: TrainConfig, log
) -> Path:
    ckpt_dir = out_dir / f"ckpt-{step:07d}"
    raw_model.save_pretrained(ckpt_dir, exist_ok=False)
    torch.save(
        {
            "step": step,
            "optimizer": optimizer.state_dict(),
            "train_config": train_cfg.to_dict(),
            "torch_rng_state": torch.get_rng_state(),
        },
        ckpt_dir / "train_state.pt",
    )
    log.info("saved checkpoint %s", ckpt_dir)
    return ckpt_dir


def _load_checkpoint(resume_from: Path, raw_model: StrataForCausalLM, optimizer, device, log) -> int:
    # train_state.pt holds only tensors + primitives (optimizer state, RNG state,
    # config dict, step), so restrict unpickling to the safe weights-only path.
    state = torch.load(resume_from / "train_state.pt", map_location=device, weights_only=True)
    model_state = torch.load(resume_from / "model.pt", map_location=device, weights_only=True)
    raw_model.load_state_dict(model_state)
    optimizer.load_state_dict(state["optimizer"])
    step = int(state["step"])
    log.info("resumed from %s at step %d", resume_from, step)
    return step