deepmodeling / deepmodeling/deepmd-kit
feat(pt_expt): add compile-aware EMA lifecycle and checkpoint retention
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Current status
The pt_expt EMA lifecycle is implemented on `master`:
- `enable_ema`, `ema_decay`, and `ema_ckpt_keep` are consumed;
- canonical parameter names are independent of compile/DDP wrappers;
- shadows are initialized and updated after successful optimizer steps;
- restart state and materialized EMA checkpoints use canonical keys;
- live-weight and EMA checkpoint retention are independent;
- EMA weights are available to validation and deployment checkpoint flows;
- unsupported sharded configurations are rejected explicitly.
The remaining functional requirement is the intended high-performance foreach update path.
## Remaining problem
`ModelEMA.update()` currently iterates over canonical named parameters in Python and calls `shadow.lerp_()` once per tensor. This is numerically correct, but it launches one interpolation operation per parameter and repeats name/structure traversal in the per-step hot path.
For large multi-task models, the EMA overhead scales with the number of parameter tensors rather than with a small number of device/dtype groups.
## Proposed design
Build the update plan once when EMA binds to the canonical model:
1. retain deterministic canonical name-to-shadow mapping for serialization;
2. group stable parameter and shadow references by compatible device, dtype, and layout;
3. update each dense group under `torch.no_grad()` using `torch._foreach_lerp_(shadow_group, parameter_group, 1 - decay)`;
4. use a small explicit per-tensor fallback only for layouts unsupported by foreach, or reject unsupported sharded tensors before training;
5. rebuild and validate the groups only when loading/rebinding EMA state, never on every optimizer step.
The bound parameter references must point to the canonical uncompiled model so compile/DDP wrapper namespaces do not enter either the update plan or checkpoint keys. Shadows stay on device; CPU copies remain checkpoint-boundary work only.
## Performance requirements
- `enable_ema=false` performs no EMA allocation or per-step work.
- The update path performs no parameter-name inspection, CPU copy, or host synchronization.
- Dense parameters use one foreach launch per compatible group rather than one launch per tensor.
- The change must not add Dynamo guards or compile the EMA control flow.
- DDP/ZeRO-1 requires no additional communication after the optimizer step.
## Acceptance criteria
- EMA preserves `shadow = decay * shadow + (1 - decay) * parameter` exactly within dtype tolerance.
- Compiled and uncompiled models retain identical canonical EMA/checkpoint semantics.
- Restart, validation, deployment checkpoints, multi-task naming, and independent retention behavior remain unchanged.
- Dense parameters are updated through grouped foreach operations.
- Unsupported sharded or tensor-layout cases fail explicitly or use a documented non-hot fallback.
Refs #5755.
Contributor guide
Assessment
This issue has not been assessed yet.