TRN-5 — The L-BFGS full-batch training regime

[aacostadiaz]

Depends on: TRN-1 (#1575), TRN-2 (#1576), TRN-3 (#1577), TRN-4 (#1578), CFG-1 (#1574) · Blocks: GATE-3 (#1585)

Prerequisite: the Phase 0 characterization suite is committed — goldens, fixtures and the single-source tolerance table under tests/golden/, characterization tests under tests/unit/. References of the form P0-N name that work.

Context: --lbfgs (mace/tools/arg_parser.py:984-989) is a second training regime, not an optimizer choice, and no other ticket covers it — v1 without it is a regression. Legacy switches the optimizer object after the scheduler, SWA container, checkpoint handler, resume and EMA have already been built (mace/cli/run_train.py:971-976), constructing LBFGS(model.parameters(), history_size=200, max_iter=20, line_search_fn="strong_wolfe"), and the training loop then dispatches on isinstance(optimizer, LBFGS) (mace/tools/train.py:386-399) into take_step_lbfgs (:467). That step is a genuinely different algorithm: it first counts the whole dataset (total_sample_count, all-reduced under DDP), then runs one optimizer.step(closure) per epoch where the closure iterates every batch, scales each batch loss by num_graphs / total_sample_count, backpropagates per batch and accumulates — a full-batch gradient assembled in chunks. Three consequences the regime cannot work without: the dataloaders must keep their ragged tail (drop_last=(not args.lbfgs) at run_train.py:741,:757,:777); under DDP the closure broadcasts every parameter from rank 0 and all-reduces the total loss on each line-search evaluation (train.py:494-501,:544-545), because strong_wolfe calls the closure a data-dependent number of times and the ranks would otherwise diverge; and resume has no L-BFGS fallback. The restart_lbfgs retry that reloaded the checkpoint after the optimizer swap was unreachable and has been deleted; its non-execution is pinned by tests/workflows/test_cli_contracts.py:375. An Adam state dict still cannot be loaded into L-BFGS, so v1 has to decide what resume means for this regime rather than port a mechanism that never ran. Coverage today is one workflow test (tests/workflows/test_run_train.py:856 test_run_train_lbfgs) that asserts the run completes. This ticket lands the regime as a stage kind in TRN-2 (#1576)'s per-stage schedule rather than a boolean that rewrites already-built objects.

Interface & constraints:

1. L-BFGS is a stage kind, declared in config. schedule.stages = [{kind: "minibatch", optimizer: adam, ...}, {kind: "full_batch", optimizer: lbfgs, ...}]. The legacy behaviour — Adam-then-LBFGS for the whole run — is exactly a two-stage schedule; nothing is swapped in place, and no object built for an earlier stage is silently reused by a later one. The legacy defaults (history_size=200, max_iter=20, line_search_fn="strong_wolfe") become typed fields with those values.
2. Incompatible companions are rejected at config time, not ignored at run time. CFG-1 (CFG-1 — Full training config schema with cross-section validation #1574)'s validator forbids EMA and a plateau-style scheduler on an L-BFGS stage. Legacy builds both and then replaces the optimizer, leaving lr_scheduler.step(metrics=valid_loss) driving an object the optimizer no longer is and ema shadowing weights that move once per epoch; v1 refuses the combination instead of producing a number nobody can interpret.
3. The full-batch step is explicit and reproduces the legacy arithmetic. One step(closure) per epoch; the closure zeroes grads, iterates the stage's loader in order, weights each batch loss by num_graphs / total_sample_count, backpropagates per batch, applies max_grad_norm clipping once over the accumulated gradient, and returns the summed loss. The per-batch weighting is what makes the result equal to a single loss over the concatenated dataset — assert that equivalence in a test at fp64 on a tiny fixture, since it is the property the regime exists for.
4. drop_last comes from the stage, not from a flag. A full-batch stage requests drop_last=False through TRN-3 (TRN-3 — Multi-dataloader balancing and metrics/logging (error tables, optional wandb) #1577)'s per-stage loader property; the mini-batch stage keeps drop_last=True. No args.lbfgs reaches the loader construction.
5. DDP protocol. Under a distributed run the closure broadcasts parameters from rank 0 and all-reduces the total loss before returning, exactly as legacy does, so every rank evaluates the same point during the line search. The parameter broadcast is the contract; the signal tensor legacy uses to keep non-zero ranks in step (train.py:497-502) is an implementation detail v1 may replace, provided the invariant "all ranks see the same parameters and the same loss at every closure evaluation" is asserted by the smoke.
6. Resume across the regime boundary is explicit. A checkpoint records which stage wrote it (TRN-4 (TRN-4 — Safetensors checkpointing/resume and the DDP port #1578)'s stage counter), so resuming into an L-BFGS stage from a mini-batch checkpoint restores weights and declares the optimizer state as not transferable (ResumeResult.optimizer_state = "reinitialized"), instead of legacy's three nested except Exception blocks inferring it from a load failure.

Task:

1. Add the full_batch stage kind and its typed L-BFGS fields to TRN-2 (TRN-2 — Composable losses from observable specs, per-stage schedules, transform registry, and the SCF model-transform hook #1576)'s schedule; wire CFG-1 (CFG-1 — Full training config schema with cross-section validation #1574)'s EMA/plateau-scheduler validator.
2. Implement mace_torch/train/full_batch.py: the closure-based step with per-batch loss weighting, gradient accumulation, single clipping, and the loss reduction.
3. Wire the per-stage drop_last request through TRN-3 (TRN-3 — Multi-dataloader balancing and metrics/logging (error tables, optional wandb) #1577)'s loader layer.
4. Implement the DDP closure protocol (parameter broadcast + loss all-reduce) on TRN-4 (TRN-4 — Safetensors checkpointing/resume and the DDP port #1578)'s primitives.
5. Regime-aware resume: stage-tagged checkpoints and the typed non-transferable-optimizer outcome; remove the need for any load-failure inference.
6. Tests: full-batch/concatenated-loss equivalence at fp64; a two-stage Adam→L-BFGS run that decreases loss and writes a resumable checkpoint at the boundary; a 2-rank gloo L-BFGS smoke asserting identical parameters and loss on every closure evaluation; port tests/workflows/test_run_train.py::test_run_train_lbfgs under --engine v1.

Out of scope: the schedule machinery itself (TRN-2 (#1576)); loaders and balancing (TRN-3 (#1577)); checkpoint format and the DDP primitives (TRN-4 (#1578)); any new optimizer beyond the three legacy choices (adam, adamw, schedulefree, arg_parser.py:887-892).

Acceptance criteria:

• A full_batch L-BFGS stage runs one step(closure) per epoch and its accumulated gradient equals the gradient of the loss over the concatenated dataset at fp64 on a tiny fixture.
• An L-BFGS stage combined with EMA or a plateau scheduler fails config validation, naming both fields.
• The full-batch stage's loaders keep every configuration (drop_last=False); the mini-batch stage's do not — asserted by sample counts, with no lbfgs boolean reaching loader construction.
• 2-rank gloo L-BFGS smoke: parameters and total loss are identical on every closure evaluation across ranks, and the run's final weights match a single-process run of the same data at fp64.
• Resuming an L-BFGS stage from a mini-batch checkpoint restores weights and reports the optimizer state as reinitialized; no except Exception fallback exists in the v1 resume path.
• test_run_train_lbfgs passes under --engine v1.

Inventory gaps assigned here:

• --lbfgs — one completion smoke today (tests/workflows/test_run_train.py:856); nothing pins the full-batch arithmetic, the drop_last inversion, or the DDP closure protocol.

Verify:

python -m pytest packages/mace-torch/tests/train/test_full_batch.py -v      # closure arithmetic, gradient equivalence, drop_last
python -m pytest packages/mace-torch/tests/train/test_full_batch_ddp.py -v  # 2-rank gloo closure protocol
MACE_ENGINE=v1 python -m pytest tests/workflows/test_run_train.py -k lbfgs -v

Review focus: the per-batch loss weighting (get it wrong and the regime trains on a differently-scaled objective without failing anything) and the DDP closure protocol, since strong_wolfe evaluates the closure a data-dependent number of times and a missed broadcast diverges the ranks silently.

---