[zhangxin81]

Summary

• Coalesce scalar training metric all-reduces in EnvironMeterCallback into a single collective when possible.
• Preserve the previous per-key fallback for non-numeric/custom metric values.
• Add focused tests for coalesced numeric reductions and fallback behavior.

Performance Rationale

EnvironMeterCallback previously reduced total_loss, every loss entry, and grad_norm with one scalar all-reduce per key on every step. These are small control-plane collectives on the hot path. Coalescing numeric scalar metrics into one list all-reduce preserves the logged values while reducing per-step collective count.

Validation

• PYTHONPATH=/Users/bytedance/Documents/VeOmni-metric-reduce PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH pytest tests/trainer/test_trace_callback_metrics.py -q -> 2 passed
• PYTHONPATH=/Users/bytedance/Documents/VeOmni-metric-reduce PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH make quality -> ruff check and format passed
• /veomni-review verdict: safe

Notes: local CPU cannot run a meaningful distributed trainer smoke. The hardware-side follow-up is a small distributed run that exercises EnvironMeterCallback with tqdm/wandb consuming step_train_metrics and step_env_metrics.

[gemini-code-assist]

Code Review

This pull request introduces a helper function _reduce_training_metrics in trace_callback.py to coalesce and reduce scalar training metrics in a single collective call, along with corresponding unit tests. The reviewer identified critical issues with this implementation: it will crash if there is only a single numeric metric (due to all_reduce returning a scalar instead of an iterable), and coalescing is disabled if any metric is a torch.Tensor. Additionally, calling all_reduce on non-reducible types will cause runtime crashes. The reviewer provided robust code suggestions to safely handle single-element reductions, tensor types, and non-reducible values, as well as updates to the unit tests to reflect these improvements.

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.

[gemini-code-assist]

There are two major issues with the current implementation of _reduce_training_metrics:

1. Single-Element Crash (Critical): If metrics contains only a single numeric element, all_reduce is called with a list of length 1. According to all_reduce implementation in dist_utils.py, if data.numel() == 1, it returns data.item() (a scalar float) instead of a list. This causes zip(keys, reduced_values) to raise a TypeError: zip argument #2 must support iteration and crash the training loop.

2. Coalescing Disablement & Crash on Non-Reducible Types (High): If any metric is a torch.Tensor (which is common for metrics like grad_norm or loss), all(isinstance(value, (int, float))) evaluates to False. This triggers the fallback, disabling coalescing entirely for all other numeric metrics. Furthermore, if an actual non-reducible object (like a string or custom object) is present, calling all_reduce on it will crash at runtime since all_reduce cannot convert arbitrary objects to tensors.

We can resolve all these issues by cleanly separating coalescable numeric/scalar metrics from other metrics, and handling single-element reductions safely.

def _reduce_training_metrics(metrics: Dict[str, Any], group) -> Dict[str, Any]:
    """Reduce scalar training metrics with one collective when possible."""
    if not metrics:
        return {}

    import torch

    coalesced_keys = []
    coalesced_values = []
    results = {}

    for key, value in metrics.items():
        if isinstance(value, (int, float)):
            coalesced_keys.append(key)
            coalesced_values.append(value)
        elif isinstance(value, torch.Tensor) and value.numel() == 1:
            coalesced_keys.append(key)
            coalesced_values.append(value.item())
        elif isinstance(value, torch.Tensor) or (isinstance(value, list) and all(isinstance(x, (int, float)) for x in value)):
            results[key] = all_reduce(value, group=group)
        else:
            results[key] = value

    if coalesced_values:
        if len(coalesced_values) == 1:
            results[coalesced_keys[0]] = all_reduce(coalesced_values[0], group=group)
        else:
            reduced_coalesced = all_reduce(coalesced_values, group=group)
            for k, v in zip(coalesced_keys, reduced_coalesced):
                results[k] = v

    return results

[gemini-code-assist]

Update the test to reflect the safe and robust behavior where non-reducible values are preserved without calling all_reduce (which would crash in real life), while numeric values are still successfully coalesced and reduced.

Suggested change

	def test_reduce_training_metrics_falls_back_for_non_numeric_values(monkeypatch):
	calls = []
	def fake_all_reduce(value, group=None):
	calls.append((value, group))
	return value
	monkeypatch.setattr(trace_callback, "all_reduce", fake_all_reduce)
	metrics = {"loss": 2.0, "custom": object()}
	reduced = trace_callback._reduce_training_metrics(metrics, group="fsdp")
	assert reduced == metrics
	assert calls == [(metrics["loss"], "fsdp"), (metrics["custom"], "fsdp")]
	def test_reduce_training_metrics_handles_non_reducible_values(monkeypatch):
	calls = []
	def fake_all_reduce(value, group=None):
	calls.append((value, group))
	if isinstance(value, list):
	return [item + 1 for item in value]
	return value + 1
	monkeypatch.setattr(trace_callback, "all_reduce", fake_all_reduce)
	custom_obj = object()
	metrics = {"loss": 2.0, "custom": custom_obj}
	reduced = trace_callback._reduce_training_metrics(metrics, group="fsdp")
	assert reduced == {"loss": 3.0, "custom": custom_obj}
	assert calls == [(2.0, "fsdp")]

