From 590c9e74ec71c07d10bde45c83d5bd8db1bec84f Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Wed, 15 Jul 2026 08:55:00 +0000 Subject: [PATCH] [FSDP][Loss] Switch gradient reduction to SUM Flip gradient reduction from mean to pure SUM, so each parameter receives the sum of per-rank local-component gradients, i.e. the global loss gradient, with no compensating divides. This is the single atomic semantic change; the three cancelling x world_size factors and the reduce divides are removed together. - BaseModel/MoE/Dense.fully_shard call set_gradient_reduce_sum() at the end. - MoE.scale_and_reduce_grad drops the expert div_(ep_size) and the replicated div_(flat_mesh.size()); only the coalesced SUM all_reduce remains. - CE: drop the WORLD autograd all_reduce; the loss stays this rank's local component (display global value restored by the C2 detached pipeline). - Balancing: use local_gating_sum directly instead of all_reduce_autograd; the global detached statistics (tokens_global/seqlen_global/scale_global) are kept unchanged. - Z-loss: drop the x world_size in the global-average branch. Verified on torch 2.10 (bf16 force-sum): distributed full gradient reproduces a single-process full-batch token-mean CE reference at EP=1 and EP=2 (norm-ratio median 0.9998, EP-invariant); global display loss unchanged through the flip; balancing+z backward stays finite. Regression tests cover the bf16 reduce-sum mechanism, token-mean parity with grad-acc=2, aux-loss finite gradients, and an isolated fp32 balancing-only gate-gradient A/B (new local+SUM vs old all_reduce_autograd+AVG) asserting element-wise equality. --- tests/model/test_reduce_sum_grad.py | 385 +++++++++++++++++- xtuner/v1/loss/ce_loss.py | 9 +- xtuner/v1/loss/moe_loss.py | 34 +- xtuner/v1/model/base.py | 19 + xtuner/v1/model/compose/base.py | 6 + .../compose/intern_s1/modeling_intern_s1.py | 4 + xtuner/v1/model/dense/dense.py | 4 + xtuner/v1/model/moe/moe.py | 14 +- 8 files changed, 443 insertions(+), 32 deletions(-) diff --git a/tests/model/test_reduce_sum_grad.py b/tests/model/test_reduce_sum_grad.py index 77c6d7da12..64ec7c2a8c 100644 --- a/tests/model/test_reduce_sum_grad.py +++ b/tests/model/test_reduce_sum_grad.py @@ -1,10 +1,20 @@ +from contextlib import contextmanager + import torch import torch.distributed as dist import torch.nn as nn -from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard +from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard +from torch.distributed.tensor import DTensor from xtuner._testing.testcase import DeterministicDDPTestCase -from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.loss.moe_loss import BalancingLossConfig, ZLossConfig +from xtuner.v1.model.base import BaseModel, ModelItem, XTunerBaseModelConfig +from xtuner.v1.model.moe.moe import MoE, MoEConfig, SequenceContext +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.router import NoAuxRouterConfig class _ReduceSumToyConfig(XTunerBaseModelConfig): @@ -74,3 +84,374 @@ def test_bf16_reduce_sum_equals_local_grad_sum(self): rel_to_mean = ((full_grad - g_mean).norm() / (g_mean.norm() + 1e-12)).item() assert rel_to_sum < 1e-2, f"expected SUM of local grads, rel_to_sum={rel_to_sum}" assert rel_to_mean > 0.1, f"gradient matched MEAN, reduce-sum not applied, rel_to_mean={rel_to_mean}" + + +@contextmanager +def _fake_dist_uninitialized(): + # Make the single-process reference skip its WORLD all_reduce so the CE denominator counts each + # token once, i.e. a plain token-mean CE over the whole (concatenated) global batch. + orig = dist.is_initialized + dist.is_initialized = lambda: False + try: + yield + finally: + dist.is_initialized = orig + + +def _tiny_moe_config(ep_size: int, balancing: bool, z_loss: bool) -> MoEConfig: + return MoEConfig( + vocab_size=1024, + max_position_embeddings=512, + pad_token_id=0, + eos_token_id=0, + num_hidden_layers=2, + hidden_size=256, + intermediate_size=512, + rms_norm_eps=1e-6, + rope_theta=1e6, + hidden_act="silu", + attention=MHAConfig(num_attention_heads=8, num_key_value_heads=8, head_dim=32), + tie_word_embeddings=False, + n_routed_experts=8, + n_shared_experts=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + hidden_factor=1.0, + moe_intermediate_size=256, + compile_cfg=False, + router=NoAuxRouterConfig( + scoring_func="sigmoid", router_scaling_factor=1.0, n_group=1, topk_group=1, norm_topk_prob=True + ), + ep_size=ep_size, + balancing_loss_cfg=BalancingLossConfig() if balancing else None, + z_loss_cfg=ZLossConfig() if z_loss else None, + ) + + +def _full_grads(model) -> dict[str, torch.Tensor]: + out = {} + for name, p in model.named_parameters(): + if p.grad is None: + continue + g = p.grad + g = g.full_tensor() if isinstance(g, DTensor) else g + out[name.replace("._checkpoint_wrapped_module", "")] = g.detach().float().cpu() + return out + + +class TestReduceSumEndToEnd(DeterministicDDPTestCase): + @property + def world_size(self) -> int: + return 2 + + def test_reduce_sum_moe_matches_token_mean_reference_ep1(self): + # EP=1 keeps every param replicated, so scale_and_reduce_grad's no-divide SUM branch is on + # the critical path. + self._check_token_mean_parity(ep_size=1) + + def test_reduce_sum_moe_matches_token_mean_reference_ep2(self): + # EP=2 exercises the expert path: removing the expert div_(ep_size) (the §6 high-risk change) + # must still reproduce the token-mean reference. The experts_fsdp reduce-scatter is the only + # aggregation their grads get, so a wrong factor here would show as an ep_size scaling. + self._check_token_mean_parity(ep_size=2) + + def _check_token_mean_parity(self, ep_size: int): + # The reduce-sum path (FSDP SUM reduce-scatter + no CE WORLD all_reduce + SUM-only + # scale_and_reduce_grad) must reproduce a single-process, full-batch token-mean CE gradient. + # grad-acc=2 exercises SUM composability across micro-batch backwards. + self.create_pg("cuda") + device = "cuda" + rank = dist.get_rank() + seq_len = 32 + n_microbatch = 2 + efsdp = self.world_size // ep_size + + config = _tiny_moe_config(ep_size=ep_size, balancing=False, z_loss=False) + torch.manual_seed(0) + fsdp_cfg = FSDPConfig(cpu_offload=False, ep_size=ep_size, reduce_dtype=torch.bfloat16) + engine = TrainEngine(model_cfg=config, optim_cfg=AdamWConfig(), fsdp_cfg=fsdp_cfg) + engine.init_model_weights() + + gold_weights = { + name.replace("._checkpoint_wrapped_module", ""): ( + p.full_tensor() if isinstance(p, DTensor) else p.detach() + ) + .detach() + .float() + .cpu() + for name, p in engine.model.named_parameters() + } + + # One distinct sequence per (fsdp shard, micro-batch); concatenated they form the reference + # global batch. EP replicas (same fsdp position) consume identical data. + gen = torch.Generator().manual_seed(1234) + shards = [torch.randint(0, 512, (1, seq_len + 1), generator=gen) for _ in range(efsdp * n_microbatch)] + + def build_items(ids_list): + loss_data = [] + for ids in ids_list: + ids = ids.to(device) + seq_ctx = SequenceContext.from_input_ids(input_ids=(ids[:, :-1],), device=device) + loss_data.append({"seq_ctx": seq_ctx, "shifted_labels": ids[:, 1:]}) + loss_ctx_list = engine.model.build_loss_ctx_batch(loss_data, sp_mesh=None) + return [ModelItem(seq_ctx=d["seq_ctx"], loss_ctx=lc) for d, lc in zip(loss_data, loss_ctx_list)] + + fsdp_idx = rank // ep_size + my_shards = shards[fsdp_idx * n_microbatch : (fsdp_idx + 1) * n_microbatch] + engine.model.zero_grad(set_to_none=True) + engine.train_step(build_items(my_shards)) + engine.model.scale_and_reduce_grad() + dist_grads = _full_grads(engine.model) + + if rank == 0: + ref = MoE(config=_tiny_moe_config(ep_size=1, balancing=False, z_loss=False)).to(torch.bfloat16).to(device) + missing, unexpected = ref.load_state_dict( + {k: v.to(torch.bfloat16).to(device) for k, v in gold_weights.items()}, strict=False + ) + assert not unexpected, f"unexpected keys: {unexpected}" + ref.zero_grad(set_to_none=True) + loss_cfg = CELossConfig() + seq_ctxs, loss_ctxs = [], [] + for ids in shards: + ids = ids.to(device) + seq_ctxs.append(SequenceContext.from_input_ids(input_ids=(ids[:, :-1],), device=device)) + loss_ctxs.append(loss_cfg.build(data={"shifted_labels": ids[:, 1:]}, sp_mesh=None)) + with _fake_dist_uninitialized(): + loss_ctxs = loss_cfg.loss_ctx_cls.build_batches(loss_ctxs) + total = torch.zeros((), device=device) + for seq_ctx, lc in zip(seq_ctxs, loss_ctxs): + total = total + ref(seq_ctx=seq_ctx, loss_ctx={"lm": lc})["loss"] + total.backward() + ref_grads = {n: p.grad.detach().float().cpu() for n, p in ref.named_parameters() if p.grad is not None} + + ratios = [] + for name, rg in ref_grads.items(): + dg = dist_grads.get(name) + if dg is None or dg.shape != rg.shape: + continue + ratios.append((dg.norm() / rg.norm().clamp_min(1e-12)).item()) + ratios_t = torch.tensor(ratios) + median = ratios_t.median().item() + assert abs(median - 1.0) < 0.05, f"ep={ep_size} reduce-sum grad norm ratio median={median} (want ~1.0)" + assert (ratios_t - 1.0).abs().max().item() < 0.3, f"ep={ep_size} a param grad ratio drifted: {ratios_t}" + + def test_reduce_sum_with_aux_losses_produces_finite_grads(self): + # Balancing + z-loss enabled: after dropping their `× world_size` injection, backward must + # still produce finite, non-zero gradients (the aux-loss backward flows through the router). + self.create_pg("cuda") + device = "cuda" + rank = dist.get_rank() + seq_len = 32 + + config = _tiny_moe_config(ep_size=1, balancing=True, z_loss=True) + torch.manual_seed(0) + fsdp_cfg = FSDPConfig(cpu_offload=False, ep_size=1, reduce_dtype=torch.bfloat16) + engine = TrainEngine(model_cfg=config, optim_cfg=AdamWConfig(), fsdp_cfg=fsdp_cfg) + engine.init_model_weights() + + gen = torch.Generator().manual_seed(1234 + rank) + loss_data = [] + for _ in range(2): + ids = torch.randint(0, 512, (1, seq_len + 1), generator=gen).to(device) + seq_ctx = SequenceContext.from_input_ids(input_ids=(ids[:, :-1],), device=device) + loss_data.append({"seq_ctx": seq_ctx, "shifted_labels": ids[:, 1:]}) + loss_ctx_list = engine.model.build_loss_ctx_batch(loss_data, sp_mesh=None) + items = [ModelItem(seq_ctx=d["seq_ctx"], loss_ctx=lc) for d, lc in zip(loss_data, loss_ctx_list)] + + engine.model.zero_grad(set_to_none=True) + info = engine.train_step(items) + engine.model.scale_and_reduce_grad() + + assert torch.isfinite(torch.tensor(info["total_loss"])), "total_loss is not finite" + grads = _full_grads(engine.model) + gate_grads = [g for n, g in grads.items() if "gate" in n] + assert gate_grads, "router gate has no gradient; aux-loss backward path is broken" + for name, g in grads.items(): + assert torch.isfinite(g).all(), f"non-finite gradient in {name}" + assert any(g.abs().sum() > 0 for g in gate_grads), "router gate gradients are all zero" + + +class TestBalancingLossReduceSum(DeterministicDDPTestCase): + @property + def world_size(self) -> int: + # >1 rank forms the replicate group over which the router/gate gradient is aggregated. + return 4 + + def test_balancing_local_sum_matches_global(self): + # Isolate the balancing-loss gate gradient (CE off) and prove the reduce-sum rewrite is + # exact for it: aggregate grad-norm parity cannot single out this term. The new scheme + # (local_gating_sum + global detached coefficients, gradients SUM-reduced across ranks) + # must reproduce the old scheme (all_reduce_autograd + FSDP mean-reduce) element-wise. In + # fp32 with a fixed router this is a clean equality (no bf16 routing jitter), so the + # tolerance is tight; if it fails, the "linear + global detached coefficient => local-sum + # equals global" reasoning is wrong and the reduce-sum switch is unsound. + from torch.distributed.nn.functional import all_reduce as all_reduce_autograd + + self.create_pg("cuda") + device = "cuda" + rank = dist.get_rank() + world = dist.get_world_size() + + n_layers, n_experts, topk, n_tokens, hidden = 2, 8, 2, 16, 32 + # Replicated gate: identical init on every rank. Distinct per-rank token features so each + # rank owns a different local_gating_sum, exactly the asymmetry the reduce path must handle. + w_gen = torch.Generator().manual_seed(123) + w0 = torch.randn(hidden, n_experts, generator=w_gen, dtype=torch.float32).to(device) + x_gen = torch.Generator().manual_seed(1000 + rank) + x = torch.randn(n_layers, n_tokens, hidden, generator=x_gen, dtype=torch.float32).to(device) + + def router_weights(w: torch.Tensor) -> torch.Tensor: + return torch.sigmoid(torch.einsum("lth,he->lte", x, w)) + + # Non-differentiable per-expert token counts; global view via a detached SUM all_reduce. + rw_detached = router_weights(w0) + _, selected = torch.topk(rw_detached, topk, dim=-1) + tpe_local = torch.stack( + [torch.histc(selected[layer].float(), bins=n_experts, min=0, max=n_experts) for layer in range(n_layers)] + ).long() + tpe_global = tpe_local.clone() + dist.all_reduce(tpe_global, op=dist.ReduceOp.SUM) + + cfg = BalancingLossConfig() + alpha = cfg.balancing_loss_alpha + tokens_global = tpe_global.sum(-1) + seqlen_global = tokens_global // topk + scale_global = n_experts / tokens_global + + # New scheme via the real BalancingLossContext: per-rank local loss, gradients SUM-reduced. + w_new = w0.clone().requires_grad_(True) + rw_new = router_weights(w_new) + ctx = cfg.build() + for layer in range(n_layers): + ctx.accumulate(router_weights=rw_new[layer]) + loss_new, _ = ctx.finalize( + tokens_per_expert_local=tpe_local, + tokens_per_expert_global=tpe_global, + n_routed_experts=n_experts, + num_experts_per_tok=topk, + non_pad_token=n_tokens, + ) + loss_new.backward() + g_new = w_new.grad.clone() + dist.all_reduce(g_new, op=dist.ReduceOp.SUM) # FSDP/scale_and_reduce SUM over the replicate group + + # Old scheme: all_reduce_autograd over the gating sum, then FSDP mean-reduce (divide by world). + w_old = w0.clone().requires_grad_(True) + rw_old = router_weights(w_old) + gating_sum = torch.stack([rw_old[layer].sum(dim=0) for layer in range(n_layers)]) + routing_weights_sum_global = all_reduce_autograd(gating_sum, op=dist.ReduceOp.SUM) + routing_weights_mean = routing_weights_sum_global / seqlen_global.unsqueeze(-1) + loss_old = (scale_global * (tpe_global * routing_weights_mean).sum(-1)).sum() * alpha + loss_old.backward() + g_old = w_old.grad.clone() + g_old.div_(world) # FSDP mean-reduce (old default AVG) + dist.all_reduce(g_old, op=dist.ReduceOp.SUM) + + if rank == 0: + assert g_new.abs().sum() > 0, "balancing gate gradient is all zero" + torch.testing.assert_close(g_new, g_old, rtol=1e-4, atol=1e-4) + + def test_zloss_local_matches_global(self): + # Same isolation for z-loss (CE off). The new scheme drops the `× world_size` factor and + # relies on the SUM gradient reduction across the replicate group; it must reproduce the old + # scheme (`× world_size` in the loss + FSDP mean-reduce) on the router gate, element-wise. + self.create_pg("cuda") + device = "cuda" + rank = dist.get_rank() + world = dist.get_world_size() + + n_layers, n_experts, n_tokens, hidden = 2, 8, 16, 32 + w_gen = torch.Generator().manual_seed(321) + w0 = torch.randn(hidden, n_experts, generator=w_gen, dtype=torch.float32).to(device) + x_gen = torch.Generator().manual_seed(2000 + rank) + x = torch.randn(n_layers, n_tokens, hidden, generator=x_gen, dtype=torch.float32).to(device) + + def router_logits(w: torch.Tensor) -> torch.Tensor: + return torch.einsum("lth,he->lte", x, w) + + num_tokens_local = n_tokens + num_tokens_global = torch.tensor(num_tokens_local, device=device, dtype=torch.int64) + dist.all_reduce(num_tokens_global, op=dist.ReduceOp.SUM) # detached global token count + denom_global = torch.clamp(num_tokens_global, min=1) + + cfg = ZLossConfig() + alpha = cfg.z_loss_alpha + + # New scheme via the real ZLossContext: per-layer local z-loss (no × world_size), SUM-reduced. + w_new = w0.clone().requires_grad_(True) + logits_new = router_logits(w_new) + ctx = cfg.build() + z_new = torch.zeros((), device=device) + for layer in range(n_layers): + z_new = z_new + ctx.accumulate( + router_logits=logits_new[layer], + num_tokens_local=num_tokens_local, + num_tokens_global=num_tokens_global, + ) + z_new.backward() + g_new = w_new.grad.clone() + dist.all_reduce(g_new, op=dist.ReduceOp.SUM) + + # Old scheme: the same per-layer z-loss but multiplied by world_size, then FSDP mean-reduce. + w_old = w0.clone().requires_grad_(True) + logits_old = router_logits(w_old) + z_old = torch.zeros((), device=device) + for layer in range(n_layers): + base = torch.logsumexp(logits_old[layer], dim=-1).square().sum() / max(num_tokens_local, 1) + z_old = z_old + base * num_tokens_local * world / denom_global * alpha + z_old.backward() + g_old = w_old.grad.clone() + g_old.div_(world) # FSDP mean-reduce (old default AVG) + dist.all_reduce(g_old, op=dist.ReduceOp.SUM) + + if rank == 0: + assert g_new.abs().sum() > 0, "z-loss gate gradient is all zero" + torch.testing.assert_close(g_new, g_old, rtol=1e-4, atol=1e-4) + + +def _fsdp_reduce_cfg(module: FSDPModule): + # Read back the reduce-scatter reduction config that set_gradient_divide_factor / + # set_force_sum_reduction_for_comms write onto the FSDP param group. + param_group = module._get_fsdp_state()._fsdp_param_group # type: ignore[attr-defined] + return param_group.gradient_divide_factor, param_group.force_sum_reduction_for_comms + + +class TestComposeReduceSumHook(DeterministicDDPTestCase): + @property + def world_size(self) -> int: + return 2 + + def test_root_pass_sets_sum_on_independently_sharded_children(self): + # Compose/VLM models shard vision_tower / multi_modal_projector via their own fully_shard + # overrides that do NOT set reduce-sum, plus a root self._fully_shard wrap. The fix relies on + # a single root-level set_gradient_reduce_sum() covering every nested FSDPModule via + # self.modules(). This reproduces that topology with independently sharded children and + # asserts they all flip from the FSDP AVG default to divide_factor=1 + force_sum. + self.create_pg("cuda") + model = _ReduceSumToyConfig(hidden_size=8, compile_cfg=False).build().cuda() + # Stand-ins for vision_tower / multi_modal_projector, sharded independently without reduce-sum. + vision_like = nn.Linear(8, 8, bias=False).cuda() + projector_like = nn.Linear(8, 8, bias=False).cuda() + model.add_module("vision_like", vision_like) + model.add_module("projector_like", projector_like) + + mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16) + fully_shard(vision_like, mp_policy=mp_policy) + fully_shard(projector_like, mp_policy=mp_policy) + fully_shard(model, mp_policy=mp_policy) # root wrap, FSDP AVG default + + fsdp_modules = [m for m in model.modules() if isinstance(m, FSDPModule)] + assert len(fsdp_modules) >= 3, "expected root + two independently sharded children" + # Precondition: none are reduce-sum yet (default AVG has divide_factor None / force_sum False). + for m in fsdp_modules: + factor, force_sum = _fsdp_reduce_cfg(m) + assert not (factor == 1.0 and force_sum), "precondition failed: module already reduce-sum" + + # The compose fix: one root-level pass. + model.set_gradient_reduce_sum() + + for m in fsdp_modules: + factor, force_sum = _fsdp_reduce_cfg(m) + assert factor == 1.0 and force_sum is True, ( + f"nested FSDPModule not switched to SUM: divide_factor={factor} force_sum={force_sum}" + ) diff --git a/xtuner/v1/loss/ce_loss.py b/xtuner/v1/loss/ce_loss.py index eba945fae7..fb2fa6045b 100644 --- a/xtuner/v1/loss/ce_loss.py +++ b/xtuner/v1/loss/ce_loss.py @@ -6,7 +6,6 @@ import torch.nn.functional as F from cyclopts import Parameter from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.nn.functional import all_reduce from xtuner.v1.utils.device import get_device @@ -282,10 +281,10 @@ def forward( extra_info["local_base_loss"] = loss.detach().clone() - # Step 2.c in the loss calculation: reduce the loss over all ranks using all_reduce with autograd support - if dist.is_initialized(): - loss = all_reduce(loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD) - + # Under reduce-sum gradients the loss stays as this rank's local component (local token sum + # over the global token denominator). Cross-rank aggregation happens on the gradients via the + # FSDP SUM reduce-scatter, so no autograd WORLD all_reduce is injected here. The global loss + # scalar for logging is restored separately by the detached display pipeline (§5.4). return loss, (logits, extra_info) diff --git a/xtuner/v1/loss/moe_loss.py b/xtuner/v1/loss/moe_loss.py index 62b8c54bd2..3ee6f38f81 100644 --- a/xtuner/v1/loss/moe_loss.py +++ b/xtuner/v1/loss/moe_loss.py @@ -155,22 +155,18 @@ def finalize( alpha = self.loss_cfg.balancing_loss_alpha if self.loss_cfg.balancing_loss_global_average and dist.is_initialized(): - group = dist.group.WORLD - assert group is not None tokens_global = tokens_per_expert_global.sum(-1) seqlen_global = tokens_global // num_experts_per_tok scale_global = n_routed_experts / tokens_global - routing_weights_sum_global = all_reduce_autograd(local_gating_sum, "sum", group) - routing_weights_mean_global = routing_weights_sum_global / seqlen_global.unsqueeze(-1) - loss_vec = scale_global * (tokens_per_expert_global * routing_weights_mean_global).sum(-1) - - # Detached local component: same global denominators, but this rank's own gating sum in - # place of the cross-rank sum. Because `all_reduce_autograd` is a plain SUM and every - # other factor here (scale_global, seqlen_global, tokens_per_expert_global) is detached - # and global, summing `local_vec` over ranks reproduces `loss_vec` exactly. - routing_weights_mean_local = local_gating_sum.detach() / seqlen_global.unsqueeze(-1) - local_vec = scale_global * (tokens_per_expert_global * routing_weights_mean_local).sum(-1) + # Under reduce-sum gradients the loss stays as this rank's local component: use this + # rank's own gating sum with the global denominators, without any cross-rank all_reduce. + # Cross-rank aggregation happens on the gradients (FSDP / scale_and_reduce_grad SUM); + # since every other factor here is detached and global, summing the loss over ranks + # reproduces the global balancing loss. + routing_weights_mean = local_gating_sum / seqlen_global.unsqueeze(-1) + loss_vec = scale_global * (tokens_per_expert_global * routing_weights_mean).sum(-1) + local_vec = loss_vec.detach() else: valid_tokens = max(non_pad_token, 1) scale_global = n_routed_experts / (valid_tokens * num_experts_per_tok) @@ -298,20 +294,18 @@ def accumulate( denom_local = max(num_tokens_local, 1) base = torch.logsumexp(router_logits, dim=-1).square().sum() / denom_local - local_loss = base loss = base if self.loss_cfg.z_loss_global_average and num_tokens_global is not None: - # Local component: this rank's share of the global-average z-loss, without the - # `× world_size` factor. The backward path keeps `× world_size` (removed in the - # reduce-sum switch); summing `local_loss` over ranks reproduces the global z-loss. + # Under reduce-sum gradients the injected z-loss stays as this rank's local component + # (its share of the global-average z-loss, WITHOUT any `× world_size`). Cross-rank + # aggregation happens on the gradients via the FSDP SUM reduce-scatter; summing this + # local component over ranks reproduces the global z-loss. denom_global = torch.clamp(num_tokens_global, min=1) - local_loss = base * num_tokens_local / denom_global - loss = local_loss * world_size + loss = base * num_tokens_local / denom_global - local_loss = local_loss * self.loss_cfg.z_loss_alpha / self._batch_size loss = loss * self.loss_cfg.z_loss_alpha / self._batch_size self._update_running(loss.detach()) - self._update_local_running(local_loss.detach()) + self._update_local_running(loss.detach()) return loss def finalize(self) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 7bd171d26d..a1f5e10125 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -676,6 +676,10 @@ def fully_shard( reshard_after_forward=fsdp_config.reshard_after_forward, offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, ) + # Reduce-scatter gradients with pure SUM (no divide). Combined with the loss forwards no + # longer injecting x world_size, each param's gradient is the sum of per-rank local-component + # gradients, i.e. the global loss gradient. Covers nested/child FSDP modules via self.modules(). + self.set_gradient_reduce_sum() return self def _fully_shard( @@ -1388,6 +1392,21 @@ def post_micro_batch_forward(self, batch_outputs: Sequence[ModelOutputs]) -> Bat if "reduced_base_loss" in reduced_other_losses: reduced_other_losses["reduced_llm_loss"] = reduced_other_losses.pop("reduced_base_loss") + # Safety net: every `*loss` tensor field an output exposes must have produced a display curve + # from a registered local component. Otherwise a newly added loss term would silently vanish + # from the logged curves while still driving backward, hiding the regression. + for output in batch_outputs: + for name in output.model_fields: + field_value = getattr(output, name, None) + if "loss" in name and isinstance(field_value, torch.Tensor): + expected = "reduced_llm_loss" if name == "loss" else f"reduced_{name}" + if expected not in reduced_other_losses: + raise RuntimeError( + f"Loss field '{name}' has no display curve ('{expected}' missing from reduced " + f"logs): register its detached local component on extra_info (see " + f"`_store_local_display_losses` / CE's `local_base_loss`)." + ) + ret = BatchForwardInfo( logs_info=reduced_other_losses, extra_info=train_engine_extra_info, diff --git a/xtuner/v1/model/compose/base.py b/xtuner/v1/model/compose/base.py index 51eb1fa02b..4823764cb1 100644 --- a/xtuner/v1/model/compose/base.py +++ b/xtuner/v1/model/compose/base.py @@ -138,6 +138,12 @@ def fully_shard( self.language_model.set_modules_to_forward_prefetch([self.language_model.layers["0"]]) # type: ignore self._to_empty_meta() + # Reduce-scatter gradients with pure SUM for every sharded submodule. The vision tower, + # projector, and this compose root are sharded by their own fully_shard overrides / the root + # wrap above, none of which set reduce-sum; a single root-level pass over self.modules() + # covers them all (and is idempotent for the language model, already set). Without this the + # vision/projector grads silently fall back to FSDP AVG and lose a 1/fsdp_size factor. + self.set_gradient_reduce_sum() return self def from_hf(self, hf_path: str | Path, strict=True): diff --git a/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py b/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py index 7a7c4387c6..619d6d59b4 100644 --- a/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py +++ b/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py @@ -98,6 +98,10 @@ def fully_shard( self.language_model.set_modules_to_forward_prefetch([self.language_model.layers["0"]]) # type: ignore self._to_empty_meta() + # Reduce-scatter gradients with pure SUM for every sharded submodule (vision tower, projector, + # compose root); their own fully_shard overrides do not set reduce-sum. One root-level pass + # over self.modules() covers them all (idempotent for the already-set language model). + self.set_gradient_reduce_sum() return self def extract_feature(self, pixel_values): diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index ef1ad4c7e9..4a90e795b4 100644 --- a/xtuner/v1/model/dense/dense.py +++ b/xtuner/v1/model/dense/dense.py @@ -300,6 +300,10 @@ def fully_shard( # Make sure it works properly when using fsdp if self.config.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight + # Reduce-scatter gradients with pure SUM (no divide) for every sharded submodule; combined + # with the loss forwards no longer injecting x world_size, this yields the global loss + # gradient. See BaseModel.set_gradient_reduce_sum. + self.set_gradient_reduce_sum() return self # TODO: 支持 tp diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index fccc75c646..c9f3370847 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -1193,6 +1193,10 @@ def fully_shard( module.forward = types.MethodType(self.patched_emb_forward, module) # type: ignore self._to_empty_meta() + # Reduce-scatter gradients with pure SUM (no divide) for every sharded submodule; the + # expert / replicated grads not covered by reduce-scatter are handled without division in + # scale_and_reduce_grad. See BaseModel.set_gradient_reduce_sum. + self.set_gradient_reduce_sum() return self @property @@ -1222,10 +1226,10 @@ def scale_and_reduce_grad(self): if param.grad is None: continue - # Expert parameters live on a unique EP rank, so no cross-rank reduction - # is needed — just rescale by `ep_size` to keep the effective average. + # Expert parameters live on a unique EP rank; their FSDP sharding is only over the + # experts_fsdp sub-dim, already SUM-reduced by reduce-scatter. No cross-rank reduction + # and no rescaling: under reduce-sum the local-component gradient is what we keep. if ep_enabled and ".experts" in name: - param.grad.div_(self.ep_mesh.size()) # type: ignore continue if not isinstance(param, DTensor): @@ -1252,8 +1256,8 @@ def scale_and_reduce_grad(self): flat_mesh = param.device_mesh[replicate_dim_names[0]] grad = param.grad.to_local() if isinstance(param.grad, DTensor) else param.grad - # Pre-scale locally so the SUM all_reduce below yields the mean across replicas. - grad.div_(flat_mesh.size()) # type: ignore + # Replicated params get no reduce-scatter; SUM their per-rank local-component grads + # across the replicate group with NO pre-divide, matching the reduce-sum invariant. grads_by_group.setdefault(flat_mesh.get_group(), []).append(grad) # type: ignore # One coalesced all_reduce per process group covers all replicated grads.