fix(disp): div 0 error when using disp_avg under multitask - #5809
fix(disp): div 0 error when using disp_avg under multitask#5809OutisLi wants to merge 3 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (25)
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughTraining display averaging now uses declared per-task metric names and a shared accumulator. PyTorch and PyTorch Exportable trainers average detached training metrics per display interval, render ChangesTraining metric display averaging
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~45 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant Trainer
participant TrainingMetricAccumulator
participant LearningCurveWriter
Trainer->>TrainingMetricAccumulator: Add detached metrics for each task step
Trainer->>TrainingMetricAccumulator: Request per-task interval averages
TrainingMetricAccumulator-->>Trainer: Return averages or NaN for unsampled tasks
Trainer->>LearningCurveWriter: Log training and validation results
Suggested reviewers: Merge Risk: ⚪ Minimal · up to The display-averaging changes cover unsampled tasks, metric schemas, interval resets, and learning-curve alignment without an identified merge-blocking risk. 🚥 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: 2
🤖 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/pt/train/training.py`:
- Around line 1619-1632: Update the multitask display-step handling in run()
around the disp_avg branch: perform a dummy forward pass for each unsampled
model key to populate self.train_loss_accu with its loss keys before generating
train_results, excluding l2_ fields as in the proposed flow. Move the validation
and rank-zero console logging loop outside the disp_avg conditional so
log_loss_valid and progress messages execute for both averaged and non-averaged
modes.
In `@source/tests/pt/test_multitask.py`:
- Around line 273-277: Update the lcurve.out assertions in the multitask test to
read and tokenize the header row, then assert every data row has exactly the
header’s column count. Preserve the existing displayed_steps and “nan” checks
while ensuring the unsampled task does not produce a shorter row.
🪄 Autofix (Beta)
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
Run ID: 5ac6ebae-2ba6-49ed-b897-c162751bd612
📒 Files selected for processing (2)
deepmd/pt/train/training.pysource/tests/pt/test_multitask.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Requesting changes because the filtered-batch path can still leave an unsampled task without the metric schema required by learning-curve output.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5809 +/- ##
==========================================
- Coverage 77.25% 77.06% -0.19%
==========================================
Files 1153 1154 +1
Lines 138930 138975 +45
Branches 5056 5056
==========================================
- Hits 107328 107100 -228
- Misses 29717 29990 +273
Partials 1885 1885 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/pt/train/training.py`:
- Around line 1715-1716: Update the unsampled-task initialization around the `if
not task_input` early return so filtered `{}` batches cannot leave
`train_results[_key]` without its metric keys. Populate the expected metric
schema independently of a consumable training batch, or continue fetching until
`get_data()` returns a valid batch, while preserving the existing `disp_avg` and
validation display behavior.
🪄 Autofix (Beta)
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
Run ID: 28d64a76-12e4-4629-83a0-3c8a68c14e06
📒 Files selected for processing (2)
deepmd/pt/train/training.pysource/tests/pt/test_multitask.py
🚧 Files skipped from review as they are similar to previous changes (1)
- source/tests/pt/test_multitask.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Request changes: the unsampled-task schema is still not guaranteed when a training batch is filtered out by min_pair_dist.
initialize_task_loss_accumulator() returns immediately on an empty task_input. That leaves train_results[_task_key] empty. At the first display, print_on_training() iterates the validation metric keys and indexes the corresponding missing training keys, causing a KeyError; without validation, the generated lcurve.out header is incomplete and later rows can become misaligned.
Please populate the placeholder metric schema without depending on a single consumable batch (or keep fetching until one is usable), and add a regression test covering an unsampled task whose initialization batch is fully filtered.
All CI checks are otherwise passing.
— OpenClaw 2026.6.11
njzjz
left a comment
njzjz
left a comment
There was a problem hiding this comment.
Thanks — the second half of this is a clear bug fix, but the accumulator seeding worries me.
The logging hoist is right
On master the for _key in self.model_keys: loop that calls log_loss_valid and emits the _trn/_val lines sits inside the else: (non-disp_avg) branch, so a multi-task run with disp_avg: true produced no per-task log lines at all — train_results was filled and then dropped on the floor. Hoisting it to its own loop over model_keys fixes that, and threading check_total_rmse_nan=False for a task with step_count_per_task == 0 is the right way to keep the deliberate NaN from tripping the NaN guard. dict.fromkeys(task_losses, float("nan")) for the zero-step case is a clean replacement for the silent empty dict.
initialize_task_loss_accumulator runs a real training step to learn column names
self.optimizer.zero_grad(set_to_none=True)
task_input, task_label, _ = self.get_data(is_train=True, task_key=_task_key)
if not task_input:
return
_, _, task_more_loss = self.wrapper(**task_input, cur_lr=pref_lr, label=task_label, task_key=_task_key)I understand why: lcurve.out writes its header once, so every task's column set has to be known at the first display step, and more_loss keys only exist after a forward. But this pays for that with four side effects on a path that is supposed to be pure reporting:
- It consumes a training batch.
get_data(is_train=True, task_key=...)advances that task's training iterator. So the task the display is about to report as "not sampled this interval" is sampled — the batch is drawn, fed forward, and thrown away. That shifts the data stream and epoch bookkeeping for that task by one batch per display step until it is first sampled naturally. Your own test has to mockdp_random.choicewith an exact[0, 1]sequence, which is a hint at how sensitive this is. - No
torch.no_grad(). The enclosing block runs afterself.wrapper.eval(), so this forward builds a full autograd graph in eval mode and discards it — wasted memory and time on every display step until the task is seeded, and for a large task that is not cheap. self.optimizer.zero_grad(set_to_none=True)mutates optimizer state from inside the display path. The non-disp_avgbranch does the same thing on master, so there is precedent, but that branch is at least computing numbers it then reports; here the loss is discarded.- DDP. The
if not task_input: returnearly exit is evaluated per rank. If it can ever be true on some ranks and not others, the ranks disagree on whether to run a forward and the collectives desynchronize. Worth confirmingget_datacannot return a falsytask_inputon a subset of ranks.
Two directions that avoid all four:
- Keep the keys, not the counters. The reset at the end of the display block already zeroes values and preserves keys, so a task only needs seeding if it has never been sampled since training started. If the header could be deferred until every task has been seen once, or written with a per-task placeholder set derived from the loss configuration (which terms are enabled: energy / force / virial / …) rather than from an executed
more_loss, no forward is needed at all. - If a forward really is unavoidable, at minimum wrap it in
torch.no_grad(), and draw from the validation loader rather than the training one so the training stream is untouched.
Smaller points
initialize_task_loss_accumulatoris redefined on every display step; it does not close over anything that changes exceptpref_lr, so it could be a method or moved above the loop.- The test asserts
self.assertIn("nan", data_lines[1])— that checks the literal token appears somewhere in the row. Asserting that the NaN falls in that task's columns (viaheader_columns.index(...)) would pin the actual contract; as written it would still pass if the NaN landed in the wrong task's column. - Worth a line in the test docstring or a comment saying which task is unsampled and why, so the
[0, 1]mock sequence is not load-bearing but unexplained.
Summary by CodeRabbit
NaNvalues rather than incomplete results.disp_avgtraining option.