# Dist Lifecycle > 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". - Skill: `nvidia-bionemo/dist-lifecycle` (Agent Skill) - Install (CLI): `npx skillmds@latest add nvidia-bionemo/dist-lifecycle` - Raw SKILL.md: https://api.skillmd.com/api/skills/nvidia-bionemo/dist-lifecycle/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: AI & ML - Author: nvidia-bionemo (https://skillmd.com/u/nvidia-bionemo) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/nvidia-bionemo/dist-lifecycle --- # 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.