Skip to content

fix(disp): div 0 error when using disp_avg under multitask #5809

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
OutisLi wants to merge 3 commits into
base: master
Choose a base branch
from
+87 −33
Open

fix(disp): div 0 error when using disp_avg under multitask #5809

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
620847a
fix(disp): div 0 error when using disp_avg under multitask
OutisLi Jul 10, 2026
abfddfb
fix(pt): align multitask display metrics
OutisLi Jul 16, 2026
a73ebb0
Merge branch 'master' into pr/disp
OutisLi Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 56 additions & 33 deletions deepmd/pt/train/training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1616,26 +1616,20 @@ def fake_model() -> dict:
def log_loss_train(
_loss: Any, _more_loss: Any, _task_key: str = "Default"
) -> dict:
results = {}
if not self.multi_task:
# Use accumulated average loss for single task
for item in self.train_loss_accu:
results[item] = (
self.train_loss_accu[item]
/ self.step_count_in_interval
)
else:
# Use accumulated average loss for multi-task
if (
_task_key in self.train_loss_accu
and _task_key in self.step_count_per_task
):
for item in self.train_loss_accu[_task_key]:
results[item] = (
self.train_loss_accu[_task_key][item]
/ self.step_count_per_task[_task_key]
)
return results
return {
item: value / self.step_count_in_interval
for item, value in self.train_loss_accu.items()
}

task_losses = self.train_loss_accu.get(_task_key, {})
step_count = self.step_count_per_task.get(_task_key, 0)
if step_count == 0:
return dict.fromkeys(task_losses, float("nan"))
return {
item: value / step_count
for item, value in task_losses.items()
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
else:

def log_loss_train(
Expand Down Expand Up @@ -1709,7 +1703,31 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
valid_results = {_key: {} for _key in self.model_keys}
if self.disp_avg:
# For multi-task, use accumulated average loss for all tasks
def initialize_task_loss_accumulator(
_task_key: str,
) -> None:
if self.train_loss_accu[_task_key]:
return
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
Comment thread
OutisLi marked this conversation as resolved.
Comment on lines +1715 to +1716

@coderabbitai coderabbitai Bot Jul 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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.

_, _, task_more_loss = self.wrapper(
**task_input,
cur_lr=pref_lr,
label=task_label,
task_key=_task_key,
)
self.train_loss_accu[_task_key] = {
item: 0.0
for item in task_more_loss
if "l2_" not in item
}

for _key in self.model_keys:
initialize_task_loss_accumulator(_key)
train_results[_key] = log_loss_train(
loss, more_loss, _task_key=_key
)
Expand All @@ -1732,25 +1750,30 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
train_results[_key] = log_loss_train(
loss, more_loss, _task_key=_key
)
valid_results[_key] = log_loss_valid(_task_key=_key)
if self.rank == 0:
for _key in self.model_keys:
valid_results[_key] = log_loss_valid(_task_key=_key)
if self.rank == 0:
log.info(
format_training_message_per_task(
batch=display_step_id,
task_name=_key + "_trn",
rmse=train_results[_key],
learning_rate=cur_lr,
check_total_rmse_nan=not (
self.disp_avg
and self.step_count_per_task.get(_key, 0) == 0
),
)
)
if valid_results[_key]:
log.info(
format_training_message_per_task(
batch=display_step_id,
task_name=_key + "_trn",
rmse=train_results[_key],
learning_rate=cur_lr,
task_name=_key + "_val",
rmse=valid_results[_key],
learning_rate=None,
)
)
if valid_results[_key]:
log.info(
format_training_message_per_task(
batch=display_step_id,
task_name=_key + "_val",
rmse=valid_results[_key],
learning_rate=None,
)
)
self.wrapper.train()

if self.disp_avg:
Expand Down
31 changes: 31 additions & 0 deletions source/tests/pt/test_multitask.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@
from pathlib import (
Path,
)
from unittest.mock import (
patch,
)

import torch

Expand Down Expand Up @@ -252,6 +255,34 @@ def setUp(self) -> None:
self.config["model"]
)

def test_disp_avg_handles_unsampled_interval(self) -> None:
config = deepcopy(self.config)
config["training"]["numb_steps"] = 2
config["training"]["disp_freq"] = 1
config["training"]["disp_avg"] = True
config = update_deepmd_input(config, warning=True)
config = normalize(config, multi_task=True)
trainer = get_trainer(config, shared_links=self.shared_links)

with patch(
"deepmd.pt.train.training.dp_random.choice",
side_effect=[0, 1],
):
trainer.run()

with open("lcurve.out") as f:
lines = f.readlines()
header_lines = [line.split() for line in lines if line.startswith("#")]
data_lines = [line.split() for line in lines if not line.startswith("#")]
self.assertTrue(header_lines)
header_columns = header_lines[0][1:]
self.assertTrue(any("_val_" in column for column in header_columns))
for columns in data_lines:
self.assertEqual(len(columns), len(header_columns))
displayed_steps = [int(columns[0]) for columns in data_lines]
self.assertEqual(displayed_steps, [1, 2])
self.assertIn("nan", data_lines[1])
Comment thread
coderabbitai[bot] marked this conversation as resolved.

def tearDown(self) -> None:
MultiTaskTrainTest.tearDown(self)

Expand Down
Loading