[zhangxin81]

Summary

• Compute the shifted-label valid-token denominator once in chunk_loss_function.
• Reuse that scalar across all chunk loss calls instead of recomputing (labels != ignore_index).sum() for every chunk.
• Add a focused regression test that counts denominator reductions across multiple chunks.

Performance Rationale

The chunked CE path uses the full post-shift label denominator for every chunk. Previously each chunk rebuilt the same boolean mask sum, adding repeated tensor work in the inner chunk loop. This keeps the math and scaling identical while reducing repeated work to one denominator calculation per loss call.

Validation

• PYTHONPATH=/Users/bytedance/Documents/VeOmni-chunk-loss-denom PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH pytest tests/ops/test_chunk_loss.py -q -> 1 passed
• PYTHONPATH=/Users/bytedance/Documents/VeOmni-chunk-loss-denom PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH make quality -> ruff check and format passed
• /veomni-review verdict: safe

Notes: tests/ops/test_seqcls_loss.py is blocked locally because liger_kernel is not installed; this is unrelated to the chunk-loss path.

[gemini-code-assist]

Code Review

This pull request optimizes the chunk loss function by computing the number of items in the batch once and reusing it across chunks, rather than recalculating it for each chunk. A corresponding unit test was added to verify this behavior. The review feedback points out that unconditionally overwriting num_items_in_batch ignores any user-provided value, and suggests conditionally assigning it only when it is None.

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]

The chunk_loss_function accepts an optional num_items_in_batch argument. Unconditionally overwriting it with (labels != ignore_index).sum() ignores any user-provided value for num_items_in_batch. This leads to incorrect loss scaling when a custom or global denominator is passed (e.g., in distributed training). Using a conditional assignment preserves the passed argument when it is not None.

Suggested change

	num_items_in_batch = (labels != ignore_index).sum()
	if num_items_in_batch is None:
	num_items_in_batch = (labels != ignore_index).sum()

[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 — approve (one naming follow-up)

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, small but real

On main, (labels != ignore_index).sum() sits inside the loss_kwargs_chunks list comprehension, so it is recomputed once per chunk over the full label tensor even though the value is identical every time. Hoisting it is a straightforward win. I measured it with the PR's own counter: on main the reduction fires 3 times for a 6-token sequence at chunk_size=2; with the PR it fires once. At production shapes (32k sequence, chunk_size=1024) that is ~32 redundant full-length reductions plus 32 throwaway allocations per micro-batch.

The rewrite is equivalence-preserving: it is the same expression, labels is not mutated between iterations, and the resulting int64 tensor is requires_grad=False, so sharing one object across all chunks introduces no autograd retention.

2. Comments / docs / tests

No docs needed — this changes no API or config. The added tests/ops/test_chunk_loss.py is a genuine regression test, not a tautology: reverting only chunk_loss.py to main makes it fail with assert 3 == 1, and it also asserts both gradients are populated so the refactor cannot silently break the backward. Monkeypatching torch.Tensor.sum globally and filtering on dtype == torch.bool and shape == (1, 5) is a bit coupled to the internal labels[..., 1:] shift, but a shape change would drive the counter to 0 and fail loudly rather than pass vacuously, so it degrades safely.

3. Cleanliness / conflicts

ruff check veomni tests and ruff format --check are clean; tests/ops/test_chunk_loss.py passes.

• Rebase. 18 commits behind main.
• Conflicts with #964. That PR rewrites ChunkLoss.forward and adds a block right after ChunkLoss.apply in this same function, a few lines from your hunk. Whichever lands second needs a manual rebase — worth cross-linking the two.
• Related to #941 (reuse reduced loss denominators) from the same author; the two are conceptually the same optimization at different layers.

See the inline note on the variable name — it is the only thing I would change before merge, and it is cosmetic.

  

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

[cursor]

Reusing the name num_items_in_batch rebinds chunk_loss_function's own parameter (line 95, num_items_in_batch: Optional[int] = None). The behavior is unchanged — main already discarded the caller's value by inlining the same expression — but this line now makes the parameter look like it is being consumed when it is in fact being overwritten.

That matters because the value can actually arrive: _chunk_loss_dispatch forwards *args, **kwargs verbatim, and num_items_in_batch is the 4th positional slot in the loss-function contract (see _LOSS_CALL_POSITIONAL_KEYS in veomni/trainer/callbacks/channel_loss_callback.py:54), so a caller doing gradient-accumulation-correct normalization would have its denominator silently dropped.

Two options, either is fine:

• Cosmetic: name the local something like valid_token_count so the dead parameter stays visibly dead.
• Or actually honor it: denominator = num_items_in_batch if num_items_in_batch is not None else (labels != ignore_index).sum().

If the intent is that chunk_loss_function deliberately owns its own denominator, a one-line comment saying so would prevent the next reader from filing this as a bug.