From 9e4bf98e46a020beeb28170d0cebfbb71414c9de Mon Sep 17 00:00:00 2001 From: gakkiri Date: Sat, 22 Aug 2026 11:05:08 +0800 Subject: [PATCH 1/2] perf(megatron): defer loss all-reduce to log time --- swift/megatron/trainers/trainer.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/swift/megatron/trainers/trainer.py b/swift/megatron/trainers/trainer.py index a5fa72bc68..088b35ebd4 100644 --- a/swift/megatron/trainers/trainer.py +++ b/swift/megatron/trainers/trainer.py @@ -9,7 +9,7 @@ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from typing import List, Optional -from swift.utils import get_logger +from swift.utils import get_current_device, get_logger from .base import BaseMegatronTrainer logger = get_logger() @@ -62,9 +62,9 @@ def loss_func(self, losses = losses * loss_scale loss = torch.cat([torch.sum(losses * loss_mask).view(1), loss_mask.sum().view(1)]) - # Reduce loss for logging. + # Reduce loss for logging; the DP all-reduce is deferred to log time + # (once per logging event) to avoid a global sync point per microbatch. reporting_loss = loss.detach().clone() - torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group(with_context_parallel=True)) lm_loss = loss[0] lm_loss = lm_loss.clone() @@ -111,6 +111,24 @@ def _compute_channel_loss(self, losses, loss_mask, channels, packed_seq_params=N new_metrics = self._all_reduce_metric(new_metrics, torch.distributed.ReduceOp.SUM, group=dp_cp_group) return new_metrics + def _log_callback(self, logs, n_steps): + # loss_func defers the logging-loss DP all-reduce from per-microbatch + # to here (once per logging event); numerically identical by linearity. + # All last-stage ranks must enter the collective unconditionally: + # a rank whose whole logging window had zero valid tokens (e.g. a fully + # masked CP shard) contributes zeros instead of skipping the call, which + # would otherwise deadlock the group. + if self.args.task_type == 'causal_lm' and mpu.is_pipeline_last_stage(ignore_virtual=True): + v = logs.get('loss') + if v is None: + v = torch.zeros(2, dtype=torch.float32, device=get_current_device()) + dist.all_reduce(v, op=dist.ReduceOp.SUM, group=mpu.get_data_parallel_group(with_context_parallel=True)) + if v[1].item() > 0: + logs['loss'] = v + else: + logs.pop('loss', None) + super()._log_callback(logs, n_steps) + def forward_step(self, data_iterator, model): vp_stage = model.module.module.vp_stage data = self.get_batch(data_iterator, vp_stage) From 6bff4105e014d3dba3f1d0e7abef8ff4b96274fb Mon Sep 17 00:00:00 2001 From: gakkiri Date: Sat, 22 Aug 2026 11:26:01 +0800 Subject: [PATCH 2/2] docs(megatron): clarify deferred loss reduction semantics --- swift/megatron/trainers/trainer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/swift/megatron/trainers/trainer.py b/swift/megatron/trainers/trainer.py index 088b35ebd4..b0c5a9fbe1 100644 --- a/swift/megatron/trainers/trainer.py +++ b/swift/megatron/trainers/trainer.py @@ -62,7 +62,7 @@ def loss_func(self, losses = losses * loss_scale loss = torch.cat([torch.sum(losses * loss_mask).view(1), loss_mask.sum().view(1)]) - # Reduce loss for logging; the DP all-reduce is deferred to log time + # Keep local loss stats for logging; defer the DP all-reduce to log time # (once per logging event) to avoid a global sync point per microbatch. reporting_loss = loss.detach().clone() @@ -113,7 +113,8 @@ def _compute_channel_loss(self, losses, loss_mask, channels, packed_seq_params=N def _log_callback(self, logs, n_steps): # loss_func defers the logging-loss DP all-reduce from per-microbatch - # to here (once per logging event); numerically identical by linearity. + # to here (once per logging event); mathematically equivalent by + # linearity, modulo floating-point reduction order. # All last-stage ranks must enter the collective unconditionally: # a rank whose whole logging window had zero valid tokens (e.g. a fully # masked CP shard) contributes zeros instead of skipping the call, which