Skip to content

[trainer, perf] feat: reuse reduced loss denominators #941

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

Merged
Coach257 merged 1 commit into from
Aug 5, 2026
Merged

[trainer, perf] feat: reuse reduced loss denominators #941

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
71 changes: 71 additions & 0 deletions tests/utils/test_loss_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
from types import SimpleNamespace

import pytest
import torch

import veomni.utils.loss_utils as loss_utils


def test_reduce_global_loss_token_reduces_each_denominator_once(monkeypatch):
calls = []

def fake_all_reduce(value, op="mean", group=None):
calls.append((value, op, group))
return value + 10

monkeypatch.setattr(loss_utils, "all_reduce", fake_all_reduce)

reduced = loss_utils.reduce_global_loss_token(
{
"foundation_tokens": torch.tensor(3),
"image_decoder_tokens": torch.tensor(0),
}
)

assert reduced == {"foundation_tokens": 13, "image_decoder_tokens": 10}
assert calls == [(3, "sum", None), (0, "sum", None)]
Comment on lines +9 to +26

@gemini-code-assist gemini-code-assist Bot Jul 17, 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.

high

Update the unit test to match the optimized single-call packed all_reduce implementation and mock the distributed initialization check.

def test_reduce_global_loss_token_reduces_each_denominator_once(monkeypatch):\n    calls = []\n\n    def fake_all_reduce(value, op=\"mean\", group=None):\n        calls.append((value, op, group))\n        return [v + 10 for v in value]\n\n    monkeypatch.setattr(loss_utils, \"all_reduce\", fake_all_reduce)\n    import torch.distributed as dist\n    monkeypatch.setattr(dist, \"is_initialized\", lambda: True)\n\n    reduced = loss_utils.reduce_global_loss_token(\n        {\n            \"foundation_tokens\": torch.tensor(3),\n            \"image_decoder_tokens\": torch.tensor(0),\n        }\n    )\n\n    assert reduced == {\"foundation_tokens\": 13, \"image_decoder_tokens\": 10}\n    assert calls == [([3, 0], \"sum\", None)]



def test_mean_global_loss_reuses_pre_reduced_denominator(monkeypatch):
def fail_all_reduce(*args, **kwargs):
raise AssertionError("denominator should already be reduced")

monkeypatch.setattr(loss_utils, "all_reduce", fail_all_reduce)
monkeypatch.setattr(
loss_utils,
"get_parallel_state",
lambda: SimpleNamespace(sp_enabled=False, fsdp_size=2),
)

loss_dict = loss_utils.mean_global_loss(
{"foundation_loss": torch.tensor(4.0)},
{"foundation_tokens": torch.tensor(3)},
{"foundation_tokens": torch.tensor(5)},
{"foundation_tokens": 6},
)

assert loss_dict["foundation_loss"].item() == pytest.approx(4.0)


def test_mean_global_loss_falls_back_to_reducing_denominator(monkeypatch):
calls = []

def fake_all_reduce(value, op="mean", group=None):
calls.append((value, op, group))
return value * 2

monkeypatch.setattr(loss_utils, "all_reduce", fake_all_reduce)
monkeypatch.setattr(
loss_utils,
"get_parallel_state",
lambda: SimpleNamespace(sp_enabled=False, fsdp_size=1),
)

loss_dict = loss_utils.mean_global_loss(
{"foundation_loss": torch.tensor(4.0)},
{"foundation_tokens": torch.tensor(2)},
{"foundation_tokens": torch.tensor(5)},
)

assert loss_dict["foundation_loss"].item() == pytest.approx(0.8)
assert calls == [(5, "sum", None)]
8 changes: 6 additions & 2 deletions veomni/trainer/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@
is_nccl_backend,
synchronize,
)
from ..utils.loss_utils import count_loss_token, mean_global_loss
from ..utils.loss_utils import count_loss_token, mean_global_loss, reduce_global_loss_token
from ..utils.model_utils import pretty_print_trainable_parameters
from .callbacks import (
ChannelLossCallback,
Expand Down Expand Up @@ -675,7 +675,10 @@ def postforward(
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
"""Postprocess model outputs after forward pass."""
loss_dict: Dict[str, torch.Tensor] = mean_global_loss(
outputs.loss, self.micro_batch_token_len, self.micro_batches_token_len
outputs.loss,
self.micro_batch_token_len,
self.micro_batches_token_len,
getattr(self, "global_micro_batches_token_len", None),

@cursor cursor Bot Jul 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

getattr(self, "global_micro_batches_token_len", None) is a soft contract in two directions.

It never raises, so a trainer that sets micro_batches_token_len but forgets the new line silently loses the optimization rather than failing — fine for safety, invisible for maintenance. More importantly the attribute is only ever set, never cleared, so once train_step has run it stays populated on the trainer for the object's lifetime. Any future postforward reached outside a train_step (an eval path, a debug forward) would pair a fresh micro_batch_token_len with a stale step-level denominator and silently mis-normalize the loss, with no error to point at.

Suggest declaring global_micro_batches_token_len: dict[str, float] | None = None alongside step_train_metrics at line 289 so the plain attribute read works and the getattr default can go, and clearing it back to None at the end of train_step so a stale value can never leak.

)
loss = torch.stack(list(loss_dict.values())).sum()
return loss, loss_dict
Expand Down Expand Up @@ -766,6 +769,7 @@ def train_step(

# token num for fixed_ce_loss in postforward
self.micro_batches_token_len = count_loss_token(micro_batches)
self.global_micro_batches_token_len = reduce_global_loss_token(self.micro_batches_token_len)
num_micro_steps = len(micro_batches)
# forward and backward pass with gradient_accumulationsteps
for micro_step, micro_batch in enumerate(micro_batches):
Expand Down
3 changes: 2 additions & 1 deletion veomni/trainer/text_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
from ..models import build_tokenizer
from ..utils import helper
from ..utils.device import synchronize
from ..utils.loss_utils import count_loss_token
from ..utils.loss_utils import count_loss_token, reduce_global_loss_token
from .base import BaseTrainer, VeOmniIter


Expand Down Expand Up @@ -127,6 +127,7 @@ def train_step(

# token num for fixed_ce_loss in postforward
self.base.micro_batches_token_len = count_loss_token(micro_batches)
self.base.global_micro_batches_token_len = reduce_global_loss_token(self.base.micro_batches_token_len)
num_micro_steps = len(micro_batches)
# forward and backward pass with gradient_accumulationsteps
for micro_step, micro_batch in enumerate(micro_batches):
Expand Down
3 changes: 2 additions & 1 deletion veomni/trainer/vlm_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
from ..optim import build_optimizer
from ..utils import helper
from ..utils.device import synchronize
from ..utils.loss_utils import count_loss_token
from ..utils.loss_utils import count_loss_token, reduce_global_loss_token
from ..utils.model_utils import pretty_print_trainable_parameters
from .base import BaseTrainer, VeOmniIter, _collect_muon_kwargs

Expand Down Expand Up @@ -301,6 +301,7 @@ def train_step(

# token num for fixed_ce_loss in postforward
self.base.micro_batches_token_len = count_loss_token(micro_batches)
self.base.global_micro_batches_token_len = reduce_global_loss_token(self.base.micro_batches_token_len)
num_micro_steps = len(micro_batches)
# forward and backward pass with gradient_accumulationsteps
for micro_step, micro_batch in enumerate(micro_batches):
Expand Down
11 changes: 10 additions & 1 deletion veomni/utils/loss_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ def mean_global_loss(
losses: Union[dict[str, torch.Tensor], torch.Tensor],
micro_batch_token_len: dict[str, torch.Tensor],
micro_batches_token_len: dict[str, torch.Tensor],
global_micro_batches_token_len: dict[str, float] | None = None,
):
"""Calcuate the global mean loss. Avg on all_reduced_token_num instead of on dp_size.
- cur_losses[key] = cur_loss * cur_token_num / global_batches_token_num * get_parallel_state().fsdp_size
Expand All @@ -72,7 +73,10 @@ def mean_global_loss(
if get_parallel_state().sp_enabled:
cur_token_len = all_reduce(cur_token_len.item(), op="sum", group=get_parallel_state().sp_group)

all_reduced_len = all_reduce((micro_batches_token_len[f"{loss_name}_tokens"].item()), op="sum")
if global_micro_batches_token_len is None:
all_reduced_len = all_reduce((micro_batches_token_len[f"{loss_name}_tokens"].item()), op="sum")
else:
all_reduced_len = global_micro_batches_token_len[f"{loss_name}_tokens"]

if all_reduced_len != 0:
cur_loss = cur_loss * cur_token_len / all_reduced_len * get_parallel_state().fsdp_size
Expand All @@ -88,3 +92,8 @@ def mean_global_loss(
loss_dict[key] = cur_loss

return loss_dict


def reduce_global_loss_token(micro_batches_token_len: dict[str, torch.Tensor]) -> dict[str, float]:
"""All-reduce per-step loss-token denominators once for reuse across micro-batches."""
return {key: all_reduce(value.item(), op="sum") for key, value in micro_batches_token_len.items()}
Comment on lines +97 to +99

@gemini-code-assist gemini-code-assist Bot Jul 17, 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.

high

Performance & Robustness Improvement: Pack All-Reduce Calls and Add Distributed Guard\n\n1. Performance: Currently, reduce_global_loss_token performs a separate sequential all_reduce call for each key in micro_batches_token_len. Since there are always at least two keys (foundation_tokens and image_decoder_tokens), this results in multiple sequential communication round-trips per training step. We can pack all values into a single list and perform a single all_reduce call, which significantly reduces communication overhead.\n2. Robustness: If training is run in a non-distributed environment (where torch.distributed.is_initialized() is False), calling all_reduce will raise a RuntimeError. Adding a guard ensures the code falls back gracefully to local values.

def reduce_global_loss_token(micro_batches_token_len: dict[str, torch.Tensor]) -> dict[str, float]:\n    \"\"\"All-reduce per-step loss-token denominators once for reuse across micro-batches.\"\"\"\n    if not micro_batches_token_len:\n        return {}\n\n    import torch.distributed as dist\n\n    keys = list(micro_batches_token_len.keys())\n    values = [micro_batches_token_len[k].item() for k in keys]\n\n    if not dist.is_initialized():\n        return dict(zip(keys, values))\n\n    reduced_values = all_reduce(values, op=\"sum\")\n    if len(keys) == 1:\n        return {keys[0]: reduced_values}\n    return dict(zip(keys, reduced_values))

@cursor cursor Bot Jul 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Two things about iterating micro_batches_token_len here rather than letting mean_global_loss drive.

Latent: the collective sequence becomes data-dependent. On main the reduce at line 75 is driven by losses.items() — the model's loss keys, which are the same on every rank. This helper instead reduces every key in count_loss_token's output, and that key set is data-dependent: count_loss_token starts from two fixed keys but then adds f"{name}_tokens" for any batch key ending in _labels (loss_utils.py:37-40). If two DP ranks' micro-batches disagree on those keys, they issue a different number of world-group all-reduces and the run desynchronizes.

I checked and no collator in veomni/data/ currently emits a <modality>_labels key, so today the set is always exactly {foundation_tokens, image_decoder_tokens} on every rank and this cannot fire. But that scan clearly exists so modalities can add keys, and whoever adds one will not connect it to a hang here.

Minor: it reduces keys nobody consumes. A text-only run has one loss key but two token keys, so this issues 2 collectives per step where 1 suffices.

Both go away with a slightly different shape — keep the reduce driven by losses, but memoize it per step inside mean_global_loss, e.g. cache on the identity of the micro_batches_token_len dict (it is rebuilt once per train_step) or pass an explicit per-step cache dict that starts empty. Each denominator is then still reduced exactly once per step, the collective order stays exactly as on main, and veomni/trainer/base.py, text_trainer.py, and vlm_trainer.py need no changes at all. Given constraint §24 on keeping composed trainers in sync, removing three duplicated call sites is worth something on its own.

If you prefer to keep the current shape, please at least reduce only the keys mean_global_loss will actually look up, and add a short comment stating that every rank must produce an identical key set.

Loading