[Bugfix] Fix loss missing from logs when context parallelism is enabled

[Zhikaiiii]

PR type

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

Summary

• Fix loss key missing from training logs when context parallelism (CP) is enabled
• The root cause is that reporting_loss was all-reduced only over get_data_parallel_group() (excludes CP ranks), so CP ranks whose sequence chunk has no valid tokens (all labels=-100) would report loss_mask.sum()=0, causing the loss key to be skipped in _aggregated_metrics
• Also fix the same issue in _compute_channel_loss which used the DP-only group for key synchronization and metric reduction

Root Cause

This issue is caused by how Swift splits sequence data across CP ranks. Swift uses a zigzag pattern to distribute sequence chunks (e.g., with CP=2: GPU0 gets chunk_0+chunk_3, GPU1 gets chunk_1+chunk_2). In SFT scenarios where prompts are long and responses are short, the loss-bearing tokens (response part) tend to concentrate at one end of the sequence. As a result, certain CP ranks may receive chunks where all labels are -100 (prompt-only), making loss_mask.sum()=0 on those ranks. Since reporting_loss was only all-reduced within the DP group (not across CP ranks), the aggregated token count remains zero on these ranks, and the loss key gets skipped entirely during metric aggregation.

Changes

• trainer.py:69: Use get_data_parallel_group(with_context_parallel=True) for reporting_loss all-reduce
• trainer.py:98-104: Use DP+CP group for channel loss key synchronization and metric reduction

[gemini-code-assist]

Code Review

This pull request updates the Megatron trainer to include context parallel groups in data parallel operations for loss reduction and metric synchronization. Specifically, it modifies loss_func and _compute_channel_loss to use with_context_parallel=True when retrieving the data parallel group. A potential RuntimeError was identified in _compute_channel_loss where _all_reduce_metric could be called with an empty dictionary, and a code suggestion was provided to add a safety check.

[gemini-code-assist]

The _all_reduce_metric method in the base class uses torch.stack on the dictionary values. If new_metrics is empty (which can happen if no ranks in the group have any channel-specific tokens in the current micro-batch), torch.stack will raise a RuntimeError. It is safer to check if the dictionary is non-empty before attempting the reduction.

Suggested change

	new_metrics = self._all_reduce_metric(new_metrics, torch.distributed.ReduceOp.SUM, group=dp_cp_group)
	if new_metrics:
	new_metrics = self._all_reduce_metric(new_metrics, torch.distributed.ReduceOp.SUM, group=dp_cp_group)

[Jintao-Huang]

thanks!