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 4 commits into
Conversation
📝 WalkthroughWalkthroughMultitask training-loss logging now handles unsampled intervals by emitting NaN values for affected loss keys. Display averaging initializes missing task accumulators, computes per-task validation results, and a regression test verifies the resulting learning-curve output. ChangesMultitask loss logging
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 79.06% 78.84% -0.23%
==========================================
Files 1067 1067
Lines 123805 123815 +10
Branches 4510 4515 +5
==========================================
- Hits 97891 97624 -267
- Misses 24304 24581 +277
Partials 1610 1610 ☔ 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
| if not task_input: | ||
| return |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Handle filtered batches when initializing unsampled metrics.
When disp_avg is enabled and a task was not sampled in the interval, get_data() can return {} if min_pair_dist rejects the fetched batch. Returning early leaves train_results[_key] without metric keys. With validation configured, print_on_training() then indexes those missing training metrics and raises a KeyError at the first display; without validation, the header is incomplete and later rows become misaligned.
Please populate the metric schema without relying on one consumable training batch, or loop until a valid batch is returned.
🛠️ Proposed fix
- if not task_input:
- return
+ while not task_input:
+ task_input, task_label, _ = self.get_data(
+ is_train=True, task_key=_task_key
+ )🤖 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 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.
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
| if self.train_loss_accu[_task_key]: | ||
| return | ||
| self.optimizer.zero_grad(set_to_none=True) | ||
| task_input, task_label, _ = self.get_data( |
There was a problem hiding this comment.
[P1] Do not consume a real training batch merely to discover the metric schema
initialize_task_loss_accumulator() calls get_data(is_train=True, task_key=_task_key) for an unsampled task, and then runs only a forward pass. However, get_data() unconditionally advances that task's iterator with next(iterator) (line 2310); this batch is neither cached nor used in backward() / optimizer.step().
For example, with two tasks, disp_avg=true, and disp_freq=1, if step 1 trains model_1, the display path consumes model_2's first batch. If step 2 selects model_2, optimization starts from its second batch. Thus changing a logging-only setting (disp_freq, or whether a task happens to be sampled before the first display) changes the batches used for optimization and breaks reproducibility for a fixed seed.
Please obtain the loss-key schema without advancing the training iterator, or buffer the probe batch so that the next optimization step for this task consumes it. Add a regression test that records fid (or an equivalent batch identifier) and verifies that schema initialization does not alter the subsequent training-batch sequence.
— 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 instead of missing/empty results.disp_avglogging to keep displayed indices and columns consistent, includingNaNentries when averages can’t be computed.disp_avgcorrectly handles unsampled intervals and thatlcurve.outshowsNaNin the expected row/columns.