feat(pr2): grad-accum + checkpoint eval-barrier in the eager loop (memory fix)

apstenku123 · apstenku123 · commit 05e53d8bb697 · 2026-06-02T01:08:45.000+02:00
PR-2 from FUSED-PIPELINE-ROADMAP §4/§5: bring the seq=4096 training-step peak down by fixing the LOOP, not the kernels. - one_step_train: add grad_accum_steps (kwarg or GRAD_ACCUM_STEPS env, default 1 = byte-identical to prior single-eval step). When >1, split the batch into N microbatches, run value_and_grad per micro, mx.eval the token-weighted accumulator as a free-barrier, and drop each micro's forward/backward graph before the next. Peak activation collapses from full-batch to ~one microbatch. Token-weighting makes the accumulated grad numerically equal to the full-batch grad (a token-mean loss => full grad = sum_i (n_i/N) grad_i). Optional clear_cache hook after the step. - hybrid_lm: opt-in mx.eval(hidden_states) free-barrier after each grad-checkpointed layer (CPPMEGA_MLX_CHECKPOINT_EVAL_BARRIER=1, default OFF), so the stored checkpoint input is forced+freed before recompute — converting the previously no-op checkpoint into a real O(sbh) footprint. Numeric equivalence verified on Mac (scratch/pr2_grad_accum_equivalence.py): N=4 micro vs full-batch -> loss abs-diff 0.0, worst grad rel-diff 2.8e-7, post-step param rel-diff 1.2e-7; checkpoint barrier ON vs OFF -> grads bitwise identical. All loop/compiled/gradient/checkpoint tests pass (165). Both levers env-gated, default = current behavior (RULE #1, no silent change).
diff --git a/cppmega_mlx/models/hybrid_lm.py b/cppmega_mlx/models/hybrid_lm.py
@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import math
+import os
 from dataclasses import asdict, dataclass, field
 from typing import (
     TYPE_CHECKING,
@@ -61,6 +62,13 @@
 if TYPE_CHECKING:
     from cppmega_mlx.runtime.path_c_fusion import PathCFusionRegion
 
+# PR-2 (FUSED-PIPELINE-ROADMAP §4): opt-in free-barrier after each
+# grad-checkpointed layer. Default OFF so behavior is unchanged unless
+# CPPMEGA_MLX_CHECKPOINT_EVAL_BARRIER=1 is set.
+_CHECKPOINT_EVAL_BARRIER = (
+    os.environ.get("CPPMEGA_MLX_CHECKPOINT_EVAL_BARRIER") == "1"
+)
+
 HybridBackend = Literal["attention", "mamba3", "moe", "m2rnn", "engram", "concept"]
 HybridBlockModule = (
     CausalSelfAttention
@@ -2288,6 +2296,16 @@ def decoder_hidden_states(
                 else:
                     fn = lambda hs, _layer=layer, _mask=mask: _layer(hs, _mask)
                 hidden_states = nn.utils.checkpoint(layer, fn)(hidden_states)
+                if _CHECKPOINT_EVAL_BARRIER:
+                    # PR-2 (FUSED-PIPELINE-ROADMAP §4): force+free each
+                    # checkpointed layer's stored input before the next layer
+                    # builds its graph, so the recompute working set never
+                    # coexists with the full set of stored inputs. Without this
+                    # barrier MLX's lazy tape keeps every layer's checkpoint
+                    # input live simultaneously and checkpoint is a no-op for
+                    # peak memory. Default OFF; opt in via
+                    # CPPMEGA_MLX_CHECKPOINT_EVAL_BARRIER=1.
+                    mx.eval(hidden_states)
         else:
             attention_layer_idx = 0
             for layer_index, layer in enumerate(self.layers):
diff --git a/cppmega_mlx/training/loop.py b/cppmega_mlx/training/loop.py
@@ -4,22 +4,27 @@
 
 import os
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Callable, Mapping
 
 import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optim
-from mlx.utils import tree_flatten
+from mlx.utils import tree_flatten, tree_map
 
-from cppmega_mlx.data.batch import LMTokenBatch
+from cppmega_mlx.data.batch import LMTokenBatch, ensure_lm_batch
 from cppmega_mlx.training.loss import next_token_cross_entropy
 from cppmega_mlx.training.optimizers import make_adamw
 
 
 LossFn = Callable[[nn.Module, LMTokenBatch | Mapping[str, mx.array] | mx.array], tuple[mx.array, mx.array]]
 
 _STRICT_DTYPE_CONTRACT_ENV = "STRICT_DTYPE_CONTRACT"
+# PR-2 (FUSED-PIPELINE-ROADMAP §4/§5): gradient accumulation in the loop.
+# Default OFF (grad_accum_steps=1) so behavior is byte-identical to the prior
+# single-eval step unless explicitly opted in. ``GRAD_ACCUM_STEPS`` env var
+# provides a process-wide default for callers that do not pass the kwarg.
+_GRAD_ACCUM_STEPS_ENV = "GRAD_ACCUM_STEPS"
 
 
 @dataclass(frozen=True)
@@ -68,32 +73,162 @@ def assert_grad_dtype_matches_param_dtype(grads, params) -> None:
         )
 
 
+def _resolve_grad_accum_steps(grad_accum_steps: int | None) -> int:
+    """Resolve the microbatch count from kwarg then env, default 1 (OFF)."""
+
+    if grad_accum_steps is None:
+        env_value = os.environ.get(_GRAD_ACCUM_STEPS_ENV)
+        grad_accum_steps = int(env_value) if env_value not in (None, "") else 1
+    if not isinstance(grad_accum_steps, int) or grad_accum_steps < 1:
+        raise ValueError(
+            f"grad_accum_steps must be a positive integer, got {grad_accum_steps!r}"
+        )
+    return grad_accum_steps
+
+
+def _microbatch_slices(batch_size: int, num_micro: int) -> list[tuple[int, int]]:
+    """Split ``batch_size`` rows into ``num_micro`` contiguous, near-equal spans.
+
+    Raises if the batch is too small to provide one row per microbatch — a
+    silent no-op split would mask a misconfigured ``grad_accum_steps`` (RULE #1).
+    """
+
+    if num_micro > batch_size:
+        raise ValueError(
+            f"grad_accum_steps={num_micro} exceeds batch size {batch_size}; "
+            "each microbatch needs at least one row"
+        )
+    base, extra = divmod(batch_size, num_micro)
+    spans: list[tuple[int, int]] = []
+    start = 0
+    for i in range(num_micro):
+        length = base + (1 if i < extra else 0)
+        spans.append((start, start + length))
+        start += length
+    return spans
+
+
+def _slice_lm_batch(batch: LMTokenBatch, start: int, stop: int) -> LMTokenBatch:
+    """Return a row-slice ``[start:stop)`` of every per-example array field."""
+
+    updates: dict[str, mx.array] = {}
+    for name, value in vars(batch).items():
+        if isinstance(value, mx.array) and value.shape and value.shape[0] == batch.tokens.shape[0]:
+            updates[name] = value[start:stop]
+    return replace(batch, **updates)
+
+
+def _accumulate_grads_token_weighted(
+    accumulator: dict | None,
+    grads: dict,
+    weight: mx.array,
+) -> dict:
+    """Add ``weight * grads`` into ``accumulator`` (token-weighted running sum).
+
+    Token weighting makes the accumulated gradient numerically equal to the
+    full-batch gradient: the per-microbatch loss is a token-mean, so the
+    full-batch grad is ``sum_i (n_i / N) * grad_i`` by linearity.
+    """
+
+    weighted = tree_map(
+        lambda g: (g * weight).astype(g.dtype) if isinstance(g, mx.array) else g,
+        grads,
+    )
+    if accumulator is None:
+        return weighted
+    return tree_map(
+        lambda acc, g: acc + g if isinstance(acc, mx.array) else acc,
+        accumulator,
+        weighted,
+    )
+
+
 def one_step_train(
     model: nn.Module,
     optimizer: optim.Optimizer,
     batch: LMTokenBatch | Mapping[str, mx.array] | mx.array,
     *,
     loss_fn: LossFn = next_token_cross_entropy,
+    grad_accum_steps: int | None = None,
+    clear_cache: bool = False,
 ) -> TrainStepResult:
-    """Run one eager MLX AdamW-compatible training step."""
+    """Run one eager MLX AdamW-compatible training step.
+
+    ``grad_accum_steps`` (default 1, also settable via the ``GRAD_ACCUM_STEPS``
+    env var) splits the batch into N microbatches and accumulates token-weighted
+    grads with an ``mx.eval(grads)`` free-barrier between microbatches, dropping
+    each microbatch's forward/backward graph before the next is built. This
+    collapses peak activation memory from full-batch to ~one microbatch while
+    remaining numerically equivalent to the full-batch step (FUSED-PIPELINE
+    ROADMAP §4/§5, PR-2). With N=1 the path is byte-identical to the prior step.
+    """
 
     model.train()
-    loss_and_grad = nn.value_and_grad(model, loss_fn)
+    num_micro = _resolve_grad_accum_steps(grad_accum_steps)
 
     start = time.perf_counter()
-    (loss, ntokens), grads = loss_and_grad(model, batch)
+    if num_micro == 1:
+        loss_and_grad = nn.value_and_grad(model, loss_fn)
+        (loss, ntokens), grads = loss_and_grad(model, batch)
+        if _strict_dtype_contract_enabled():
+            assert_grad_dtype_matches_param_dtype(grads, model.parameters())
+        optimizer.update(model, grads)
+        mx.eval(model.parameters(), optimizer.state, loss, ntokens)
+        elapsed = time.perf_counter() - start
+        tokens = int(ntokens.item())
+        if clear_cache:
+            mx.clear_cache()
+        return TrainStepResult(
+            loss=float(loss.item()),
+            ntokens=tokens,
+            seconds=elapsed,
+            tokens_per_second=tokens / elapsed if elapsed > 0 else float("inf"),
+        )
+
+    lm_batch = ensure_lm_batch(batch)
+    batch_size = int(lm_batch.tokens.shape[0])
+    spans = _microbatch_slices(batch_size, num_micro)
+
+    loss_and_grad = nn.value_and_grad(model, loss_fn)
+
+    # Pass 1: total live-token count across the whole batch, used to
+    # token-weight each microbatch grad so the accumulated grad equals the
+    # full-batch grad. ``target_mask`` is pure data (no forward graph), so this
+    # costs nothing in activation memory.
+    total_ntokens = lm_batch.target_mask.sum().astype(mx.float32)
+    mx.eval(total_ntokens)
+    total_tokens_int = int(total_ntokens.item())
+    if total_tokens_int <= 0:
+        raise ValueError("grad accumulation requires at least one live target token")
+
+    accumulator: dict | None = None
+    weighted_loss = mx.array(0.0, dtype=mx.float32)
+    for span_start, span_stop in spans:
+        micro = _slice_lm_batch(lm_batch, span_start, span_stop)
+        (micro_loss, micro_ntok), micro_grads = loss_and_grad(model, micro)
+        weight = (micro_ntok.astype(mx.float32) / total_ntokens)
+        accumulator = _accumulate_grads_token_weighted(accumulator, micro_grads, weight)
+        weighted_loss = weighted_loss + micro_loss.astype(mx.float32) * weight
+        # Free-barrier: force the accumulator + running loss, then drop the
+        # microbatch's forward/backward graph before building the next one.
+        mx.eval(accumulator, weighted_loss)
+        del micro_grads, micro, micro_loss, micro_ntok
+
+    if accumulator is None:
+        raise RuntimeError("grad accumulation produced no gradients")
     if _strict_dtype_contract_enabled():
-        assert_grad_dtype_matches_param_dtype(grads, model.parameters())
-    optimizer.update(model, grads)
-    mx.eval(model.parameters(), optimizer.state, loss, ntokens)
+        assert_grad_dtype_matches_param_dtype(accumulator, model.parameters())
+    optimizer.update(model, accumulator)
+    mx.eval(model.parameters(), optimizer.state, weighted_loss)
     elapsed = time.perf_counter() - start
 
-    tokens = int(ntokens.item())
+    if clear_cache:
+        mx.clear_cache()
     return TrainStepResult(
-        loss=float(loss.item()),
-        ntokens=tokens,
+        loss=float(weighted_loss.item()),
+        ntokens=total_tokens_int,
         seconds=elapsed,
-        tokens_per_second=tokens / elapsed if elapsed > 0 else float("inf"),
+        tokens_per_second=total_tokens_int / elapsed if elapsed > 0 else float("inf"),
     )
 
 
diff --git a/scratch/pr2_grad_accum_equivalence.py b/scratch/pr2_grad_accum_equivalence.py
@@ -0,0 +1,125 @@
+"""PR-2 numeric-equivalence check (FUSED-PIPELINE-ROADMAP §4/§5).
+
+Verify that token-weighted grad accumulation over N microbatches produces
+gradients and a loss numerically equal (within fp tolerance) to the
+full-batch step. RULE #1: a mismatch here means the loop fix is wrong.
+
+Run: python3 scratch/pr2_grad_accum_equivalence.py
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten
+
+from cppmega_mlx.data.batch import synthetic_token_batch
+from cppmega_mlx.models.hybrid_lm import HybridTinyConfig, HybridTinyLM
+from cppmega_mlx.training.loss import next_token_cross_entropy
+from cppmega_mlx.training.loop import one_step_train
+
+
+def _build(grad_checkpoint: bool) -> HybridTinyLM:
+    mx.random.seed(0)
+    cfg = HybridTinyConfig(
+        vocab_size=64,
+        hidden_size=16,
+        pattern="A",
+        depth=2,
+        num_attention_heads=4,
+        max_seq_length=16,
+        grad_checkpoint=grad_checkpoint,
+    )
+    model = HybridTinyLM(config=cfg)
+    mx.eval(model.parameters())
+    return model
+
+
+def _grads_only(model: HybridTinyLM, batch) -> dict:
+    loss_and_grad = nn.value_and_grad(model, next_token_cross_entropy)
+    (loss, ntokens), grads = loss_and_grad(model, batch)
+    mx.eval(grads, loss, ntokens)
+    return loss, ntokens, grads
+
+
+def _flat(grads) -> dict:
+    return {k: v for k, v in tree_flatten(grads)}
+
+
+def _max_rel_diff(a: dict, b: dict) -> float:
+    worst = 0.0
+    worst_name = ""
+    for k in a:
+        ga, gb = a[k], b[k]
+        denom = mx.maximum(mx.abs(ga).max(), mx.array(1e-6))
+        rel = float((mx.abs(ga - gb).max() / denom).item())
+        if rel > worst:
+            worst, worst_name = rel, k
+    print(f"  worst grad rel-diff: {worst:.3e} at {worst_name}")
+    return worst
+
+
+def main() -> None:
+    batch_size = 4
+    batch = synthetic_token_batch(batch_size=batch_size, seq_length=16, vocab_size=64)
+
+    # Reference full-batch grads (single forward/backward).
+    ref_model = _build(grad_checkpoint=False)
+    full_loss, full_ntok, full_grads = _grads_only(ref_model, batch)
+    full_flat = _flat(full_grads)
+    print(f"full-batch loss={float(full_loss.item()):.8f} ntokens={int(full_ntok.item())}")
+
+    # Manual token-weighted accumulation matching one_step_train's math.
+    acc_model = _build(grad_checkpoint=False)
+    total_ntok = batch.target_mask.sum().astype(mx.float32)
+    mx.eval(total_ntok)
+    num_micro = 4
+    base = batch_size // num_micro
+    acc = None
+    acc_loss = mx.array(0.0, dtype=mx.float32)
+    from cppmega_mlx.training.loop import _slice_lm_batch
+    loss_and_grad = nn.value_and_grad(acc_model, next_token_cross_entropy)
+    for i in range(num_micro):
+        micro = _slice_lm_batch(batch, i * base, (i + 1) * base)
+        (ml, mn), mg = loss_and_grad(acc_model, micro)
+        w = mn.astype(mx.float32) / total_ntok
+        wg = {k: v * w for k, v in tree_flatten(mg)}
+        if acc is None:
+            acc = wg
+        else:
+            acc = {k: acc[k] + wg[k] for k in acc}
+        acc_loss = acc_loss + ml.astype(mx.float32) * w
+        mx.eval(acc, acc_loss)
+    print(f"accum   loss={float(acc_loss.item()):.8f}")
+    print(f"loss abs-diff: {abs(float(acc_loss.item()) - float(full_loss.item())):.3e}")
+    worst = _max_rel_diff(full_flat, acc)
+
+    tol = 2e-4
+    assert abs(float(acc_loss.item()) - float(full_loss.item())) < tol, "loss mismatch"
+    assert worst < tol, f"grad mismatch {worst} >= {tol}"
+
+    # End-to-end through one_step_train (optimizer applied; compare param deltas).
+    from cppmega_mlx.training.optimizers import make_adamw
+    m_full = _build(grad_checkpoint=False)
+    opt_full = make_adamw(learning_rate=1e-3)
+    r_full = one_step_train(m_full, opt_full, batch, grad_accum_steps=1)
+
+    m_acc = _build(grad_checkpoint=False)
+    opt_acc = make_adamw(learning_rate=1e-3)
+    r_acc = one_step_train(m_acc, opt_acc, batch, grad_accum_steps=4)
+    print(f"one_step_train loss full={r_full.loss:.8f} accum={r_acc.loss:.8f} "
+          f"diff={abs(r_full.loss - r_acc.loss):.3e}")
+    pf = _flat(m_full.parameters())
+    pa = _flat(m_acc.parameters())
+    pworst = _max_rel_diff(pf, pa)
+    assert abs(r_full.loss - r_acc.loss) < tol, "one_step_train loss mismatch"
+    assert pworst < tol, f"post-step param mismatch {pworst}"
+    assert r_full.ntokens == r_acc.ntokens, "ntokens mismatch"
+    print("PASS: grad accumulation is numerically equivalent to full-batch step")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/pr2_seq4096_memory_probe.py b/scripts/pr2_seq4096_memory_probe.py