[Coach257]

Agree with the single-element concern: all_reduce([x]) returns a scalar via numel()==1 → .item(), so zip(keys, reduced_values) will TypeError. Current EnvironMeter path usually has ≥2 float metrics so it may not hit today, but the helper should still handle len==1 (and ideally coerce 0-dim tensors with .item() before coalescing). Also, the non-numeric fallback that calls all_reduce on every value is unsafe for non-reducible objects — please coalesce only numeric scalars and leave/skip the rest.

[cursor]

🤖 Thanks for the contribution! Cursor has started reviewing this PR and will post the full review shortly.

  

_{Sent by Cursor Automation: VeOmni-Review-Agent(Outer)}

[cursor]

Review summary — needs changes (one-line fix)

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 the optimization does fire

On main, on_step_end calls all_reduce once per metric. Each call builds a one-element tensor, runs a collective, and ends in .item() — a full device sync. Coalescing to one collective plus one sync per step is a genuine win, small in absolute terms since this runs per optimizer step rather than per micro-batch.

I checked the thing that would have made this a no-op: whether the values are actually Python numbers at runtime. They are. total_loss accumulates loss.item() (veomni/trainer/base.py:816), total_loss_dict accumulates v.item() (:818), and veomni_clip_grad_norm returns grad_norm.item() if hasattr(grad_norm, "item") else float(grad_norm) (veomni/distributed/clip_grad_norm.py:19). So the isinstance(value, (int, float)) gate passes and the batched path is taken in real training — it does not silently fall back.

The list input is also genuinely supported: all_reduce wraps non-tensors with torch.tensor(data, ...) (veomni/utils/dist_utils.py:50-51), so a list of floats reduces elementwise. That part is sound.

2. Comments / docs / tests — the gap that hides the bug below

No docs needed. But both tests monkeypatch all_reduce away, so they only assert the shape of the call (calls == [([2.0, 3.0], "fsdp")]) and never exercise the real function. That is exactly why the single-metric defect below is not caught. all_reduce returns data.item() when data.numel() == 1 and data.tolist() otherwise (dist_utils.py:63-66) — a return-type switch that only the real implementation exhibits.

A test against the real all_reduce is cheap here: dist.init_process_group(backend="gloo") with WORLD_SIZE=1 is enough to exercise the whole helper on CPU. That is how I reproduced the crash below.

3. Cleanliness / conflicts

ruff check veomni tests and ruff format --check are clean; both new tests pass. 18 commits behind main — please rebase. Two nits: group has no type annotation while the rest of the signature is annotated, and the declared return type Dict[str, float] does not hold on the non-numeric fallback branch, which returns whatever all_reduce gives back.

Related: #941 and #943 from the same author are the same class of optimization at other layers — none of the three touch overlapping files, so they can land independently.

  

_{Sent by Cursor Automation: VeOmni-Review-Agent(Outer)}

[cursor]

Latent crash when metrics has exactly one entry.

all_reduce returns a scalar rather than a list when the tensor has a single element:

if data.numel() == 1:
    return data.item()
else:
    return data.tolist()

(veomni/utils/dist_utils.py:63-66)

So a one-key dict makes all_reduce([v]) return a bare float, and line 42's dict(zip(keys, reduced_values)) then tries to iterate it. Reproduced against the real all_reduce with a single-process gloo group:

two keys : {'total_loss': 2.0, 'grad_norm': 3.0}
one key  : ERROR TypeError 'float' object is not iterable
empty    : {}

This is not reachable from the current call site — step_train_metrics always carries at least total_loss and grad_norm (lines 213-217) — but _reduce_training_metrics is a module-level helper with a generic name and a docstring that reads as general-purpose, so the next caller will hit it. Worth hardening now since it is one line:

reduced_values = all_reduce(values, group=group)
if not isinstance(reduced_values, list):
    reduced_values = [reduced_values]
return dict(zip(keys, reduced_values))

or simply if len(keys) == 1: return {keys[0]: all_reduce(values[0], group=group)} before the batched call. Please add a test for the single-key case that uses the real all_reduce rather than the mock.