feat(pt_expt): align the training runtime with pt - #5958
Conversation
…ith pt
A pt_expt run could not be operated under the same conventions as a pt
run: it ignored the checkpoint directory and retention options, kept no
EMA, never reported its parameter count, and printed progress in its own
format with a systematically inflated estimate of the remaining time.
Four features are brought over, and everything about them that is not
specific to a backend is described once so that both backends execute the
same implementation instead of two drifting copies.
Checkpointing. `save_dir` and `ckpt_keep_ratio` are honoured, and the
on-disk layout -- naming, prefix symlinks, pointer file and retention --
moves into `CheckpointStore` (deepmd/dpmodel/train/checkpoint.py), which
both backends now use; pt loses four hand-rolled copies of the publish
sequence. The store drops checkpoints numbered above the one being
written before it applies the retention window. Those are remnants of a
longer earlier run over the same directory, and leaving them in place let
the window discard the checkpoint that was just written, so restarting a
run in a finished directory kept no result at all. A disabled window
(`max_ckpt_keep < 1`) now retains every checkpoint, as the jax and tf2
backends already do, rather than deleting all of them including the
current one. `resolve_keep_ckpt_count` also handles `save_freq <= 0`,
which previously raised `ZeroDivisionError` when combined with a
retention ratio.
EMA. `enable_ema`, `ema_decay` and `ema_ckpt_keep` are honoured. The
shadow weights are updated after every optimizer step, written as a
separate family of checkpoints carrying neither optimizer nor EMA state,
and restored on restart. `deepmd/pt/train/ema.py` is reused as is rather
than copied.
Full validation. `build_full_validators` configures the live-weight and
the EMA-weight flow together, since they differ only in the weights they
read, the log they write and the prefix of the checkpoints they select.
pt_expt thereby gains `ema_full_validation`, and the per-flow eligibility
check stays with the backend that knows what it supports.
`compiled_infer` and `amp_infer` reach the models through the shared
`infer_env_defaults` translation. `tf32_infer` remains unimplemented in
pt_expt, which has no TF32 path yet.
Training report. The display now prints the losses before the wall-clock
line and omits the per-step average, matching pt, and the run ends with
the average step time over the representative intervals. The remaining
time is extrapolated from the interval that just ended rather than from
the average since the run began: the latter carries the one-off cost of
the first steps, such as graph compilation, and therefore never stops
overestimating. This accounting moves into `TrainingTimer`
(deepmd/dpmodel/train/timing.py), replacing pt's three loose counters.
The parameter-count report moves to `deepmd/loggers/training.py`, the
home of the other training log messages, and reads counts a backend
supplies.
Relocation. `deepmd/pt/train/{utils,validation,ema}.py` move under
`deepmd/pt_expt/train/`. pt is being retired, so the shared training code
belongs with the backend that outlives it and the dependency arrow is
reversed. While moving, the validator recognizes its validation data by
the surface it exposes rather than by pt's dataset types, and reads the
environment constants from pt_expt; only `AutoBatchSize` and
`to_torch_tensor` still come from `deepmd/pt/utils`, which is the next
unit to migrate.
Two gaps separated pt_expt from pt in distributed training: a run was always plain data parallel, and a step was taken without ever inspecting its gradient. Both are closed here, and the part that is not specific to PyTorch is shared with pt rather than duplicated. `training.zero_stage` now selects the same four strategies as in pt: plain DDP, DDP over a redundancy-sharded optimizer, and FSDP2 sharding the gradients or the parameters as well. What a stage implies -- which wrapper holds the model, how the optimizer is built, how a checkpoint is assembled, whether a gradient norm may be reduced locally -- follows from the stage alone, so `ShardingPolicy` in the backend-independent train layer states it once and both backends query it, rather than each comparing the stage against numbers wherever a decision is due; pt sheds twenty such comparisons. A single-process run drops the requested stage instead of failing, so one configuration stays usable whether or not it is launched across ranks. Assembling a checkpoint out of shards is a collective operation, which the shared training loop had no notion of: it called `save_checkpoint` on the chief alone, which would leave the other ranks waiting at the next barrier. The trainer gained `checkpoint_is_collective`, false by default so that tf2 and jax are unaffected; a backend that opts in is called on every rank and gates the write itself. A run is restored before the model is distributed. A checkpoint records whole tensors, and those cannot be copied into parameters that FSDP2 has already cut into shards. The optimizer is still built after distribution, so its state is restored separately, through the distributed-checkpoint API when the stage calls for it. That reorder also removes a second construction of the learning-rate schedule, and with it a defect it was covering: the schedule was built before the resumed step was known and rebuilt with the true value only when optimizer state was present, so restarting from a checkpoint that carried a step but no optimizer state -- a frozen model, or one saved without it -- resumed at the wrong learning rate. Sharding is rejected alongside multi-task training, EMA from stage two, and `change_bias_after_training`, as in pt. pt_expt additionally rejects `enable_compile`, whose graph is traced from the parameters that FSDP2 replaces with DTensors. Display-time validation is skipped once the parameters are sharded, because its forward gathers them while the display runs on the chief alone; the full-validation flow, which every rank enters together, stays available. pt_expt clipped gradients with the stock `torch.nn.utils.clip_grad_norm_` and never inspected the resulting norm, so a run could write a checkpoint of a model that had already diverged, and a gradient that was large but still representable could be misread as infinite when the sum of squares of the naive reduction overflowed. pt has carried safeguards against both for a while; they now serve pt_expt as well. The two safeguards move out of the training utilities, which had accreted four unrelated concerns, into `deepmd/pt_expt/train/gradient.py`. They share one rationale -- keep the reduction overflow-safe, and keep the verdict off the host until it is needed -- which the module can now state once. `deepmd/pt_expt/train/utils.py` keeps the trainer setup helpers. pt_expt feeds the norm to the guard on every step and consults the guard at the checkpoint boundary. The verdict is deliberately not read anywhere else: the check resets the accumulated state, so a second caller between two boundaries would consume a divergence that the checkpoint about to be written should have seen. Reading it once per boundary also keeps the step free of host synchronization, which is why the state is accumulated on device in the first place. The reduction itself is overflow-safe except where the parameters are sharded, since that reduction has to propagate DTensor sharding instead. Verified on two gloo ranks: every stage trains, checkpoints and restarts, recording whole tensors and restoring the optimizer state. One test pins a defect the sharded path invites, in that a redundancy-sharded optimizer turns each constructor keyword into a param-group default and would therefore record a second copy of the model in every checkpoint of a name-routed optimizer.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe PR adds shared checkpoint, sharding, and timing APIs. It integrates ZeRO and FSDP2 training, EMA checkpoints and validation, guarded gradient handling, parameter logging, generalized validation, and updated configuration support. ChangesTraining infrastructure
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/utils/argcheck.py (1)
6028-6033: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winKeep
compiled_inferbehindmodel.use_compilein thept_exptdocs.The
pt_exptbackend exportsDP_COMPILE_INFER, but DPA4 models still sample it only as a descriptor activation-checkpoint switch, andpt_exptraises onmodel.use_compile. TheArgument("compiled_infer", ..., doc=supported_backends("pt", "pt_expt") + doc_compiled_infer)makescompiled_inferlook available for the eval torch.compile path inpt_expt, which users cannot enable. Label itptonly, or add the actualpt_expttorch_compile path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/utils/argcheck.py` around lines 6028 - 6033, Update the compiled_infer Argument declaration so its documentation advertises the option only for the pt backend, unless a real pt_expt torch.compile implementation is added; do not expose it as an eval torch.compile option for pt_expt.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/dpmodel/train/timing.py`:
- Around line 71-117: Update the timing implementation around __init__ and
record to use time.monotonic() for _interval_start and elapsed wall_time
calculations, preventing clock adjustments from producing negative durations or
forecasts. Keep the displayed timestamp based on a separate time.time() reading
converted to the local timezone.
In `@deepmd/pt_expt/train/ema.py`:
- Around line 76-90: Update apply_shadow to snapshot every parameter’s original
data before performing any EMA copy, including aliased parameters returned by
_named_model_parameters. Separate backup collection from shadow-value
application, then restore all originals in finally so shared parameters retain
their pre-EMA training values after the context exits.
In `@deepmd/pt/train/training.py`:
- Around line 1227-1229: The checkpoint_dir parameter passed to
resolve_best_checkpoint_dir uses Path(self.save_ckpt).parent as a default, but
this may differ from the active checkpoint store directory
(self.ckpt_store.directory) when training.save_dir is set and
validating.save_best_dir is unset. Update the resolve_best_checkpoint_dir call
to pass self.save_dir alongside validating_params and self.save_ckpt, so it can
construct the default validation checkpoint directory consistently from
self.ckpt_store.directory rather than inferring the parent directory from the
checkpoint file path alone.
- Line 463: Move the pretrained_model construction into the
scoped_env_defaults(eval_env_defaults) context manager block. Currently,
pretrained_model is built outside this context, which causes
get_model_for_wrapper to sample environment flags like DP_COMPILE_INFER,
DP_TF32_INFER, and DP_AMP_INFER using ambient settings rather than the intended
eval defaults. Ensure the pretrained_model is fully constructed and available
within the scoped_env_defaults context so that all downstream model building
operations use consistent environment configuration.
In `@source/tests/pt_expt/test_training.py`:
- Line 2253: Update the assertion for model_ema.ckpt.pt in the affected training
test to first verify the file exists, then check os.path.islink only when
platform.system() is not Windows. Add the platform import if needed, matching
the equivalent logic in the PT training test.
---
Outside diff comments:
In `@deepmd/utils/argcheck.py`:
- Around line 6028-6033: Update the compiled_infer Argument declaration so its
documentation advertises the option only for the pt backend, unless a real
pt_expt torch.compile implementation is added; do not expose it as an eval
torch.compile option for pt_expt.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 50fdceaf-7e0f-4d86-a25b-a86afa80e3f5
📒 Files selected for processing (26)
deepmd/dpmodel/train/__init__.pydeepmd/dpmodel/train/checkpoint.pydeepmd/dpmodel/train/sharding.pydeepmd/dpmodel/train/timing.pydeepmd/dpmodel/train/trainer.pydeepmd/loggers/training.pydeepmd/pt/train/training.pydeepmd/pt_expt/train/ema.pydeepmd/pt_expt/train/gradient.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/train/utils.pydeepmd/pt_expt/train/validation.pydeepmd/utils/argcheck.pydoc/train/training-advanced.mdsource/tests/common/dpmodel/test_train_abstract_trainer.pysource/tests/common/dpmodel/test_train_checkpoint.pysource/tests/common/dpmodel/test_train_sharding.pysource/tests/common/dpmodel/test_train_timing.pysource/tests/common/test_argcheck_training.pysource/tests/common/test_loggers_training.pysource/tests/pt/test_training.pysource/tests/pt/test_validation.pysource/tests/pt_expt/test_entrypoint.pysource/tests/pt_expt/test_train_gradient.pysource/tests/pt_expt/test_training.pysource/tests/pt_expt/test_training_ddp.py
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/utils/argcheck.py (1)
6028-6033: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winKeep
compiled_inferbehindmodel.use_compilein thept_exptdocs.The
pt_exptbackend exportsDP_COMPILE_INFER, but DPA4 models still sample it only as a descriptor activation-checkpoint switch, andpt_exptraises onmodel.use_compile. TheArgument("compiled_infer", ..., doc=supported_backends("pt", "pt_expt") + doc_compiled_infer)makescompiled_inferlook available for the eval torch.compile path inpt_expt, which users cannot enable. Label itptonly, or add the actualpt_expttorch_compile path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/utils/argcheck.py` around lines 6028 - 6033, Update the compiled_infer Argument declaration so its documentation advertises the option only for the pt backend, unless a real pt_expt torch.compile implementation is added; do not expose it as an eval torch.compile option for pt_expt.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/dpmodel/train/timing.py`:
- Around line 71-117: Update the timing implementation around __init__ and
record to use time.monotonic() for _interval_start and elapsed wall_time
calculations, preventing clock adjustments from producing negative durations or
forecasts. Keep the displayed timestamp based on a separate time.time() reading
converted to the local timezone.
In `@deepmd/pt_expt/train/ema.py`:
- Around line 76-90: Update apply_shadow to snapshot every parameter’s original
data before performing any EMA copy, including aliased parameters returned by
_named_model_parameters. Separate backup collection from shadow-value
application, then restore all originals in finally so shared parameters retain
their pre-EMA training values after the context exits.
In `@deepmd/pt/train/training.py`:
- Around line 1227-1229: The checkpoint_dir parameter passed to
resolve_best_checkpoint_dir uses Path(self.save_ckpt).parent as a default, but
this may differ from the active checkpoint store directory
(self.ckpt_store.directory) when training.save_dir is set and
validating.save_best_dir is unset. Update the resolve_best_checkpoint_dir call
to pass self.save_dir alongside validating_params and self.save_ckpt, so it can
construct the default validation checkpoint directory consistently from
self.ckpt_store.directory rather than inferring the parent directory from the
checkpoint file path alone.
- Line 463: Move the pretrained_model construction into the
scoped_env_defaults(eval_env_defaults) context manager block. Currently,
pretrained_model is built outside this context, which causes
get_model_for_wrapper to sample environment flags like DP_COMPILE_INFER,
DP_TF32_INFER, and DP_AMP_INFER using ambient settings rather than the intended
eval defaults. Ensure the pretrained_model is fully constructed and available
within the scoped_env_defaults context so that all downstream model building
operations use consistent environment configuration.
In `@source/tests/pt_expt/test_training.py`:
- Line 2253: Update the assertion for model_ema.ckpt.pt in the affected training
test to first verify the file exists, then check os.path.islink only when
platform.system() is not Windows. Add the platform import if needed, matching
the equivalent logic in the PT training test.
---
Outside diff comments:
In `@deepmd/utils/argcheck.py`:
- Around line 6028-6033: Update the compiled_infer Argument declaration so its
documentation advertises the option only for the pt backend, unless a real
pt_expt torch.compile implementation is added; do not expose it as an eval
torch.compile option for pt_expt.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 50fdceaf-7e0f-4d86-a25b-a86afa80e3f5
📒 Files selected for processing (26)
deepmd/dpmodel/train/__init__.pydeepmd/dpmodel/train/checkpoint.pydeepmd/dpmodel/train/sharding.pydeepmd/dpmodel/train/timing.pydeepmd/dpmodel/train/trainer.pydeepmd/loggers/training.pydeepmd/pt/train/training.pydeepmd/pt_expt/train/ema.pydeepmd/pt_expt/train/gradient.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/train/utils.pydeepmd/pt_expt/train/validation.pydeepmd/utils/argcheck.pydoc/train/training-advanced.mdsource/tests/common/dpmodel/test_train_abstract_trainer.pysource/tests/common/dpmodel/test_train_checkpoint.pysource/tests/common/dpmodel/test_train_sharding.pysource/tests/common/dpmodel/test_train_timing.pysource/tests/common/test_argcheck_training.pysource/tests/common/test_loggers_training.pysource/tests/pt/test_training.pysource/tests/pt/test_validation.pysource/tests/pt_expt/test_entrypoint.pysource/tests/pt_expt/test_train_gradient.pysource/tests/pt_expt/test_training.pysource/tests/pt_expt/test_training_ddp.py
🛑 Comments failed to post (1)
deepmd/pt_expt/train/ema.py (1)
76-90: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Take all parameter backups before applying EMA values.
_named_model_parameterscan return multiple names for aliased parameters in a model dictionary.apply_shadowbacks up and overwrites each name in one pass. The second alias is then backed up after the first alias already contains EMA values. Thefinallyblock can leave the shared parameter EMA-weighted after the context exits. The next training step then uses incorrect weights.Collect all backups before the first
copy_, then apply shadow values in a second pass.Proposed fix
backups: dict[str, torch.Tensor] = {} + named_parameters = self._named_model_parameters(model) try: with torch.no_grad(): - for name, param in self._named_model_parameters(model): + for name, param in named_parameters: backups[name] = param.detach().clone() + for name, param in named_parameters: param.copy_( self.shadow_params[name].to( device=param.device, dtype=param.dtype, @@ finally: with torch.no_grad(): - for name, param in self._named_model_parameters(model): + for name, param in named_parameters: if name in backups: param.copy_(backups[name])Also applies to: 178-200
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt_expt/train/ema.py` around lines 76 - 90, Update apply_shadow to snapshot every parameter’s original data before performing any EMA copy, including aliased parameters returned by _named_model_parameters. Separate backup collection from shadow-value application, then restore all originals in finally so shared parameters retain their pre-EMA training values after the context exits.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The shape of this is right and I want to say that before the problems. Extracting checkpoint layout, retention, timing and sharding into deepmd/dpmodel/train/ follows the precedent #5603 set, the three new modules import no torch so the backend-independence is real rather than nominal, and the alternative -- reimplementing the same policies inside pt_expt -- would have produced exactly the drift the PR body describes. Several things I went looking for turned out clean: the ZeRO collective ordering is correct (_collect_checkpoint_states runs consolidate_state_dict before the rank != 0 return, so no rank can skip a collective), ZeRO-1 restart across a different world size still works, pt's non-finite gradient handling is a pure rename with no behavioural diff, and resolve_keep_ckpt_count actually fixes a latent ZeroDivisionError on save_freq <= 0.
Two blocking problems, both inline, and one question.
The second one is the serious one: CheckpointStore.prune deletes checkpoints it is supposed to keep, on pt as well as pt_expt. I verified it by running the new class against the implementation it replaces rather than by reading, and the divergence is not subtle -- with max_ckpt_keep=10 and nine checkpoints on disk, master keeps all nine and this branch keeps two.
I also want to flag the ordering problem this creates for review itself. Because CI never got past collection, none of the ~19958 tests ran, including all the new ones. So the checkpoint bug was not caught by the suite, and more importantly nothing else in this 2600-line diff has been exercised either -- the parts I checked by hand look right, but "looks right on inspection" is a much weaker statement than this PR deserves given it touches production pt training. Worth fixing the import first and letting a green run tell us what else is there before anyone reads the rest too closely.
A few smaller notes I am recording rather than asking you to act on. format_training_message's step_time parameter and the avg = ... s/step field are gone, along with the TestFormatTrainingMessageStepTime test; those came from #5500 whose stated goal was to fold that average into the normal log line, and TrainingTimer.format_average() only prints once at end of run, so pt_expt, jax and tf2 all lose the per-interval figure. build_checkpoint_stores gates the retention log and store.prepare() on rank == 0 but every test uses the default rank, so the non-chief branch is unexercised -- the same gap I raised on #5603 for is_chief. ema_ckpt_keep moving from a hard 3 to inheriting max_ckpt_keep silently shrinks EMA retention for anyone who set max_ckpt_keep to 1 or 2; it is documented, so this is a release-note item rather than a defect. And pt_expt now inherits pt's "delete every checkpoint numbered above the current step" rule, which is a new data-loss path for pt_expt users restarting from an older checkpoint in a directory holding a longer run -- pre-existing for pt, new for pt_expt, and not mentioned in doc/train/training-advanced.md.
Bound checkpoint retention before slicing, measure elapsed time with a monotonic clock, and preserve aliased parameters across EMA swaps. Apply eval defaults to finetune source models and align cross-platform tests and backend documentation.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/pt/train/training.py (2)
1987-1996: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftMake ZeRO-1 full-validation checkpoint saves collective.
save_modelcollects optimizer state by default, and the ZeRO-1 branch callsself.optimizer.consolidate_state_dict(to=0). This requires every rank to enter_collect_checkpoint_states.Full-validation only executes
save_checkpointwhenself.rank == 0, so ZeRO-1 non-LoRA full-validation can hang whenvalidating.full_validation=trueandvalidating.save_best=true. Use a collective best-checkpoint save path, or reject this configuration.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt/train/training.py` around lines 1987 - 1996, Update the ZeRO-1 checkpoint flow around _collect_checkpoint_states and save_checkpoint so full-validation best-checkpoint saves enter the optimizer consolidation collectively on every rank, including non-LoRA runs. Ensure rank 0 still performs the actual checkpoint output after collective state gathering, or explicitly reject the incompatible validating.full_validation and validating.save_best configuration before entering the save path.
1744-1758: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winGuard or reject
save_freq=0before the modulo.
self.save_freqdefaults to1000, buttraining_params.get("save_freq")allows arbitrary integer config values. Withsave_freq=0,display_step_id % self.save_freqraisesZeroDivisionErrorbefore any final-checkpoint save, so either validatesave_freq > 0or skip the periodic modulo when it is disabled.Proposed fix
- should_save_checkpoint = ( - (display_step_id) % self.save_freq == 0 and _step_id != self.start_step - ) or (display_step_id) == self.num_steps + should_save_checkpoint = ( + ( + self.save_freq > 0 + and display_step_id % self.save_freq == 0 + and _step_id != self.start_step + ) + or display_step_id == self.num_steps + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt/train/training.py` around lines 1744 - 1758, Validate self.save_freq before the should_save_checkpoint calculation in the training flow, rejecting or explicitly handling zero so display_step_id % self.save_freq is never evaluated with a zero divisor. Preserve the final-checkpoint condition based on display_step_id == self.num_steps, and apply the validation at the configuration or initialization point that consumes training_params.get("save_freq").
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@deepmd/pt/train/training.py`:
- Around line 1987-1996: Update the ZeRO-1 checkpoint flow around
_collect_checkpoint_states and save_checkpoint so full-validation
best-checkpoint saves enter the optimizer consolidation collectively on every
rank, including non-LoRA runs. Ensure rank 0 still performs the actual
checkpoint output after collective state gathering, or explicitly reject the
incompatible validating.full_validation and validating.save_best configuration
before entering the save path.
- Around line 1744-1758: Validate self.save_freq before the
should_save_checkpoint calculation in the training flow, rejecting or explicitly
handling zero so display_step_id % self.save_freq is never evaluated with a zero
divisor. Preserve the final-checkpoint condition based on display_step_id ==
self.num_steps, and apply the validation at the configuration or initialization
point that consumes training_params.get("save_freq").
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 92ca0272-1a21-4c24-9860-24d2aeb56d12
📒 Files selected for processing (11)
deepmd/dpmodel/train/checkpoint.pydeepmd/dpmodel/train/timing.pydeepmd/pt/train/training.pydeepmd/pt_expt/train/ema.pydeepmd/utils/argcheck.pydoc/train/training-advanced.mdsource/tests/common/dpmodel/test_train_checkpoint.pysource/tests/common/dpmodel/test_train_timing.pysource/tests/pt_expt/test_ema.pysource/tests/pt_expt/test_entrypoint.pysource/tests/pt_expt/test_training.py
💤 Files with no reviewable changes (1)
- source/tests/pt_expt/test_entrypoint.py
🚧 Files skipped from review as they are similar to previous changes (6)
- doc/train/training-advanced.md
- source/tests/common/dpmodel/test_train_timing.py
- deepmd/dpmodel/train/timing.py
- source/tests/common/dpmodel/test_train_checkpoint.py
- deepmd/utils/argcheck.py
- source/tests/pt_expt/test_training.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed against the request: are the docs updated, and do other backends break?
Docs: partially updated
- Done:
doc/train/training-advanced.mdgets amax_ckpt_keepbullet,save_dir/ckpt_keep_rationow tagged PyTorch + PyTorch Exportable, and a restart-retention note. The argcheck doc strings (auto-rendered into the options table via.. dargs::indoc/train/train-input.rst) are also updated:ema_ckpt_keepinheritance, pt_expt tags on the save/EMA/ZeRO/full-validation/amp options, thezero_stageenable_compileexclusion for stages 2/3, and thetf32_infernote. - Gap:
doc/train/parallel-training.mdis not updated, although this PR makeszero_stageavailable to the PyTorch Exportable backend. That page still says ZeRO "Works only in PyTorch backend" and lists constraints without pt_expt or theenable_compile/enable_emarestrictions on stages 2/3. Users reaching for the new feature from the parallel-training guide will get stale guidance.
Other backends: not broken (verified by installing the PR head)
- All six trainer modules import cleanly:
pt,pt_expt,jax,tf2,pd,tf. - The moves of
deepmd.pt.train.{ema,utils,validation}topt_exptleave no dangling imports anywhere in the tree. - The shared
AbstractTrainerrefactor is consumed by jax/tf2/pt_expt; ransource/tests/jax/test_training.py+source/tests/tf2/test_training.py(33 passed) andsource/tests/pt/test_training.py(52 passed) — no regressions. format_training_messagedroppedstep_time; tf/pd callers don't use it.- The
ema_ckpt_keepargcheck change (int default 3 → None, inheritsmax_ckpt_keep) affects only pt/pt_expt; common argcheck/schema tests pass.
Regression found (pt_expt)
source/tests/pt_expt/test_training.py::test_unsupported_optimizer_has_clear_errorfails: unsupported optimizer types now raiseKeyError: 'adam_beta1'instead ofValueError("Unsupported optimizer type: ..."). See the inline comment.
Behavior change worth a release note
- pt's default EMA checkpoint retention changes from 3 to inheriting
max_ckpt_keep(default 5). It is documented in the options table and training-advanced.md, but the on-disk retention for existing users changes silently.
Attribution
Coding agent: opencode
opencode version: 1.18.13
Model: ustc/deepseek-v4-flash
Reasoning effort: max
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Follow-up review of deepmd/utils/argcheck.py in this PR.
Correct / consistent:
ema_ckpt_keeptype changeint(default 3) →[int, None](default None) is sound: theextra_checkadmits None, the generated JSON schema stays valid (test_doc_train_inputpasses), and the only reader isbuild_checkpoint_stores(deepmd/dpmodel/train/checkpoint.py:277) which handles None by inheritingmax_ckpt_keep. No other code reads it. Note the deliberate behavior change: the pt backend's default EMA retention moves from 3 tomax_ckpt_keep(default 5).- The
tf32_infernote is accurate: pt_expt never readsDP_TF32_INFER(grep of deepmd/pt_expt confirms it only sets it viainfer_env_defaults), so "no effect there" is correct.
Two doc inconsistencies (not functional breaks):
doc_max_ckpt_keepwas not updated, although the PR makesmax_ckpt_keepthe default retention window for the EMA family too. Its text still reads as applying to regular checkpoints only, while the newema_ckpt_keepdoc ("When unset, it inheritsmax_ckpt_keep") anddoc/train/training-advanced.mdboth describe the inheritance. Users reading the generated options table learn thatema_ckpt_keepinherits frommax_ckpt_keep, but not thatmax_ckpt_keepgoverns the EMA family by default. Suggest amendingdoc_max_ckpt_keepto mention the EMA-family inheritance, e.g. "The maximum number of recent periodic checkpoints retained for each checkpoint family; the EMA family inherits this window by default."- Inconsistent backend tags under
validating:amp_inferis retaggedsupported_backends("pt", "pt_expt"), butcompiled_inferstayspt-only. Both go through the same newinfer_env_defaultsin pt_expt (deepmd/pt_expt/train/utils.py:67 mapscompiled_infer→DP_COMPILE_INFER, line 69 mapsamp_infer→DP_AMP_INFER), and pt_expt code actually consumes both:DP_COMPILE_INFERin deepmd/pt_expt/descriptor/dpa4_nn/block.py:107,DP_AMP_INFERin deepmd/pt_expt/descriptor/dpa4.py:208. Ifamp_infermerits the pt_expt tag,compiled_infershould be tagged identically.
Attribution
Coding agent: opencode
opencode version: 1.18.13
Model: ustc/deepseek-v4-flash
Reasoning effort: max
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5958 +/- ##
==========================================
- Coverage 79.59% 79.35% -0.24%
==========================================
Files 1081 1085 +4
Lines 126244 126405 +161
Branches 4592 4598 +6
==========================================
- Hits 100490 100315 -175
- Misses 24101 24437 +336
Partials 1653 1653 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both of my earlier findings are fixed, and I checked each against c814ddd97 rather than going by the replies.
The import removal is right, including the part of your reply I had not verified when I raised it: open_stat_file really does own nested HDF5 parent creation (stat_file.py calls target.parent.mkdir(parents=True, exist_ok=True)), and test_stat_file.py already exercises it through tmp_path / "nested" / "stat.hdf5". So the deleted test was redundant rather than coverage quietly dropped, and removing the helper was the better call.
For prune I reconstructed the pre-fix line and ran your new test_prune_keeps_every_checkpoint_below_the_window against both versions instead of reading it:
PRE-FIX FAIL -> deleted steps [1, 2, 3, 4, 5, 6, 7]
POST-FIX PASS
It fails pre-fix for exactly the right reason, so it is a genuine regression test. The training-advanced.md paragraph covers the retention note I raised in the body too.
The other changes in this round check out: _named_model_parameters returns a list, so collecting the backups up front in apply_shadow is safe to re-iterate and is what makes tied weights restore correctly; eval_env_defaults is in scope at the new pt/train/training.py use site; and the argcheck edit is an accuracy fix rather than a narrowing, since DP_COMPILE_INFER and DP_TF32_INFER are only read under deepmd/pt/ while DP_AMP_INFER is read in deepmd/kernels/utils.py and pt_expt/descriptor/dpa4.py.
One blocker left, inline. Getting collection working turned out to expose a regression this PR introduces against its own base, and it is the only red shard remaining.
Still open from my first review, as notes rather than blockers: format_training_message losing step_time and the avg = ... s/step field, which drops that figure for pt_expt, jax and tf2; and the rank == 0 branch of build_checkpoint_stores being unexercised because every test uses the default rank.
…backend" This reverts the pt_expt TF32 policy (99d33ea plus its comment edits in ae72043), restoring the warn-and-ignore behavior on master. The knob is unrelated to this PR's measured speedup (the benchmark card has no TF32 silicon; the whole 1.69x/3.01x gain comes from the contraction fix), its benefit was never measured, and PR deepmodeling#5958 owns the pt_expt training runtime alignment -- including the documented position that pt_expt runs at 'highest' matmul precision. Keeping a second, contradicting implementation here would split ownership of the same policy across two PRs.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
All three findings are fixed, and I checked each against a425d1e59 rather than going by the replies.
The optimizer ordering is restored the way I suggested: the whitelist runs immediately after opt_type, the HybridMuon branch became the exhaustive else, and the trailing raise is gone. What convinces me it is a real fix rather than a plausible one is the CI transition itself -- test_unsupported_optimizer_has_clear_error is pre-existing and untouched, it was the sole red job on c814ddd97, and every check now passes. So it failed pre-fix and passes post-fix on its own merits.
The two changes in that commit I had not asked for also check out. Guarding should_save_checkpoint with save_freq > 0 is logically identical to the old expression for any positive save_freq and only changes the save_freq <= 0 case, which previously divided by zero -- consistent with the resolve_keep_ckpt_count hardening already in this PR. The doc_max_ckpt_keep rewrite and the parallel-training.md edits cover the ema_ckpt_keep inheritance and the ZeRO stage-2/3 constraints, which was the release-note item from my first review.
This is the first fully green run on the branch, so the suite has now actually exercised the pt refactor rather than aborting at collection. Nice work on the checkpoint retention regression in particular -- the max_keep=10 case you added is exactly the shape that was missing.
…through pt_expt assembly (deepmodeling#5960) Users reported that compiled DPA4 training runs ~2x slower on `pt_expt` than on `pt`. This PR is the result of chasing that: a performance bug in how the SO3 contractions were lowered, and a correctness bug where a configured `use_amp: false` was lost while `pt_expt` assembled the model. ## Changes **1. Weight broadcast across the node axis (`so3.py`, `lora.py`, `grid_net.py`)** `matmul(x[..., None, :], weight[None, ...])` makes the node count `N` the matmul BATCH, so matmul broadcasts the weight to `(N, D, F, Cin, Cout)` and autograd then reduces that whole expanded gradient (`ExpandBackward0`) back to the parameter shape. At the water example's sizes a 165 K-element weight expanded to 191 M elements (~0.8 GB) per call, and the reduce was the single costliest kernel of a training step (45.6 ms, 3x per step). The fix batches the contraction over the small `(D, F)` axes so `N` stays the GEMM ROW dimension and the weight is used in place. Micro-benchmark, fwd+bwd at the real shapes: **16.48 ms -> 1.09 ms (15x)**. Two lookalike sites in `projection.py` are deliberately NOT changed: their operands are `requires_grad=False` buffers, so no backward reduce exists. Verified rather than assumed. **2. The same lowering for the frame mixers (`_degree_batched_matmul`)** Review found the `FrameContract` / `FrameExpand` mixers still on the broadcast spelling. They now share one helper, `_degree_batched_matmul`, written identically on the dpmodel side (`dpmodel/descriptor/dpa4_nn/grid_net.py`) and the pt side (`pt/model/descriptor/sezm_nn/grid_net.py`). Because it does no reshape, an empty node axis (`N == 0`) flows through unchanged instead of hitting a reshape error — pinned by a test. **3. `use_amp` was lost during `pt_expt` model assembly (correctness)** `use_amp` is a training-runtime policy, not model state, so it stays OUT of the portable serialization record (this is the deepmodeling#5963 position, and the jax deserializer actively rejects records carrying `use_amp: true`). The bug was elsewhere: `pt_expt` assembled its model by converting an already-populated dpmodel instance, and that conversion round-trips the component through `deserialize(serialize())`. Anything that is deliberately not in the portable record — `use_amp` among it — was therefore dropped, and training silently ran under bfloat16 autocast even when the input configured `use_amp: false`. The fix is at the assembly boundary, not in the record: `pt_expt` now constructs the wrapped class directly (`auto_wrapped_class(...)` in `make_model.py`, `get_model.py`, and the bridging composition path), so a live constructor-supplied component keeps its runtime state. The rule is stated once, in the `auto_wrapped_class` docstring; the call sites reference it. An earlier revision of this PR instead added `use_amp` to `serialize()`. That was reverted in review — it put a runtime knob into the portable record and would have broken the jax contract. **Also removed in review: an `enable_tf32` / `DP_TF32_INFER` implementation for `pt_expt`.** It contributes nothing to the speedup measured below (the benchmark card has no TF32 silicon), and deepmodeling#5958 owns the `pt_expt` training-runtime alignment — including the documented position that `pt_expt` always runs at `"highest"` matmul precision. `pt_expt` therefore keeps master's warn-and-ignore behavior for `enable_tf32`. ## Benchmark DPA4 water example (`examples/water/dpa4`), one Tesla T4, torch 2.11, fp32 (`use_amp: false`), batch size 6. Steady-state seconds per training step, obtained by differencing the wall time of a 33-step and a 3-step run of the same config, which cancels every one-time cost (import, data load, statistics, `torch.compile` / make_fx lowering). All five arms were measured in one session on the same machine; run-to-run variation is about 2-3%. **Provenance: measured at `ae720432b`**, the head at which this PR was opened — i.e. BEFORE the review changes (change 2, the frame-mixer lowering, and change 3's move from `serialize()` to the assembly boundary). Change 1, which is where the entire speedup comes from, is unmodified since. The numbers have not been re-measured on the current head; a re-run is pending and I will post it rather than silently reuse these. | training mode | `pt` (reference) | `pt_expt` at master | `pt_expt` at `ae720432b` | speedup vs master | |---|---|---|---|---| | eager | 0.891 s/step | 1.555 s/step | **0.921 s/step** | **1.69x** | | compiled | 0.545 s/step | 1.611 s/step | **0.535 s/step** | **3.01x** | This reproduces the reported issue at master — `pt_expt` compiled was 3.0x slower than `pt` compiled, and even slower than its own eager path, because the broadcast-weight contraction lowers to worse code under inductor than under eager cuBLAS. After the fix `pt_expt` is at parity with `pt`: eager within 3.4%, compiled within measurement noise. ## Known limitations - **Backward numerics are covered for the frame mixers, not for the SO3 / LoRA contractions.** `test_dpa4_frame_mixers.py` compares `_degree_batched_matmul`'s weight gradient against the pt module's at rtol/atol 1e-12. The rewritten SO3 and LoRA contractions are pinned on the forward against an explicit `einsum` reference (rtol/atol 1e-12, numpy and torch namespaces); their backward is still exercised only by tracing, not compared by value. - **The `pt` / `pt_expt` TF32 policy gap remains open.** On Ampere+ cards `pt` runs training matmuls under TF32 (`enable_tf32`, default `True`) while `pt_expt` ignores the key with a warning; the two backends are not speed-comparable there. Deferred to the deepmodeling#5958 training-runtime series. - **The residual compiled gap vs `pt` is not stable across sessions.** An earlier session measured `pt_expt` compiled 10.8% slower than `pt` compiled; the benchmark above measured it 1.7% faster. Both are within a couple of run-to-run standard deviations, so I treat compiled as at parity and the earlier gap as unconfirmed. - The history contains churn at the GridBranch router (`7518a417c` -> `01c58e665` -> `75459610a` -> `504bb2430` -> `157444204`): a matmul spelling introduced, reverted, reintroduced, and finally restored to master's line. The site is byte-identical to master in the final diff. The degenerate GEMM that profiling found there existed only on this branch, so it is not a fix — I have left the commits rather than rewriting pushed history, and would squash them on request. - Unrelated but found while benchmarking: **torch >= 2.11 ships no Volta (CC 7.0) kernels**, and compiled training requires >= 2.11 via `check_compile_torch_version`. Compiled DPA4 training is therefore impossible on V100 with official wheels; T4 (CC 7.5) is the oldest card that works. ## Tests - `source/tests/common/dpmodel/test_dpa4_frame_mixers.py` — `_degree_batched_matmul` vs the pt module: forward parity, the `N == 0` contract, and weight-gradient parity. - `source/tests/common/dpmodel/test_dpa4_lora.py` — new `test_lora_so3_call_matches_einsum_contract`: `LoRASO3.call` against the explicit `einsum("ndfi,difo->ndfo")` reference with a nonzero adapter, on both the numpy and torch namespaces, for `n_focus` 1 and 2. - `source/tests/pt_expt/model/test_get_model_dpa4.py` — `use_amp` survives model assembly (both branches), for the plain and the bridged/composed construction paths. - `source/tests/common/dpmodel/test_descrpt_dpa4.py` — `use_amp` is absent from the portable serialization record and defaults on deserialize. - Existing `test_grid_branch[1]`/`[2]` cover the changed SO3 contraction against the pt implementation at rtol 1e-12. - Run locally: 434 passed / 10 skipped across the dpa4 dpmodel, pt_expt and cross-backend parity suites, plus the pt_expt model suite. CUDA-gated precision-context cases were run on a T4 (29/29). ## Test status caveat — resolved An earlier revision of this description flagged two locally failing `pt_expt` AOTI-freeze tests (`test_zbl_bridging.py::test_native_spin_with_bridging_graph_freeze_and_deep_eval`, `test_dpa4_zbl_parallel.py::TestBridgedSpinGraphSelfComm::test_freeze_embeds_with_comm_artifact`) as unadjudicated. They are now adjudicated as **pre-existing and environmental, not caused by this branch**: a clean `upstream/master` worktree on the same machine fails both with the identical `InductorError: assert isinstance(index, CppCSEVariable) and index.is_vec` (torch 2.11 CPU-SIMD codegen bug on an `atomic_add` scatter buffer), and both tests pass on this branch with the known workaround `torch._inductor.config.cpp.simdlen = 1` (2 passed). The same bug is already documented in `source/tests/infer/gen_dpa4.py` / `gen_dpa2.py`. --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Summary
ptandpt_exptpt_exptsupport for checkpoint directories and retention ratios, EMA training/checkpoints, EMA full validation, consistent training reports, and restart-safe state restorationzero_stagestrategies aspt, including DDP, ZeRO-1, and FSDP2, with collective checkpoint assembly and optimizer-state restorationmax_ckpt_keepby default while preserving an explicitema_ckpt_keepoverrideMotivation
The
pt_expttrainer currently lacks several operational guarantees available inpt: equivalent checkpoint retention and restart behavior, EMA support, distributed state sharding, stable gradient checks, and consistent progress reporting. Implementing these separately would leave two training runtimes with duplicated policies that can drift.This PR keeps backend-specific serialization and execution in each trainer, while centralizing the policies that are independent of a backend. It also relocates the EMA, validation, and gradient helpers under
pt_expt, which is their continuing owner as the legacypttrainer is retired.Notable fixes
max_ckpt_keep < 1retains all checkpoints instead of deleting the current checkpointckpt_keep_ratioworks when periodic saving is disabled and overrides both regular and EMA windowsema_ckpt_keepnow gives both families themax_ckpt_keepwindowValidation
OMP_NUM_THREADS=1 DP_INTER_OP_PARALLELISM_THREADS=0 DP_INTRA_OP_PARALLELISM_THREADS=0 /Users/outisli/Software/miniforge3/envs/dpmd/bin/python -m pytest source/tests/common/dpmodel/test_train_checkpoint.py source/tests/common/test_argcheck_training.py -q(14 passed)pt_expttraining run with EMA,save_freq=1, and onlymax_ckpt_keep=2retained regular steps4, 5and EMA steps4, 5Summary by CodeRabbit
New Features
Documentation
Bug Fixes