Repository navigation
[trainer, perf] feat: reuse reduced loss denominators - #941
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes the training process by pre-reducing the loss-token denominators once per step and reusing them across micro-batches, thereby avoiding redundant communication overhead in mean_global_loss. The review feedback suggests further optimizing this by packing all values into a single all_reduce call to minimize communication round-trips, adding a guard for non-distributed environments, and updating the corresponding unit tests to reflect these changes.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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()} |
There was a problem hiding this comment.
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))| 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)] |
There was a problem hiding this comment.
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)]
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Review summary — approve (with a design suggestion worth considering first)
Baseline: main@110900b3d44d, ahead 1, behind 18. Reviewed read-only (no checkout of the PR branch; PR file versions were fetched over tracked paths, tested, then reverted).
1. Does it solve a real problem? Yes, and this is the biggest of the three
mean_global_loss is called once per micro-batch from postforward, and on main line 75 all-reduces micro_batches_token_len[...] on every one of those calls — a per-step quantity that cannot change between micro-batches. Each call also goes through .item(), so it is a device sync too. With gradient accumulation of 8 and a single loss key that is 8 collectives plus 8 syncs per step where 1 would do. Hoisting to train_step is the right call, and unlike the sibling PRs this one is on a genuinely hot path.
The reused value is provably identical: same expression, same op="sum", same default (world) group, and micro_batches_token_len is computed once per step and not mutated during the micro-batch loop.
2. Comments / docs / tests
No docs needed. The three tests are meaningful: test_mean_global_loss_reuses_pre_reduced_denominator monkeypatches all_reduce to raise, which proves no collective is issued on the fast path, and it still checks the arithmetic (4.0 * 3 / 6 * 2 == 4.0) so a silently-wrong denominator would fail. The fallback test confirms the None path is unchanged. All 3 pass; ruff check and ruff format --check are clean.
What is not covered is the trainer wiring itself — the three train_step assignments and the getattr read in postforward — which is understandable since that needs a distributed harness, but it is where the risk lives (see below).
3. Cleanliness / conflicts
18 commits behind main; please rebase. No file overlap with #942 or #943, so all three can land independently.
The wiring is currently complete: BaseTrainer, text_trainer, and vlm_trainer are the only places that set micro_batches_token_len, and all three are updated. dit_trainer overrides postforward entirely and text_dpo_trainer never touches these attributes, so neither is affected.
My main suggestion is structural and is in the inline comment: the same saving is achievable by memoizing inside mean_global_loss, which would need no trainer changes at all, would keep the collective sequence byte-for-byte identical to main, and would remove the getattr fallback and the three-way duplication that the next trainer author has to remember to replicate.
Sent by Cursor Automation: VeOmni-Review-Agent(Outer)
|
|
||
| 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()} |
There was a problem hiding this comment.
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.
| outputs.loss, | ||
| self.micro_batch_token_len, | ||
| self.micro_batches_token_len, | ||
| getattr(self, "global_micro_batches_token_len", None), |
There was a problem hiding this comment.
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.
Summary
reduce_global_loss_token()to reduce per-step loss-token denominators once per optimizer step.mean_global_loss()across gradient-accumulation micro-batches.mean_global_loss()backward-compatible by falling back to the previous per-call denominator reduction when no cached denominator is supplied.Performance Rationale
Before this change, each micro-batch loss normalization all-reduced the same global denominator for the whole optimizer step. With gradient accumulation and multiple loss keys, that repeats small control-plane reductions in the inner training loop. This PR keeps the same math but moves the step-level denominator reduction out of the micro-batch loop for Base/Text/VLM training paths.
Validation
PYTHONPATH=/Users/bytedance/Documents/VeOmni-loss-denom PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH pytest tests/utils/test_loss_utils.py -q->3 passedPYTHONPATH=/Users/bytedance/Documents/VeOmni-loss-denom PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH make quality-> ruff check and format passedPYTHONPATH=/Users/bytedance/Documents/VeOmni-loss-denom PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH python tests/special_sanity/check_device_api_usage.py --directory ./veomni-> passed/veomni-reviewverdict: safeNotes: local CPU cannot run meaningful distributed SP/VLM smoke tests. The hardware-side follow-up is an SP run with
ulysses_size>1and a multi-loss VLM/omni batch to confirm denominator reductions drop from per-micro-batch to per-step.