perf(megatron): remove per-document device syncs in channel loss aggregation

[gakkiri]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support
• Performance Optimization

PR information

When --enable_channel_loss is used together with padding_free, MegatronTrainer._compute_channel_loss aggregates the per-channel loss with a Python loop over every document in the pack. Each iteration forces three GPU -> CPU synchronizations:

• slice(cu_seqlens[i], cu_seqlens[i + 1]) converts two CUDA scalars via __index__, i.e. two implicit item() calls;
• the boolean mask index losses[0, slice_][loss_mask[0, slice_]] needs nonzero to resolve the output length, and c_loss.shape[0] depends on that same result.

So the cost of this function grows linearly with the number of documents in a pack. On the machine used for the benchmark below it is ~113 us per document, and since the function runs inside loss_func for every micro-batch on the last PP stage, the overhead is paid grad_accum times per step. Datasets made of many short samples are hit hardest, which is exactly the case where channel loss is most useful. Beyond the raw microseconds, the syncs also destroy CPU run-ahead, so the last PP stage cannot enqueue the next micro-batch in time and the pipeline bubble grows.

This PR replaces the loop with a single vectorized reduction:

• token -> channel ids are built with repeat_interleave over the segment lengths, with output_size passed explicitly so this call does not synchronize either;
• loss sum and token count are accumulated together by one index_add_ into a (num_channels, 2) buffer.

The only remaining synchronization is int(cu_seqlens[-1]), once per call instead of three per document, so the cost is now essentially independent of the document count.

Experiment results

Microbenchmark of the aggregation itself, isolated from the rest of the training step. Total pack length is fixed at 131072 tokens and split into a varying number of documents; timings are per call.

Documents	Avg doc len	Non-masked tokens	Before	After	Speedup
128	1024	25%	14.902 ms	0.308 ms	48.4x
128	1024	100%	14.776 ms	0.309 ms	47.8x
1024	128	25%	115.204 ms	0.376 ms	306.6x
1024	128	100%	115.230 ms	0.375 ms	307.2x
4096	32	25%	461.310 ms	0.589 ms	782.6x
4096	32	100%	461.487 ms	0.590 ms	781.7x
16384	8	25%	1847.029 ms	1.339 ms	1379.8x
16384	8	100%	1926.336 ms	1.346 ms	1430.9x

The 16384-document case was re-measured as a median over 10 rounds to rule out noise, and stays around 1340x: 1831.777 ms -> 1.366 ms at 25% non-masked tokens, 1840.980 ms -> 1.367 ms at 100%.

The old timings scale strictly linearly with the document count (~113 us per document across all rows), which matches the three synchronizations per iteration. The new timings are nearly flat. The lower rows are deliberately extreme to show the scaling; the realistic operating range for short-sample datasets is the 128 - 1024 document region, where the aggregation drops from 15 - 115 ms to well under 0.4 ms per micro-batch.

Correctness was verified against the original implementation on the same inputs:

• token counts match exactly;
• loss sums differ by at most ~0.2 in absolute terms, i.e. a relative error below ~5e-6, which is the expected FP32 discrepancy from a different reduction order.

Environment: NVIDIA H20, torch 2.8.0+cu128, Megatron-LM 0.18.0

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.