| name | dist_lifecycle |
| description | Stand up the distributed-model lifecycle for a CP training/inference run: device placement before DTensor wrapping, the all-trainable-params-are-DTensors invariant (with placeholder/freeze for unimplemented modules), checkpoint save (DTensor → plain) and load (realign via the live state_dict template + redistribute optimizer state to parameter placements), resume RNG seed offset, and DTensor-safe EMA. Use after dtensor_modules + build_infra, when wiring the ported CP modules into a real trainer/predictor (Lightning or custom) and you need checkpoints, resume, or EMA to work — the gap between "modules pass parity" and "training runs and resumes".
|
| argument-hint | [wrap | checkpoint | resume | ema | optimizer] |
dist_lifecycle — make a CP model train, checkpoint, and resume correctly
Porting modules (dtensor_modules) proves the math; this skill makes the trainer
work: device placement, parameter typing, checkpoint I/O, resume, and EMA under
DTensor. These steps are where a parity-passing model still fails to train or resume.
Reference files (all under $BOLTZ_CP_REPO/src/boltz/distributed):
lightning_strategy.py, train.py, predict.py, model/optim/ema.py, manager.py.
Preconditions
- Modules ported + parity-tested (
dtensor_modules + test).
docs/cp_infra.md (topology, world_size, launch recipe) and
docs/current_code_structure.md (framework: Lightning / DeepSpeed / custom).
Step 1 — Device placement before DTensor wrapping
Move the serial module to its CUDA device before wrapping/distributing it. A CPU
module wrapped against a CUDA mesh is a device mismatch. Order: build serial → .to(device)
→ DTensor-wrap (train.py:_create_distributed_model). Fail fast if torch.distributed
is not initialized — no silent single-GPU fallback (Rule 12).
Step 2 — Every trainable parameter must be a DTensor
Gradient redistribution (on_after_backward → param.grad.redistribute(...)) only fires
on DTensor parameters; a plain-tensor trainable parameter silently bypasses it and
diverges across ranks. So:
- Validate in
__init__: every requires_grad parameter is a DTensor (else raise).
- For modules you have not ported yet, use a placeholder: keep the parameters
(so
state_dict stays key-compatible with the serial checkpoint), requires_grad_(False),
and raise NotImplementedError on forward. This keeps checkpoints portable and lets
training proceed without the unimplemented submodule (boltz2.py _PlaceholderModule).
- Mirror serial attribute names and registration order so
state_dict keys line up.
Step 3 — Checkpoint SAVE: DTensor → plain tensors
Convert every DTensor in the payload to a plain (full) tensor before writing
(lightning_strategy.py). Save an identical plain checkpoint (portable across
topologies / loadable by the serial model). Do not use a callable map_location
with DTensor checkpoints.
Step 4 — Checkpoint LOAD + optimizer-state redistribution
- Load the plain checkpoint, then realign placements using the live model's
state_dict() as the template — from_local/distribute_tensor each loaded tensor
to the placement the live parameter expects.
- Redistribute optimizer buffers to match parameter placements before
optimizer.step() — plain-tensor optimizer state from the checkpoint mixed with
DTensor parameters errors (or silently mis-updates) on step.
- Pretrained / transfer load:
strict=False plus an explicit architecture check
(e.g. assert v2 flags / module presence) — strict=False silently tolerates a
confidence-module or V1/V2 mismatch that then crashes or mis-loads.
- Use FQN (fully-qualified-name) optimizer-state keys, and auto-detect legacy
integer-indexed keys on load so an old checkpoint still maps — integer param-group indices
break the moment module registration order changes.
- Prefer
from_local (with explicit shape/stride/placements) over distribute_tensor
on every same-process load/restore path, and keep map_location a string device — a
callable map_location strips the DTensor type.
- After
load_from_checkpoint, re-inject any ctor args passed to
save_hyperparameters(ignore=[...]) (validators, callbacks, non-serializable objects):
they come back None, so restore them from the live module/config before use.
Step 5 — Resume RNG
Offset the seed by global rank, and on resume additionally by epoch +
global_step, so ranks don't replay identical data/noise and a resume doesn't repeat
the pre-resume stream (train.py). This is the entropy side of the single-device RNG
rule; replicate-axis consistency is still by broadcast, and control-flow scalars
(recycling/sampling counts) are broadcast across the flat CP group. (See RULES.md
"single-device RNG entropy equivalence" + Rule 18.)
Step 6 — EMA under DTensor
EMA updates parameters with in-place .data arithmetic, which bypasses DTensor
dispatch — so EMA only works correctly when the parameters are Replicate on every
mesh axis; Shard/Partial placements give silently wrong EMA updates. When swapping
EMA weights in/out, back up each tensor's device_mesh + placements and rewrap, rather
than round-tripping through CPU (which strips the DTensor type) — see model/optim/ema.py.
Guard the EMA update with torch.inference_mode(False) (an EMA step inside an inference-mode
context errors on the in-place write), keep a CPU backup of the shadow weights for
checkpointing, and on fine-tune resume backfill any missing EMA keys from the live
parameters so a newly-added module doesn't crash the EMA load.
Step 7 — Verify with a stop-and-go test
Hand to /fold-cp:test: train N steps → checkpoint → resume → train M more, and assert
(a) the resumed run's weights/optimizer state match a non-interrupted run (within the
fp budget), (b) intermediate checkpoints differ from the final, (c) the saved checkpoint
contains plain tensors (no DTensors), (d) it loads into the serial model, and (e) when
you change the checkpoint/serialization layout, a checkpoint written by the prior format
still loads (backward-compat) — cover both save and load paths. Mirror
tests/distributed/test_dtensor_stop_and_go.py (monkeypatch the distributed cleanup to a
no-op for in-process port reuse).
Gotchas that pass parity but break training
- A detached / re-wrapped DTensor loss silently drops its gradient. Building a loss via
from_local(local.detach(), …), adding a structurally-zero term, or letting the framework's
grad-accumulation divide such a tensor creates a fresh DTensor from _local_tensor with
requires_grad=False — backward then no-ops and weights never move, with no error. Keep
the loss on the autograd graph (no detach/native re-dispatch on the differentiable path,
Rule 6); if you must accumulate manually, pin grad-accumulation to 1 and verify grad_norm > 0.
- Do not "harden" framework checkpoint loads to
weights_only=True. Lightning/torch
checkpoints carry non-tensor objects by design; forcing weights_only=True breaks the load —
it is accepted risk, not a vulnerability to patch.
Output contract
- The CP model places on device before wrapping; all trainable params are DTensors
(unimplemented ones placeholdered/frozen).
- Checkpoints save as plain tensors, load with placement realignment + optimizer-state
redistribution, and pass a stop-and-go resume parity test.
- Resume seeding offsets by rank (+ epoch/global_step); EMA (if used) is Replicate-only.
- A short report: framework, checkpoint format, what is placeholdered, and the
stop-and-go test result.
