deepmodeling / deepmodeling/deepmd-kit

feat(pt_expt): add compile-aware EMA lifecycle and checkpoint retention

Open
#5,818 0 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.