perf(megatron): defer loss all-reduce to log time

[gakkiri]

Defer logging-loss DP×CP all-reduce to log time

PR type

• Performance Optimization

PR information

Summary

MegatronTrainer.loss_func currently performs one DP×CP all-reduce per microbatch for the detached logging loss. This PR keeps the local [loss_sum, token_count] pair during forward/backward aggregation and performs the reduction once per logging event.

This only changes the logging path. Training math is unchanged, and the reported loss remains mathematically equivalent, modulo floating-point reduction-order differences.

Motivation

The current causal-LM path runs this synchronous collective once per microbatch:

reporting_loss = loss.detach().clone()
torch.distributed.all_reduce(
    reporting_loss,
    group=mpu.get_data_parallel_group(with_context_parallel=True),
)

The reduced tensor feeds only metrics['loss']:

• backward uses the local lm_loss, returned separately;
• the detached reporting tensor is not consumed by the optimizer, scheduler, or training control flow;
• _aggregated_metrics subsequently sums the [loss_sum, token_count] pairs over microbatches and logging steps.

Although the payload is only two float32 values, each call creates a synchronization point across the participating DP×CP group. On a large, latency-sensitive group, per-microbatch rank jitter can therefore be exposed repeatedly on the training critical path.

The change

Only swift/megatron/trainers/trainer.py is changed:

1. loss_func returns its local detached [loss_sum, token_count] without immediately reducing it.
2. MegatronTrainer._log_callback reduces the aggregated pair over get_data_parallel_group(with_context_parallel=True) before the base callback converts it to sum / count.
3. Every physical pipeline-last-stage rank enters the collective. If a rank has no local loss key because its entire logging window contained zero valid tokens, it contributes [0, 0]. If the group-global count is also zero, the key is dropped, matching the previous behavior.

The override is limited to MegatronTrainer with task_type == 'causal_lm'. BaseMegatronTrainer is not changed because it is also shared by DPO, GRPO, Reward, Embedding, and other trainers that already reduce their own metrics. The seq-cls path is also unchanged.

With enable_channel_loss=true, per-channel metrics still perform their existing per-microbatch all_gather_object and all-reduce. Deferring that dynamic-key path is outside the scope of this PR.

Why the result is equivalent

For a logging window containing K optimizer steps and M microbatches per step, let x[t,m] be one rank's local [loss_sum, token_count] pair. The old and new paths are mathematically equivalent because all-reduce(SUM) is linear:

old: Σ_t Σ_m all_reduce(x[t,m])
new: all_reduce(Σ_t Σ_m x[t,m])

The number of logging-loss collectives per logging window changes from K × M to 1, saving K × M - 1 calls. The final quotient is the same global token-weighted mean. Floating-point addition occurs in a different order, so last-bit differences are possible and expected.

Collective participation remains consistent:

• training logging cadence is derived from the shared iteration counter;
• evaluation reaches the same callback on every rank;
• the DP×CP group does not span pipeline stages, so only physical pipeline-last-stage ranks participate;
• participation does not depend on whether the local loss key exists.

Experiment results

Environment

• Code: ms-swift v4.4.1 (98a09c18c; the patch applies cleanly to the tag), Megatron Core 0.18.0, torch 2.8.0+cu128, NCCL 2.27.3, DeepEP 1.2.1+9af0e0d, TransformerEngine 2.16.1, Python 3.12.13, CUDA 12.9, driver 535.247.01.
• Hardware: 64 nodes × 8 NVIDIA H20 96GB (512 GPUs), intra-node NVSwitch, 8 × 200 Gb/s RoCE per node; AMD EPYC 9K84.
• Workload: 35B-A3B MoE, TP4 / PP1 / EP8 / DP128, GBS 2048, MBS 2, sequence length 16384, enable_channel_loss=false.
• Microbatches: 2048 / (2 × 128) = 8 per optimizer step. Each step in the two-step profiler window contained a logging event, so this path changes from 8 Python all-reduce invocations per step to 1.

Profiler observations

torch.profiler traces were captured for ranks 0–7 over two-step windows. The table reports trace-level stall observations associated with the reporting-loss all-reduce; profiler kernel-event counts are intentionally omitted because a kernel-event count is not necessarily one-to-one with Python collective invocations.

Metric	Before	After
Aggregate duration of long-stall (>100 ms) matching all-reduces in captured traces	13.3 s / 2 steps	None observed
Worst matching all-reduce	2.89 s	None observed
Largest GPU idle gap in the captured window	1058 ms	5.9 ms

Step time

Controlled A/B runs used the same steps, data order, hardware, and enable_channel_loss=false. Over steps 6–14 (n=9):

	Mean step time
Before	31.03 s/step
After	24.96 s/step (−19.6%)

The observed gain is specific to this latency-sensitive cluster and workload. Lower-latency fabrics, smaller DP×CP groups, or lower inter-rank variance should show a smaller absolute improvement.

Correctness validation

• Standalone functional test: a two-rank gloo harness exercised the real MegatronTrainer._log_callback, including the super() chain.
  • The deferred path matched the old per-microbatch reduction within rtol=1e-6.
  • A rank with zero valid tokens and no local loss key still entered the collective without hanging, and both ranks logged the correct global mean.
  • A group-global zero token count removed the key on both ranks, matching the original behavior.
• Training smoke comparison: three fresh-start runs used identical data and configuration. The two unpatched runs produced step-2 losses of 2.57292247 and 2.57282043; the patched run produced 2.57309008. All differences are O(1e-4) absolute, consistent in scale with expected nondeterministic variation. Step-1 loss was bit-identical across all three (2.33354068), and no anomalous grad-norm or downstream-training divergence was observed. The deterministic two-rank test above is the direct equivalence check.
• Patch hygiene: the patch applies cleanly with git apply --check on a fresh v4.4.1 checkout and contains only this change.
• Scale exercise: an earlier equivalent causal-LM implementation placed in the base callback ran for more than 800 steps at 512 ranks / DP128 without collective hangs or logging anomalies. The current scoped override was separately exercised by the functional and A/B tests above.

Expected benefit

The benefit should be largest with:

• large DP×CP groups spanning multiple nodes;
• many microbatches and/or multiple steps per logging window;
• high per-microbatch rank variance, such as MoE expert imbalance, packed variable-length sequences, host jitter, or a shared fabric.

The benefit should be smaller when the DP×CP group is small, the fabric is low-latency, or both num_microbatches and logging_steps are 1. No regression was observed in the tested configurations; the new path replaces many per-microbatch synchronization points with one two-element collective per logging event.

Summary

This PR moves a detached, logging-only DP×CP all-reduce out of the per-microbatch critical path. It reduces the logging-loss collective count from num_microbatches × logging_steps per logging window to one, preserves the global token-weighted loss up to floating-point reduction-order effects, handles zero-token CP shards without collective mismatch, and reduced mean step time by 19.6% in the measured DP128 workload.