diff --git a/tests/model/test_reduce_sum_grad.py b/tests/model/test_reduce_sum_grad.py new file mode 100644 index 0000000000..6ce0f121fe --- /dev/null +++ b/tests/model/test_reduce_sum_grad.py @@ -0,0 +1,766 @@ +from contextlib import contextmanager + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard +from torch.distributed.tensor import DTensor, Replicate, Shard + +from xtuner._testing.testcase import DeterministicDDPTestCase +from xtuner.v1.utils.misc import monkey_patch_hf_modules_cache +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, HFSaveCfg, ModelItem, XTunerBaseModelConfig +from xtuner.v1.model.dense.qwen3 import Qwen3DenseConfig +from xtuner.v1.model.moe.moe import MoE, MoEConfig, SequenceContext +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.mtp import MTPConfig +from xtuner.v1.module.router import NoAuxRouterConfig + + +class _ReduceSumDDPTestCase(DeterministicDDPTestCase): + """DDP test base for reduce-sum tests. + + These are gradient-PARITY tests that assert with tolerances, not bitwise reproducibility. + ``enable_full_determinism`` (``torch.use_deterministic_algorithms``) NaNs the MoE EP backward on + this base, so it is skipped here; the project's own EP parity scratchpad likewise runs without + it. Everything else (hf module cache patch, prepare) is preserved. + """ + + def run_func(self, test_name): + monkey_patch_hf_modules_cache() + self.prepare() + return getattr(self, test_name)() + + +class _ReduceSumToyConfig(XTunerBaseModelConfig): + hidden_size: int = 8 + + def build(self) -> "_ReduceSumToyModel": + return _ReduceSumToyModel(self) + + +class _ReduceSumToyModel(BaseModel): + config: _ReduceSumToyConfig + + def __init__(self, config: _ReduceSumToyConfig): + super().__init__(config) + self.fc = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self._init_load_spec() + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc(x) + + +class TestReduceSumGradient(DeterministicDDPTestCase): + @property + def world_size(self) -> int: + return 2 + + def test_bf16_reduce_sum_equals_local_grad_sum(self): + # Regression guard for the torch 2.10 bf16 FSDP reduce-sum path: `set_gradient_reduce_sum` + # must make the reduce-scatter yield the exact SUM of per-rank local gradients, not the + # AVG default and not the all-zero result of the NCCL PreMulSum bf16 bug (the failure that + # appears when only the divide factor is set). Assertions run in bf16 on purpose; fp32 + # reduction would mask the PreMulSum zeroing. + self.create_pg("cuda") + + torch.manual_seed(0) + dim = 8 + model = _ReduceSumToyConfig(hidden_size=dim, compile_cfg=False).build().cuda() + # Gradient of `(x @ w.T).sum()` w.r.t. `w` is independent of the weight value, so the bf16 + # weight cast under FSDP does not perturb the reference; only the reduction dtype matters. + ref_weight = model.fc.weight.detach().clone() + + mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16) + fully_shard(model, mp_policy=mp_policy) + model.set_gradient_reduce_sum() + + rank = dist.get_rank() + world = dist.get_world_size() + # Distinct input per rank so local gradients differ; SUM and MEAN are then clearly separable. + x = (torch.arange(dim, dtype=torch.float32, device="cuda") + rank + 1).reshape(1, dim) + model(x).sum().backward() + full_grad = model.fc.weight.grad.full_tensor().float() + + # Independent reference: this rank's local gradient on an unsharded copy, all-gathered. + w = ref_weight.clone().detach().requires_grad_(True) + (x @ w.T).sum().backward() + g_local = w.grad.float() + gathered = [torch.zeros_like(g_local) for _ in range(world)] + dist.all_gather(gathered, g_local) + g_sum = sum(gathered) + g_mean = g_sum / world + + assert not torch.all(full_grad.abs() < 1e-9), "bf16 reduce-sum returned all-zero gradients" + rel_to_sum = ((full_grad - g_sum).norm() / (g_sum.norm() + 1e-12)).item() + 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, + mtp: bool = False, + dispatcher: str | None = None, + router_bias_update_speed: float = 0.001, +) -> MoEConfig: + # hidden_size / moe_intermediate_size must be >= 512: the EP grouped-gemm kernel JIT-compiled by + # the all2all/DeepEP dispatcher fails ("PassManager::run failed") on smaller shapes. ep>1 uses the + # all2all dispatcher (naive is ep=1 only); ep=1 keeps the default (naive) dispatcher. + resolved_dispatcher = dispatcher if dispatcher is not None else ("all2all" if ep_size > 1 else None) + return MoEConfig( + vocab_size=1024, + max_position_embeddings=512, + pad_token_id=0, + eos_token_id=0, + num_hidden_layers=2, + hidden_size=512, + intermediate_size=1024, + rms_norm_eps=1e-6, + rope_theta=1e6, + hidden_act="silu", + attention=MHAConfig(num_attention_heads=16, num_key_value_heads=16, 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=512, + compile_cfg=False, + router=NoAuxRouterConfig( + scoring_func="sigmoid", + router_scaling_factor=1.0, + n_group=1, + topk_group=1, + norm_topk_prob=True, + router_bias_update_speed=router_bias_update_speed, + ), + ep_size=ep_size, + dispatcher=resolved_dispatcher, + balancing_loss_cfg=BalancingLossConfig() if balancing else None, + z_loss_cfg=ZLossConfig() if z_loss else None, + mtp_config=MTPConfig(num_layers=1, loss_scaling_factor=1.0) if mtp 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(_ReduceSumDDPTestCase): + @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" + + def _run_display_step(self, per_rank_seed: bool): + # Build a tiny MoE (balancing + z on), run one train_step, and return this rank's displayed + # losses. `per_rank_seed` chooses distinct (True) or identical (False) data across ranks. + device = "cuda" + rank = dist.get_rank() + config = _tiny_moe_config(ep_size=1, balancing=True, z_loss=True) + torch.manual_seed(0) + engine = TrainEngine( + model_cfg=config, optim_cfg=AdamWConfig(), fsdp_cfg=FSDPConfig(cpu_offload=False, reduce_dtype=torch.bfloat16) + ) + engine.init_model_weights() + gen = torch.Generator().manual_seed(1234 + (rank if per_rank_seed else 0)) + ids = torch.randint(0, 512, (1, 49), generator=gen).to(device) + seq_ctx = SequenceContext.from_input_ids(input_ids=(ids[:, :-1],), device=device) + lcs = engine.model.build_loss_ctx_batch([{"seq_ctx": seq_ctx, "shifted_labels": ids[:, 1:]}], sp_mesh=None) + info = engine.train_step([ModelItem(seq_ctx=seq_ctx, loss_ctx=lcs[0])]) + return info + + def test_display_loss_is_per_rank_no_all_reduce(self): + # Each rank must display ITS OWN loss (calibrate() -> per-token / per-rank mean, NO cross-rank + # all_reduce). Distinct data => the displayed values differ across ranks; identical data => + # every rank shows the same value (which, being the whole batch on each rank, is the global + # per-token mean, i.e. the world-size-1 value). + self.create_pg("cuda") + rank = dist.get_rank() + world = dist.get_world_size() + + info = self._run_display_step(per_rank_seed=True) + vals = [torch.zeros((), device="cuda") for _ in range(world)] + dist.all_gather(vals, torch.tensor(info["logs_info"]["reduced_llm_loss"], device="cuda")) + if rank == 0: + assert torch.isfinite(torch.stack(vals)).all(), "non-finite per-rank display loss" + assert (vals[0] - vals[1]).abs() > 1e-3, "distinct-data per-rank display losses are identical (all_reduced?)" + + info_sym = self._run_display_step(per_rank_seed=False) + vals_sym = [torch.zeros((), device="cuda") for _ in range(world)] + dist.all_gather(vals_sym, torch.tensor(info_sym["logs_info"]["reduced_llm_loss"], device="cuda")) + if rank == 0: + # Identical data on every rank => each rank's per-token mean equals the global one (up to + # bf16 cross-rank forward noise, since the base runs non-deterministically). + rel = ((vals_sym[0] - vals_sym[1]).abs() / vals_sym[0].abs().clamp_min(1e-6)).item() + assert rel < 0.02, f"identical-data per-rank display losses differ by {rel:.4f} (>2%)" + assert 0.0 < vals_sym[0].item() < 20.0, f"unreasonable per-token-mean CE {vals_sym[0].item()}" + + def test_reduce_sum_ep2_balancing_integration(self): + # EP=2 + balancing integration. The router gate sits on a Replicate placement over the EP + # sub-mesh; balancing loss flows ONLY into the gate (tokens_per_expert is detached, so every + # other param receives only CE gradient). Verify (a) all non-gate params still match the + # token-mean CE reference tightly -- i.e. EP=2 + balancing does not corrupt the main gradient + # flow -- and (b) the router gate gets a finite, non-zero gradient. The exact correctness of + # the gate's replicate-group SUM reduction with balancing is proven element-wise by + # TestBalancingLossReduceSum (that reduce group is the same kind as this EP sub-mesh); a tight + # dist-vs-single-process gate ratio is unreliable here because the ep-replicated + # tokens_per_expert_global statistics differ from an ep=1 reference by construction. + self.create_pg("cuda") + device = "cuda" + rank = dist.get_rank() + ep_size = 2 + n_microbatch = 2 + seq_len = 32 + efsdp = self.world_size // ep_size + + config = _tiny_moe_config(ep_size=ep_size, balancing=True, 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() + } + + 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_loss_data(ids_list): + out = [] + for ids in ids_list: + ids = ids.to(device) + seq_ctx = SequenceContext.from_input_ids(input_ids=(ids[:, :-1],), device=device) + out.append({"seq_ctx": seq_ctx, "shifted_labels": ids[:, 1:]}) + return out + + fsdp_idx = rank // ep_size + my_loss_data = build_loss_data(shards[fsdp_idx * n_microbatch : (fsdp_idx + 1) * n_microbatch]) + my_lcs = engine.model.build_loss_ctx_batch(my_loss_data, sp_mesh=None) + items = [ModelItem(seq_ctx=d["seq_ctx"], loss_ctx=lc) for d, lc in zip(my_loss_data, my_lcs)] + engine.model.zero_grad(set_to_none=True) + engine.train_step(items) + engine.model.scale_and_reduce_grad() + dist_grads = _full_grads(engine.model) + + # Router gate must receive a finite, non-zero gradient on every rank (balancing reaches it). + gate_names = [n for n in dist_grads if n.endswith(".gate.weight")] + assert gate_names, "no router gate params found" + for n in gate_names: + assert torch.isfinite(dist_grads[n]).all(), f"non-finite gate gradient {n}" + assert dist_grads[n].abs().sum() > 0, f"gate gradient {n} is all zero" + + if rank == 0: + # Reference: CE-only, ep=1, full batch. Non-gate params receive only CE gradient, so they + # must match tightly even with balancing enabled on the distributed side. + ref = MoE(config=_tiny_moe_config(ep_size=1, balancing=False, z_loss=False)).to(torch.bfloat16).to(device) + _, 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) + ref_loss_data = build_loss_data(shards) + loss_cfg = CELossConfig() + with _fake_dist_uninitialized(): + ref_lcs = loss_cfg.loss_ctx_cls.build_batches( + [loss_cfg.build(data={"shifted_labels": d["shifted_labels"]}, sp_mesh=None) for d in ref_loss_data] + ) + total = torch.zeros((), device=device) + for d, lc in zip(ref_loss_data, ref_lcs): + total = total + ref(seq_ctx=d["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(): + if name.endswith(".gate.weight"): # gate carries the balancing contribution; checked above + continue + 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"ep2+balancing non-gate grad ratio median={median} (want ~1.0)" + assert (ratios_t - 1.0).abs().max().item() < 0.3, f"ep2+balancing a non-gate grad drifted: {ratios_t}" + + def test_mtp_micro_batch_forward_runs(self): + # Guards the domino-EP MTP path in _micro_batch_forward, which calls aux_loss.accumulate for + # the MTP routed experts. Before the world_size cleanup fix, that call still passed the removed + # world_size arg (referencing a deleted z_world_size) and raised NameError/TypeError. The + # domino micro-batch path needs ep>1 + a real dispatcher (naive is ep=1 only); exercise it + # with mtp_config + intra_layer_micro_batch>1 + balancing/z on. It must run and log mtp loss. + self.create_pg("cuda") + device = "cuda" + seq_len = 32 + + # router_bias_update_speed=0 disables NoAuxRouter's update_bias, which is unrelated to this + # path and crashes on the extra MTP row that the domino block adds to tokens_per_expert + # (a pre-existing base issue, not part of the reduce-sum change under test here). + config = _tiny_moe_config( + ep_size=2, balancing=True, z_loss=True, mtp=True, dispatcher="all2all", router_bias_update_speed=0.0 + ) + torch.manual_seed(0) + fsdp_cfg = FSDPConfig(cpu_offload=False, ep_size=2, reduce_dtype=torch.bfloat16) + engine = TrainEngine( + model_cfg=config, optim_cfg=AdamWConfig(), fsdp_cfg=fsdp_cfg, intra_layer_micro_batch=2 + ) + engine.init_model_weights() + + # ep replicas (world=2, ep=2 -> one fsdp position) must consume identical data. + gen = torch.Generator().manual_seed(1234) + 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) # routes through _micro_batch_forward (domino-MTP accumulate) + engine.model.scale_and_reduce_grad() + + assert torch.isfinite(torch.tensor(info["total_loss"])), "total_loss is not finite" + assert "reduced_mtp_loss" in info["logs_info"], "MTP loss missing from logged curves" + grads = _full_grads(engine.model) + assert grads and all(torch.isfinite(g).all() for g in grads.values()), "non-finite MTP-path gradient" + + +class TestBalancingLossReduceSum(_ReduceSumDDPTestCase): + @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(_ReduceSumDDPTestCase): + @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}" + ) + + +class TestDenseFp32IgnoredParamReduce(_ReduceSumDDPTestCase): + @property + def world_size(self) -> int: + return 2 + + def test_fp32_ignored_param_grad_summed_across_ranks(self): + # fp32 ignored_params (matched by fp32_keys_pattern) are Replicate DTensors excluded from FSDP, + # so they get no reduce-scatter. Under reduce-sum, scale_and_reduce_grad must SUM their per-rank + # local grads over the replicate group -- otherwise the replicated copies carry different grads + # (diverge) and the effective gradient is local, not the global sum. Heterogeneous per-rank data + # makes the missing reduction observable. + self.create_pg("cuda") + device = "cuda" + rank = dist.get_rank() + world = dist.get_world_size() + fp32_name = "norm.weight" # HF key model.norm.weight + + cfg = Qwen3DenseConfig( + vocab_size=1024, + max_position_embeddings=512, + bos_token_id=1, + eos_token_id=2, + pad_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, + compile_cfg=False, + hf_save_cfg=HFSaveCfg(fp32_keys_pattern=[r"model\.norm\.weight"]), + ) + torch.manual_seed(0) + engine = TrainEngine( + model_cfg=cfg, optim_cfg=AdamWConfig(), fsdp_cfg=FSDPConfig(cpu_offload=False, reduce_dtype=torch.bfloat16) + ) + engine.init_model_weights() + + p = dict(engine.model.named_parameters())[fp32_name] + assert p.dtype == torch.float32, "fp32_keys_pattern param should stay fp32" + assert isinstance(p, DTensor) and any(isinstance(pl, Replicate) for pl in p.placements), ( + "fp32 ignored param should be a Replicate DTensor excluded from FSDP" + ) + + gen = torch.Generator().manual_seed(1234) + seqs = [torch.randint(0, 512, (1, 33), generator=gen) for _ in range(world)] + ids = seqs[rank].to(device) + seq_ctx = SequenceContext.from_input_ids(input_ids=(ids[:, :-1],), device=device) + loss_cfg = CELossConfig() + lc = loss_cfg.loss_ctx_cls.build_batches([loss_cfg.build(data={"shifted_labels": ids[:, 1:]}, sp_mesh=None)])[0] + engine.model.zero_grad(set_to_none=True) + engine.train_step([ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": lc})]) + + def local(g): + return (g.to_local() if isinstance(g, DTensor) else g).detach().float() + + # Pre-reduction per-rank local grad; the correct reduced value is the SUM over ranks. + pre = local(p.grad).clone() + gathered = [torch.zeros_like(pre) for _ in range(world)] + dist.all_gather(gathered, pre.contiguous()) + expected_sum = sum(gathered) + if rank == 0: + assert (gathered[0] - gathered[1]).abs().max() > 1e-6, "per-rank grads identical; test is vacuous" + + # A FSDP-sharded (Shard placement) param, to co-test alongside the Replicate one. Its exact + # reduce-sum value (== single-process global-batch gradient) is covered by the EP1/EP2 + # token-mean parity tests; here we only confirm the Shard path produces a finite, non-zero, + # cross-rank-consistent gradient (i.e. the reduce-scatter ran and did not corrupt it). + shard_name = "embed_tokens.weight" + sp = dict(engine.model.named_parameters())[shard_name] + assert isinstance(sp, DTensor) and any(isinstance(pl, Shard) for pl in sp.placements), ( + "expected a Shard-placement param to co-test with the Replicate one" + ) + shard_full = sp.grad.full_tensor().detach().float() + assert torch.isfinite(shard_full).all() and shard_full.abs().sum() > 0, "shard-param grad invalid" + + engine.model.scale_and_reduce_grad() + post = local(p.grad) + + # Replicate param: every rank's reduced grad equals the global SUM (and is therefore consistent). + torch.testing.assert_close(post, expected_sum, rtol=1e-3, atol=1e-3) + allpost = [torch.zeros_like(post) for _ in range(world)] + dist.all_gather(allpost, post.contiguous()) + shard_all = [torch.zeros_like(shard_full) for _ in range(world)] + dist.all_gather(shard_all, shard_full.contiguous()) + if rank == 0: + assert (allpost[0] - allpost[1]).abs().max() < 1e-4, "reduced fp32 (Replicate) grad not consistent" + assert (shard_all[0] - shard_all[1]).abs().max() < 1e-4, "Shard-param full grad not consistent" diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 74a56d4643..2d3bc07aaa 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -216,7 +216,10 @@ def train_step(self, data_batches: list[ModelItem]) -> TrainStepInfo: micro_batch_results = [] data_batch_info = self.model.pre_micro_batch_forward(data_batches) - total_loss = torch.tensor(0.0, device=DEVICE) + # Display total loss: sum of every loss term's per-rank display value (from each loss ctx's + # calibrate(), stored on output.calibrated_losses) across micro-batches -- NO cross-rank + # all_reduce, so each rank reports its own loss (equal to the global loss at world size 1). + display_total_loss = torch.tensor(0.0, device=DEVICE) for i in range(0, len(data_batches), intra_layer_micro_batch): ProberList.set_micro_batch_iter(micro_batch_iter) @@ -240,12 +243,13 @@ def train_step(self, data_batches: list[ModelItem]) -> TrainStepInfo: loss = self._get_total_loss(output) loss.backward() - total_loss += loss.detach() + for value in (output.calibrated_losses or {}).values(): + display_total_loss += value.detach().float() # call dump_forward_records after backward to record the recomputed activations ProberList.after_micro_iter_forward() batch_forward_info = self.model.post_micro_batch_forward(micro_batch_results) - return TrainStepInfo(total_loss=total_loss.item(), **data_batch_info, **batch_forward_info) + return TrainStepInfo(total_loss=display_total_loss.item(), **data_batch_info, **batch_forward_info) def from_hf(self, hf_path: str | Path, strict: bool = False): self.model.from_hf(hf_path=hf_path, strict=strict) diff --git a/xtuner/v1/loss/aux_loss.py b/xtuner/v1/loss/aux_loss.py index 80d1d22264..3128ca5932 100644 --- a/xtuner/v1/loss/aux_loss.py +++ b/xtuner/v1/loss/aux_loss.py @@ -3,10 +3,22 @@ from pydantic import BaseModel, ConfigDict from torch import distributed as dist from torch.distributed._functional_collectives import all_reduce +from typing_extensions import TypedDict from xtuner.v1.loss.moe_loss import BalancingLossContext, ZLossContext +class AuxLossFinalizeOutput(TypedDict): + """Backward-relevant outputs of :meth:`AuxLossContext.finalize`. + + Per-rank display values are not returned here; each loss context exposes them via ``calibrate()``. + """ + + balancing_loss: torch.Tensor | None + z_loss: torch.Tensor | None + tokens_per_expert_global: torch.Tensor + + class AuxLossScaler(torch.autograd.Function): """Inject an auxiliary loss into the main forward graph as a passthrough. @@ -91,7 +103,6 @@ def accumulate( z_ctx: list[ZLossContext] | ZLossContext | None = None, num_tokens_local: int = 0, num_tokens_global: torch.Tensor | None = None, - world_size: int = 1, ) -> torch.Tensor: """Accumulate routing statistics for one layer and inject z-loss into the main graph. @@ -111,9 +122,8 @@ def accumulate( num_tokens_local (int): Non-padding token count on this rank for the current forward. Required when any z-loss context is provided. num_tokens_global (torch.Tensor | None): All-reduced non-padding token count across - ranks (int64 scalar). Pass ``None`` when ``z_loss_global_average`` is off or no - process group is initialized. - world_size (int): World size that produced ``num_tokens_global``. + ranks (int64 scalar). Pass ``None`` when no process group is initialized + (single-process reference / eval). Returns: torch.Tensor: ``hidden_states`` augmented with the per-layer z-loss autograd hook. @@ -139,7 +149,6 @@ def accumulate( router_logits=selected_router_logits, num_tokens_local=num_tokens_local, num_tokens_global=num_tokens_global, - world_size=world_size, ) hidden_states = AuxLossScaler.apply(hidden_states, z_loss_l) @@ -151,15 +160,20 @@ def finalize( balancing_ctx: list[BalancingLossContext] | BalancingLossContext | None, z_ctx: list[ZLossContext] | ZLossContext | None, non_pad_token: int, - ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]: - """Finalize split auxiliary losses and expert counts from runtime - state.""" + ) -> "AuxLossFinalizeOutput": + """Finalize auxiliary losses and expert counts from runtime state. + + Returns: + AuxLossFinalizeOutput: ``balancing_loss`` / ``z_loss`` carry the backward graph and the + globally reduced ``tokens_per_expert_global``. Per-rank display values are produced by + each context's ``calibrate()``, not returned here. + """ tokens_per_expert_local, tokens_per_expert_global = self._cal_tokens_per_expert() balancing_loss: torch.Tensor | None = None balancing_list = _as_list(balancing_ctx) if balancing_list: - partials = [ + losses = [ ctx.finalize( tokens_per_expert_local=tokens_per_expert_local, tokens_per_expert_global=tokens_per_expert_global, @@ -169,15 +183,19 @@ def finalize( ) for ctx in balancing_list ] - balancing_loss = partials[0] if len(partials) == 1 else torch.stack(partials).sum(dim=0) + balancing_loss = losses[0] if len(losses) == 1 else torch.stack(losses).sum(dim=0) z_loss: torch.Tensor | None = None z_list = _as_list(z_ctx) if z_list: - partials = [ctx.finalize() for ctx in z_list] - z_loss = partials[0] if len(partials) == 1 else torch.stack(partials).sum(dim=0) + z_losses = [ctx.finalize() for ctx in z_list] + z_loss = z_losses[0] if len(z_losses) == 1 else torch.stack(z_losses).sum(dim=0) - return balancing_loss, z_loss, tokens_per_expert_global + return AuxLossFinalizeOutput( + balancing_loss=balancing_loss, + z_loss=z_loss, + tokens_per_expert_global=tokens_per_expert_global, + ) def _cal_tokens_per_expert(self) -> tuple[torch.Tensor, torch.Tensor]: """Stack per-layer expert counts and produce both local and globally diff --git a/xtuner/v1/loss/base_loss_ctx.py b/xtuner/v1/loss/base_loss_ctx.py index 531a6dfdab..b8b1681cdd 100644 --- a/xtuner/v1/loss/base_loss_ctx.py +++ b/xtuner/v1/loss/base_loss_ctx.py @@ -22,20 +22,21 @@ # 1. Compute the global loss mask sum among dp, sp and grad accumulation: # global_loss_mask_sum = all_reduce(sum([loss_mask.sum() for loss_mask in loss_masks_grad_acc]), op=dist.ReduceOp.SUM, group=world) # = (m00 + m01 + m02 + m03 + m10 + m11 + m12 + m13) -# 2. Compute the iter loss, take rank0 iter0 as an example: +# 2. Compute the iter loss (this rank's local component), take rank0 iter0 as an example: # a. loss_{rank0iter0} = (l00 * w00 * m00 + l01 * w01 * m01) # b. loss_{rank0iter0} = loss_{rank0iter0} / global_loss_mask_sum # = (l00 * w00 * m00 + l01 * w01 * m01) / # (m00 + m01 + m02 + m03 + m10 + m11 + m12 + m13) -# c. loss_{rank0iter0} = all_reduce_autograd(loss_{rank0iter0}, op=dist.ReduceOp.SUM, group=world) -# = (l00 * w00 * m00 + l01 * w01 * m01 + l02 * w02 * m02 + l03 * w03 * m03) / -# (m00 + m01 + m02 + m03 + m10 + m11 + m12 + m13) -# 3. Compute the step loss: +# Under reduce-sum gradients this stays a local component: it is NOT all-reduced across ranks. +# Each rank keeps its own numerator over the shared (detached) global denominator; the +# cross-rank aggregation happens on the gradients (FSDP / scale_and_reduce_grad SUM), whose +# sum over ranks equals the global-batch loss gradient. +# 3. Compute the step loss (this rank's local component summed over grad-acc iters): # step_loss = loss_{rank0iter0} + loss_{rank0iter1} -# = (l00 * w00 * m00 + l01 * w01 * m01 + l02 * w02 * m02 + l03 * w03 * m03 + -# l10 * w10 * m10 + l11 * w11 * m11 + l12 * w12 * m12 + l13 * w13 * m13) / +# = (l00 * w00 * m00 + l01 * w01 * m01 + l10 * w10 * m10 + l11 * w11 * m11) / # (m00 + m01 + m02 + m03 + m10 + m11 + m12 + m13) -# It's equivalent to loss calculation in sp1, dp1 and grad acc 1. +# Its SUM over all ranks equals the sp1/dp1/grad-acc-1 loss; that global value is restored for +# display by a separate detached SUM all_reduce (see the reduce-sum design's logging section). class BaseLossKwargs(BaseModel): diff --git a/xtuner/v1/loss/ce_loss.py b/xtuner/v1/loss/ce_loss.py index eba945fae7..779466c135 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 @@ -106,6 +105,11 @@ class LMHeadLossContext(BaseLossContext): loss_cfg: CELossConfig loss_kwargs: CELossKwargs + # Display state for calibrate(): coefficient stashed by build_batches (global_denom / + # step_grad_tokens) and this rank's detached backward loss stashed by forward(). + _display_coeff: torch.Tensor + _backward_loss: torch.Tensor + def __init__(self, loss_cfg: CELossConfig, loss_kwargs: CELossKwargs): super().__init__(loss_cfg, loss_kwargs) @@ -178,10 +182,21 @@ def build_batches( # type: ignore[override] if dist.is_initialized(): dist.all_reduce(global_denominator, op=dist.ReduceOp.SUM) + # Coefficient to turn each rank's local backward loss (local token sum / global denominator) + # back into this rank's per-token mean for display: multiply by global_denominator, divide by + # this rank's own grad-token count across the whole step. Summing `calibrate()` over the + # micro-batches then yields this rank's per-token-mean loss with no cross-rank all_reduce, and + # equals the global loss when world size is 1 (global_denominator == step_grad_tokens). + step_grad_tokens = torch.zeros((), device=global_denominator.device) + for loss_ctx in loss_ctx_list: + step_grad_tokens += (loss_ctx.loss_kwargs.shifted_labels != loss_cfg.ignore_idx).sum() + display_coeff = (global_denominator / step_grad_tokens.clamp_min(1.0)).detach() + for loss_ctx in loss_ctx_list: loss_ctx._batch_size = len(loss_ctx_list) assert loss_ctx.loss_kwargs.loss_weight is not None loss_ctx.loss_kwargs.loss_weight /= global_denominator + 1e-12 + loss_ctx._display_coeff = display_coeff return loss_ctx_list def loss_fn( @@ -280,14 +295,18 @@ def forward( if not isinstance(extra_info, ModelForwardExtraLogInfo): extra_info = ModelForwardExtraLogInfo(extra_info) - 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 per-rank + # display value is produced by `calibrate()`; stash the detached loss for it. + self._backward_loss = loss.detach() return loss, (logits, extra_info) + def calibrate(self) -> torch.Tensor: + """This rank's per-token-mean CE for display (detached, no cross-rank + all_reduce).""" + return self._backward_loss * self._display_coeff + # Deprecated: Use LMHeadLossContext instead. Will be removed in version 1.1.0 CELossContext = LMHeadLossContext diff --git a/xtuner/v1/loss/moe_loss.py b/xtuner/v1/loss/moe_loss.py index c100cefebb..cec7e05dc1 100644 --- a/xtuner/v1/loss/moe_loss.py +++ b/xtuner/v1/loss/moe_loss.py @@ -5,7 +5,6 @@ from cyclopts import Parameter from pydantic import BaseModel, ConfigDict from torch import distributed as dist -from torch.distributed._functional_collectives import all_reduce from xtuner.v1.utils.device import get_device @@ -13,38 +12,17 @@ DEVICE = get_device() -class _AllReduce(torch.autograd.Function): - @staticmethod - def forward(ctx, op, group, tensor): - ctx.group = group - ctx.op = op - tensor = tensor.clone(memory_format=torch.contiguous_format) - tensor = all_reduce(tensor, op, group=group) - return tensor - - @staticmethod - def backward(ctx, grad_output): - return (None, None) + (_AllReduce.apply(ctx.op, ctx.group, grad_output),) - - -def all_reduce_autograd(tensor, op, group): - return _AllReduce.apply(op, group, tensor) - - class BalancingLossConfig(BaseModel): """Balancing loss configuration for MoE models. Args: balancing_loss_alpha (float): Weight for the balancing loss. Defaults to 0.001. - balancing_loss_global_average (bool): Whether to perform global averaging across all ranks. - Defaults to True. router_scoring_func (str): Router scoring function type. Options are "sigmoid" and "softmax". Defaults to "softmax". """ model_config = ConfigDict(extra="forbid") balancing_loss_alpha: Annotated[float, Parameter(help="weight for balancing loss")] = 0.001 - balancing_loss_global_average: Annotated[bool, Parameter(help="global average for balancing loss")] = True router_scoring_func: Annotated[Literal["sigmoid", "softmax"], Parameter(help="router scoring function")] = ( "softmax" ) @@ -84,6 +62,8 @@ def __init__(self, loss_cfg: BalancingLossConfig, loss_kwargs: BalancingLossKwar # Per-layer differentiable accumulator. tokens_per_expert is owned by AuxLossContext # and passed in at finalize() time to avoid duplicate storage / duplicate all_reduce. self.routing_weights_sum_list: list[torch.Tensor] = [] + # Detached per-rank display value set by finalize(), returned by calibrate(). + self._calibrated: torch.Tensor | None = None @staticmethod def build_batches( @@ -139,34 +119,64 @@ def finalize( non_pad_token (int): Number of non-padding tokens on this rank. Returns: - torch.Tensor: Final balancing loss. + torch.Tensor: This rank's balancing loss carrying the autograd graph for backward. Under + reduce-sum it is computed from this rank's own ``local_gating_sum`` with global detached + statistics; cross-rank aggregation happens on the gradients (FSDP / scale_and_reduce_grad + SUM), so summing over ranks reproduces the global balancing loss. The per-rank display + value is computed separately by ``calibrate()`` from local statistics. """ routing_weights_sum_list = self.routing_weights_sum_list self.routing_weights_sum_list = [] if self.loss_cfg.balancing_loss_alpha == 0 or not routing_weights_sum_list: - return torch.tensor(0.0, device=tokens_per_expert_local.device, dtype=torch.float32) + self._calibrated = torch.tensor(0.0, device=tokens_per_expert_local.device, dtype=torch.float32) + return self._calibrated local_gating_sum = torch.stack(routing_weights_sum_list, dim=0) + 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 + if dist.is_initialized(): tokens_global = tokens_per_expert_global.sum(-1) seqlen_global = tokens_global // num_experts_per_tok - - routing_weights_sum_global = all_reduce_autograd(local_gating_sum, "sum", group) - routing_weights_mean_global = routing_weights_sum_global / seqlen_global.unsqueeze(-1) scale_global = n_routed_experts / tokens_global - tokens_per_expert_for_loss = tokens_per_expert_global + routing_weights_mean = local_gating_sum / seqlen_global.unsqueeze(-1) + loss_vec = scale_global * (tokens_per_expert_global * routing_weights_mean).sum(-1) else: + # Single-process path (no process group). Numerically identical to the distributed branch + # at world size 1 (tokens_per_expert_global == tokens_per_expert_local, seqlen_global == + # valid_tokens); kept so reference / eval (dist uninitialized) still works. valid_tokens = max(non_pad_token, 1) scale_global = n_routed_experts / (valid_tokens * num_experts_per_tok) routing_weights_mean_global = local_gating_sum / valid_tokens - tokens_per_expert_for_loss = tokens_per_expert_local + loss_vec = scale_global * (tokens_per_expert_local * routing_weights_mean_global).sum(-1) + + loss = loss_vec.sum() * alpha / self._batch_size + # Display value (detached, no all_reduce): this rank's balancing loss from LOCAL statistics + # (its own tokens_per_expert / seqlen), a readable per-rank number. This is a display-only + # computation, NOT a backward mode -- it does not reintroduce the removed non-global averaging. + self._calibrated = self._local_balancing_loss( + local_gating_sum, tokens_per_expert_local, n_routed_experts, num_experts_per_tok, non_pad_token + ) + return loss - loss = scale_global * (tokens_per_expert_for_loss * routing_weights_mean_global).sum(-1) - loss = loss.sum() * self.loss_cfg.balancing_loss_alpha - return loss / self._batch_size + def _local_balancing_loss( + self, + local_gating_sum: torch.Tensor, + tokens_per_expert_local: torch.Tensor, + n_routed_experts: int, + num_experts_per_tok: int, + non_pad_token: int, + ) -> torch.Tensor: + valid_tokens = max(non_pad_token, 1) + scale_local = n_routed_experts / (valid_tokens * num_experts_per_tok) + routing_weights_mean_local = local_gating_sum.detach() / valid_tokens + loss = scale_local * (tokens_per_expert_local * routing_weights_mean_local).sum(-1) + return (loss.sum() * self.loss_cfg.balancing_loss_alpha / self._batch_size).detach() + + def calibrate(self) -> torch.Tensor: + """This rank's balancing loss (from local statistics) for display + (detached, no all_reduce).""" + assert self._calibrated is not None, "finalize() must be called before calibrate()" + return self._calibrated @property def batch_size(self) -> int: @@ -178,13 +188,10 @@ class ZLossConfig(BaseModel): Args: z_loss_alpha (float): Weight for the z-loss. Defaults to 0.001. - z_loss_global_average (bool): Whether to perform global averaging across all ranks. - Defaults to True. """ model_config = ConfigDict(extra="forbid") z_loss_alpha: Annotated[float, Parameter(help="weight for z-loss")] = 0.001 - z_loss_global_average: Annotated[bool, Parameter(help="global average for z-loss")] = True def build(self) -> "ZLossContext": """Build ZLossContext. @@ -215,11 +222,13 @@ def __init__(self, loss_cfg: ZLossConfig, loss_kwargs: ZLossKwargs): self.loss_cfg = loss_cfg self.loss_kwargs = loss_kwargs self._batch_size = 1 - # Z-loss is folded into a running detached scalar for logging only. The differentiable - # per-layer scalar is injected back into the main forward graph via AuxLossScaler at - # accumulate() time, so we never need to keep per-layer logsum tensors around. This is - # the memory-saving pattern adapted from Megatron's MoEAuxLossAutoScaler. + # Z-loss is backward-only: the differentiable per-layer scalar is injected into the main graph + # via AuxLossScaler at accumulate() time (memory-saving pattern from Megatron's + # MoEAuxLossAutoScaler). For display we self-maintain a single detached running scalar holding + # THIS RANK's z-loss mean (the raw `logsumexp^2` mean per layer, without the reduce-sum + # `num_tokens_local/denom_global` factor), i.e. a clean world-size-independent per-rank value. self._running_loss_for_log: torch.Tensor | None = None + self._calibrated: torch.Tensor | None = None @staticmethod def build_batches( @@ -245,7 +254,6 @@ def accumulate( router_logits: torch.Tensor, num_tokens_local: int, num_tokens_global: torch.Tensor | None, - world_size: int, ) -> torch.Tensor: """Compute z-loss for one layer and return it as a scalar with autograd attached. @@ -261,10 +269,8 @@ def accumulate( num_tokens_local (int): Number of non-padding tokens on this rank for the current forward (constant across MoE layers in a single forward). num_tokens_global (torch.Tensor | None): All-reduced non-padding token count across - ranks, as an int64 scalar tensor. ``None`` when ``z_loss_global_average`` is off - or the process group is not initialized. - world_size (int): Number of ranks contributing to ``num_tokens_global``. Ignored when - ``num_tokens_global`` is ``None``. + ranks, as an int64 scalar tensor. ``None`` when the process group is not initialized + (single-process reference / eval). Returns: torch.Tensor: Per-layer z-loss as a 0-d tensor with autograd graph back to @@ -276,30 +282,41 @@ def accumulate( return zero denom_local = max(num_tokens_local, 1) - loss = torch.logsumexp(router_logits, dim=-1).square().sum() / denom_local - - if self.loss_cfg.z_loss_global_average and num_tokens_global is not None: - # Equivalent to scaling each layer's local loss by num_tokens_local * world_size / - # num_tokens_global, matching the original list-based finalize formula. + base = torch.logsumexp(router_logits, dim=-1).square().sum() / denom_local + + loss = base + if num_tokens_global is not None: + # Distributed: the injected z-loss stays as this rank's local component (its share of the + # global z-loss, WITHOUT any `× world_size`). Cross-rank aggregation happens on the + # gradients via the FSDP SUM reduce-scatter; summing over ranks reproduces the global + # z-loss. Single-process (num_tokens_global is None) keeps `loss = base`, its W=1 case. denom_global = torch.clamp(num_tokens_global, min=1) - loss = loss * num_tokens_local * world_size / denom_global + loss = base * num_tokens_local / denom_global loss = loss * self.loss_cfg.z_loss_alpha / self._batch_size - self._update_running(loss.detach()) + # Display value: this rank's raw z-loss mean (drop the `num_tokens_local/denom_global` + # reduce-sum factor), which is world-size independent and readable per rank. + self._update_running((base * self.loss_cfg.z_loss_alpha / self._batch_size).detach()) return loss def finalize(self) -> torch.Tensor: - """Return the accumulated z-loss as a detached scalar for logging only. + """Return this rank's accumulated z-loss mean as a detached scalar for + the output field. - The differentiable contribution has already been injected into the main forward graph at - each ``accumulate()`` call, so this value carries no autograd graph; it exists purely to - populate the logging field on the model output. + The differentiable contribution was injected into the main graph at each ``accumulate()`` + call, so this value carries no autograd graph. It is this rank's own z-loss mean (no + cross-rank all_reduce); ``calibrate()`` returns the same value for the display pipeline. """ value = self._running_loss_for_log self._running_loss_for_log = None - if value is None: - return torch.tensor(0.0, device=DEVICE, dtype=torch.float32) - return value + self._calibrated = value if value is not None else torch.tensor(0.0, device=DEVICE, dtype=torch.float32) + return self._calibrated + + def calibrate(self) -> torch.Tensor: + """This rank's z-loss mean for display (detached, no cross-rank + all_reduce).""" + assert self._calibrated is not None, "finalize() must be called before calibrate()" + return self._calibrated def _update_running(self, value: torch.Tensor) -> None: if self._running_loss_for_log is None: diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index c02c4dc9dd..0b564e0f37 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -404,6 +404,10 @@ class ModelOutputs(PydanticBaseModel): logits: torch.Tensor | None = None loss: torch.Tensor | None = None # TODO: `forward_only` mode for RL extra_info: ModelForwardExtraLogInfo | dict | None = None # TODO: `forward_only` mode for RL + # Detached per-rank display value per loss term (from each loss ctx's ``calibrate()``), keyed by + # the loss name (e.g. ``"llm_loss"``, ``"balancing_loss"``). Aggregated per-rank for logging with + # NO cross-rank all_reduce; each rank displays its own loss. + calibrated_losses: dict[str, torch.Tensor] | None = None def free_nongrad_feature(self): """Release large intermediate tensors not needed for backward or @@ -576,7 +580,31 @@ def from_hf( return loaded_keys, unloaded_keys, missing_keys def scale_and_reduce_grad(self): - return + # Params excluded from FSDP sharding (e.g. fp32 ``ignored_params`` matched by + # ``fp32_keys_pattern``) are distributed as ``Replicate`` DTensors, so they get no + # reduce-scatter. Under reduce-sum their gradient is this rank's local component; without a + # cross-rank reduction the replicated copies would carry different grads and diverge. SUM + # (no divide) their grads over the replicate group here. FSDP-sharded params are handled by + # the reduce-scatter and are skipped (no Replicate placement). + self._reduce_replicated_ignored_grads(self.trainable_parameters()) + + def _reduce_replicated_ignored_grads(self, params: Iterable[tuple[str, nn.Parameter]]) -> None: + # Bucket grads by the process group of each Replicate mesh dim, then issue one coalesced SUM + # all_reduce per group. Meshes are indexed by dim (not name): fp32 ignored params land on an + # unnamed 1D world mesh, so a name-based lookup would fail. A param replicated on multiple mesh + # dims is reduced over each dim's group in turn, which sums it over the full product of ranks. + grads_by_group: dict[dist.ProcessGroup, list[torch.Tensor]] = {} + for _, param in params: + if param.grad is None or not isinstance(param, DTensor): + continue + grad = param.grad.to_local() if isinstance(param.grad, DTensor) else param.grad + for dim, placement in enumerate(param.placements): + if isinstance(placement, Replicate): + grads_by_group.setdefault(param.device_mesh.get_group(dim), []).append(grad) + for group, grads in grads_by_group.items(): + with dist._coalescing_manager(group=group): + for grad in grads: + dist.all_reduce(grad, dist.ReduceOp.SUM, group=group) def to_hf_key_list(self, key: str) -> list[str]: raise NotImplementedError() @@ -585,6 +613,41 @@ def trainable_parameters(self): params = [(name, param) for name, param in self.named_parameters() if param.requires_grad] return params + def set_gradient_reduce_sum(self) -> None: + """Switch every sharded ``FSDPModule`` under this model to pure SUM + gradient reduction. + + FSDP2 reduce-scatters gradients with ``ReduceOp.AVG`` (divide by the sharding group size) + by default. This helper sets the gradient divide factor to 1 and forces plain + ``ReduceOp.SUM`` communication, so the reduce-scatter accumulates each rank's local-component + gradient without any division. + + The two calls must be paired. Setting only the divide factor routes FSDP through + ``_make_nccl_premul_sum(1 / factor)``, whose NCCL PreMulSum silently reduces bf16 gradients to + all-zeros on torch 2.10; ``set_force_sum_reduction_for_comms(True)`` instead keeps a plain + ``ReduceOp.SUM`` that is exact in bf16 without upcasting the reduction dtype. + + ``fully_shard`` wraps modules in place, so nested sharded children remain reachable through + ``self.modules()``; one root-level pass therefore covers every layer sharded during + ``fully_shard``. + """ + if not ( + hasattr(FSDPModule, "set_gradient_divide_factor") + and hasattr(FSDPModule, "set_force_sum_reduction_for_comms") + ): + raise RuntimeError( + "set_gradient_reduce_sum requires the FSDP2 gradient-reduction APIs " + "`FSDPModule.set_gradient_divide_factor` and `FSDPModule.set_force_sum_reduction_for_comms`, " + "available since torch 2.10. The installed torch does not expose them, and falling back to " + "`set_gradient_divide_factor` alone would silently zero bf16 gradients." + ) + for module in self.modules(): + if isinstance(module, FSDPModule): + # torch < 2.10 type stubs do not declare these FSDP2 setters; guarded by the + # hasattr check above, they exist at runtime on torch >= 2.10. + module.set_gradient_divide_factor(1.0) # type: ignore[operator] + module.set_force_sum_reduction_for_comms(True) # type: ignore[operator] + def fully_shard( self, fsdp_config: FSDPConfig, @@ -621,6 +684,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( @@ -1305,41 +1372,33 @@ def pre_micro_batch_forward(self, data_batches: Sequence[ModelItem]) -> DataBatc def post_micro_batch_forward(self, batch_outputs: Sequence[ModelOutputs]) -> BatchForwardInfo: train_engine_extra_info = ModelForwardExtraLogInfo() - - local_total_loss = torch.tensor(0.0, device=DEVICE) - reduced_other_losses: dict[str, float] = {} - + logs_info = self.reduce_display_losses(batch_outputs) for output in batch_outputs: - output_copy = output.model_copy() - for name in output_copy.model_fields: - obj = getattr(output_copy, name) - if "loss" in name and isinstance(obj, torch.Tensor): - loss_item = obj.item() - local_total_loss += loss_item - reduced_name = f"reduced_{name}" - - if reduced_name not in reduced_other_losses: - reduced_other_losses[reduced_name] = loss_item - else: - reduced_other_losses[reduced_name] += loss_item + if "extra_info" in output: + train_engine_extra_info.append(output["extra_info"]) + return BatchForwardInfo(logs_info=logs_info, extra_info=train_engine_extra_info) - if "extra_info" in output_copy: - extra_info = output["extra_info"] - train_engine_extra_info.append(extra_info) + def reduce_display_losses(self, batch_outputs: Sequence[ModelOutputs]) -> dict[str, float]: + """Aggregate each loss term's per-rank display value across micro- + batches for logging. - for name, loss in reduced_other_losses.items(): - tensor_loss = torch.tensor(loss, device=DEVICE) - dist.all_reduce(tensor_loss.div_(dist.get_world_size()), op=dist.ReduceOp.SUM) - reduced_other_losses[name] = tensor_loss.item() + Each loss context computes its own detached per-rank display value via ``calibrate()`` (this + rank's per-token / per-rank mean, no cross-rank all_reduce); the model forward stores those on + ``output.calibrated_losses``. Here we sum them over the step's micro-batches, so every rank + logs ITS OWN loss (equal to the global loss at world size 1). Subclasses may override. - if "reduced_loss" in reduced_other_losses: - reduced_other_losses["reduced_llm_loss"] = reduced_other_losses.pop("reduced_loss") + Args: + batch_outputs (Sequence[ModelOutputs]): The per-micro-batch model outputs of one step. - ret = BatchForwardInfo( - logs_info=reduced_other_losses, - extra_info=train_engine_extra_info, - ) - return ret + Returns: + dict[str, float]: ``reduced_`` -> this rank's display loss for each term. + """ + summed: dict[str, torch.Tensor] = {} + for output in batch_outputs: + for name, value in (output.calibrated_losses or {}).items(): + contribution = value.detach().float() + summed[name] = contribution if name not in summed else summed[name] + contribution + return {f"reduced_{name}": value.item() for name, value in summed.items()} def _get_save_dtype(self, name: str, dtype: torch.dtype) -> torch.dtype: patterns = self.config.hf_save_cfg.fp32_keys_pattern diff --git a/xtuner/v1/model/compose/base.py b/xtuner/v1/model/compose/base.py index 51eb1fa02b..963898c0a3 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): @@ -293,7 +299,16 @@ def post_micro_batch_forward(self, batch_outputs: Sequence[ModelOutputs]) -> Bat return self.language_model.post_micro_batch_forward(batch_outputs) def scale_and_reduce_grad(self): + # The language model reduces its own grads (MoE all_reduces replicated params; Dense uses the + # BaseModel replicate reduction). This compose model's OWN params -- vision tower, projector, + # and the root wrap -- may also hold replicated fp32 ignored_params that get no reduce-scatter, + # so SUM-reduce those here. Exclude the language model's params (handled above) to avoid a + # double reduction. self.language_model.scale_and_reduce_grad() + own_params = [ + (name, param) for name, param in self.trainable_parameters() if not name.startswith("language_model.") + ] + self._reduce_replicated_ignored_grads(own_params) @override def build_loss_ctx_batch( # type: ignore[override] 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 bf4ec01bd6..667c2717d3 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -96,6 +96,48 @@ MOE_EP_COMPILE_CFG.pop("xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward") +def _sum_calibrate(ctx: object | list | None) -> torch.Tensor | None: + """Sum the detached per-rank display value ``calibrate()`` over one or more + loss contexts. + + Aux contexts fan out to a list (one per micro-batch) whose ``finalize()`` outputs are summed for + backward; their per-rank display values sum the same way. Returns ``None`` when no context. + """ + if ctx is None: + return None + ctxs = ctx if isinstance(ctx, list) else [ctx] + if not ctxs: + return None + total = ctxs[0].calibrate() + for c in ctxs[1:]: + total = total + c.calibrate() + return total + + +def _build_calibrated_losses( + lm_display: torch.Tensor, + balancing_ctx: object | list | None, + z_ctx: object | list | None, + mtp_display: torch.Tensor | None, +) -> dict[str, torch.Tensor]: + """Collect each loss term's detached per-rank display value (from + ``calibrate()``) for logging. + + Stored on ``output.calibrated_losses`` and aggregated per-rank (no all_reduce) by the display + pipeline, so each rank logs its own loss. + """ + calibrated: dict[str, torch.Tensor] = {"llm_loss": lm_display.detach()} + balancing_display = _sum_calibrate(balancing_ctx) + if balancing_display is not None: + calibrated["balancing_loss"] = balancing_display + z_display = _sum_calibrate(z_ctx) + if z_display is not None: + calibrated["z_loss"] = z_display + if mtp_display is not None: + calibrated["mtp_loss"] = mtp_display.detach() + return calibrated + + class MoEModelOutputs(ModelOutputs): router_logits: dict[str, torch.Tensor] | None = None router_weights: dict[str, torch.Tensor] | None = None @@ -227,24 +269,22 @@ def _z_loss_dist_token_count( z_ctx: list[ZLossContext] | ZLossContext | None, num_tokens_local: int, device: torch.device | str | int, - ) -> tuple[torch.Tensor | None, int]: + ) -> torch.Tensor | None: """Compute the cross-rank non-padding token count needed by the z-loss inline path. - Returns ``(num_tokens_global, world_size)``. ``num_tokens_global`` is ``None`` (i.e. skip - global averaging) when there is no z-loss context, when the configured z-loss is not - global-average, or when no process group is initialized. + Returns the global non-padding token count, or ``None`` when there is no z-loss context or no + process group is initialized (single-process reference / eval, which is the world-size-1 + case). """ if z_ctx is None: - return None, 1 - first = z_ctx[0] if isinstance(z_ctx, list) else z_ctx - if not first.loss_cfg.z_loss_global_average or not dist.is_initialized(): - return None, 1 + return None + if not dist.is_initialized(): + return None n = torch.tensor(num_tokens_local, device=device, dtype=torch.int64) group = dist.group.WORLD assert group is not None - n_global = all_reduce(n, "sum", group) - return n_global, dist.get_world_size() + return all_reduce(n, "sum", group) def _extract_aux_loss_ctx( self, @@ -466,7 +506,7 @@ def _micro_batch_forward( [{} for _ in range(len(seq_ctx_list))] if keep_router else [] ) balancing_ctx, z_ctx = self._extract_aux_loss_ctx(loss_ctx_list) - num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, cat_mask.device) + num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, cat_mask.device) # Process through layers cat_seq_ctx: SequenceContext | None = None @@ -548,11 +588,11 @@ def _micro_batch_forward( z_ctx=z_ctx, num_tokens_local=non_pad_token, num_tokens_global=num_tokens_global, - world_size=z_world_size, ) assert hidden_states_list, "XTuner Internal Error, found empty hidden states for domino EP" + mtp_calibrated: torch.Tensor | None = None if self.mtp_block is not None: assert self.config.mtp_config is not None @@ -579,6 +619,8 @@ def _micro_batch_forward( ) mtp_losses = torch.tensor(0.0, device=DEVICE) + # Per-rank display value (from each depth's calibrate()), mirroring `mtp_losses`. + mtp_display = torch.tensor(0.0, device=DEVICE) has_mtp_loss = False for micro_batch_idx, (loss_ctx_dict, mtp_outputs) in enumerate(zip(loss_ctx_list, mtp_outputs_per_mb)): mtp_loss_ctx_list = loss_ctx_dict.get("mtp") @@ -586,15 +628,18 @@ def _micro_batch_forward( continue micro_batch_mtp_losses = torch.tensor(0.0, device=DEVICE) + micro_batch_mtp_display = torch.tensor(0.0, device=DEVICE) for mtp_idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)): mtp_hidden_states, mtp_router_results, _ = mtp_hidden mtp_loss, _ = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx)) micro_batch_mtp_losses += mtp_loss + micro_batch_mtp_display += cast(MTPLossContext, mtp_ctx).calibrate() if keep_router: router_logits_list[micro_batch_idx][f"mtp_layer{mtp_idx}"] = mtp_router_results mtp_losses += micro_batch_mtp_losses / len(mtp_loss_ctx_list) + mtp_display += micro_batch_mtp_display / len(mtp_loss_ctx_list) has_mtp_loss = True if has_mtp_loss: @@ -630,10 +675,10 @@ def _micro_batch_forward( z_ctx=z_ctx, num_tokens_local=non_pad_token, num_tokens_global=num_tokens_global, - world_size=z_world_size, ) output["mtp_loss"] = mtp_losses * self.config.mtp_config.loss_scaling_factor + mtp_calibrated = mtp_display * self.config.mtp_config.loss_scaling_factor # Apply final norm to all micro-batches cat_hidden_states = torch.cat(hidden_states_list, dim=1) @@ -643,6 +688,9 @@ def _micro_batch_forward( # Extract LM loss context from dict lm_loss_ctx_list = [loss_ctx_dict["lm"] for loss_ctx_dict in loss_ctx_list] cat_loss_ctx = type(lm_loss_ctx_list[0]).cat(lm_loss_ctx_list) + # All micro-batch lm contexts of a step share the same display coefficient (build_batches + # computes one per step); carry it onto the concatenated context so calibrate() works. + cat_loss_ctx._display_coeff = lm_loss_ctx_list[0]._display_coeff loss, (logits, extra_info) = self.lm_head(cat_hidden_states, cast(LMHeadLossContext, cat_loss_ctx)) # Aggregate losses (mean across micro-batches) @@ -652,17 +700,19 @@ def _micro_batch_forward( moe_extra_info.append(extra_info) output["extra_info"] = moe_extra_info - split_aux_output = self.aux_loss.finalize( + aux_out = self.aux_loss.finalize( balancing_ctx=balancing_ctx, z_ctx=z_ctx, non_pad_token=non_pad_token, ) - balancing_loss, z_loss, tokens_per_expert_global = split_aux_output - if balancing_loss is not None: - output["balancing_loss"] = balancing_loss - if z_loss is not None: - output["z_loss"] = z_loss - output["tokens_per_expert_global"] = tokens_per_expert_global + if aux_out["balancing_loss"] is not None: + output["balancing_loss"] = aux_out["balancing_loss"] + if aux_out["z_loss"] is not None: + output["z_loss"] = aux_out["z_loss"] + output["tokens_per_expert_global"] = aux_out["tokens_per_expert_global"] + output["calibrated_losses"] = _build_calibrated_losses( + cast(LMHeadLossContext, cat_loss_ctx).calibrate(), balancing_ctx, z_ctx, mtp_calibrated + ) if keep_router: # TODO: Returning router logits is costly. @@ -724,7 +774,7 @@ def _forward( # Hoisted out of the per-layer accumulate path: mask is constant across layers. nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1] non_pad_token = nonpad_indices.numel() - num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device) + num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device) for idx, decoder_layer in self.layers.items(): if int(idx) < self.config.first_k_dense_replace: @@ -766,7 +816,6 @@ def _forward( z_ctx=z_ctx, num_tokens_local=non_pad_token, num_tokens_global=num_tokens_global, - world_size=z_world_size, ) if self.config.return_hidden_states: @@ -783,6 +832,7 @@ def _forward( output["extra_info"] = extra_info # MTP forward pass and loss computation + mtp_calibrated: torch.Tensor | None = None if ( self.mtp_block is not None and loss_ctx is not None @@ -796,9 +846,7 @@ def _forward( # MTP uses its own mask; main mask's non-pad indices do not apply. mtp_nonpad_indices = torch.nonzero(mtp_seq_ctx.mask, as_tuple=True)[1] mtp_non_pad_token = mtp_nonpad_indices.numel() - mtp_num_tokens_global, mtp_z_world_size = self._z_loss_dist_token_count( - z_ctx, mtp_non_pad_token, mtp_seq_ctx.mask.device - ) + mtp_num_tokens_global = self._z_loss_dist_token_count(z_ctx, mtp_non_pad_token, mtp_seq_ctx.mask.device) # Forward through MTP block mtp_outputs = self.mtp_block( @@ -810,6 +858,8 @@ def _forward( # Compute MTP losses for each depth mtp_losses = torch.tensor(0.0, device=DEVICE) + # Per-rank display value (from each depth's calibrate()), mirroring `mtp_losses`. + mtp_display = torch.tensor(0.0, device=DEVICE) for idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)): mtp_hidden_states, mtp_router_results, mtp_router_weights = mtp_hidden @@ -828,10 +878,10 @@ def _forward( z_ctx=z_ctx, num_tokens_local=mtp_non_pad_token, num_tokens_global=mtp_num_tokens_global, - world_size=mtp_z_world_size, ) mtp_loss, _ = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx)) mtp_losses += mtp_loss + mtp_display += cast(MTPLossContext, mtp_ctx).calibrate() # Average MTP losses across depths and scale mtp_losses = mtp_losses / len(mtp_loss_ctx_list) @@ -839,18 +889,22 @@ def _forward( # Add to total loss output["mtp_loss"] = scaled_mtp_loss + mtp_calibrated = mtp_display / len(mtp_loss_ctx_list) * self.config.mtp_config.loss_scaling_factor # type: ignore - split_aux_output = self.aux_loss.finalize( + aux_out = self.aux_loss.finalize( balancing_ctx=balancing_ctx, z_ctx=z_ctx, non_pad_token=non_pad_token, ) - balancing_loss, z_loss, tokens_per_expert_global = split_aux_output - if balancing_loss is not None: - output["balancing_loss"] = balancing_loss - if z_loss is not None: - output["z_loss"] = z_loss - output["tokens_per_expert_global"] = tokens_per_expert_global + if aux_out["balancing_loss"] is not None: + output["balancing_loss"] = aux_out["balancing_loss"] + if aux_out["z_loss"] is not None: + output["z_loss"] = aux_out["z_loss"] + output["tokens_per_expert_global"] = aux_out["tokens_per_expert_global"] + if lm_loss_ctx is not None: + output["calibrated_losses"] = _build_calibrated_losses( + cast(LMHeadLossContext, lm_loss_ctx).calibrate(), balancing_ctx, z_ctx, mtp_calibrated + ) if keep_router: # TODO: Moving router logits to CPU is costly. @@ -1157,6 +1211,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 @@ -1186,10 +1244,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): @@ -1216,8 +1274,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. diff --git a/xtuner/v1/model/moe/qwen3vl_text.py b/xtuner/v1/model/moe/qwen3vl_text.py index fafb70a80e..b088163cf4 100644 --- a/xtuner/v1/model/moe/qwen3vl_text.py +++ b/xtuner/v1/model/moe/qwen3vl_text.py @@ -1,12 +1,14 @@ import os import re +from typing import cast import torch from xtuner.v1.data_proto import SequenceContext +from xtuner.v1.loss import LMHeadLossContext from xtuner.v1.utils.activation_offload import async_save_on_cpu -from .moe import MoELossContextDict, MoEModelOutputs +from .moe import MoELossContextDict, MoEModelOutputs, _build_calibrated_losses from .qwen3 import Qwen3MoE, Qwen3MoE30BA3Config, Qwen3MoE235BA22Config @@ -148,7 +150,7 @@ def _forward( # Hoisted out of the per-layer accumulate path: mask is constant across layers. nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1] non_pad_token = nonpad_indices.numel() - num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device) + num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device) # ===================================================== deepstack_visual_embeds = seq_ctx.deepstack_visual_embeds @@ -196,7 +198,6 @@ def _forward( z_ctx=z_ctx, num_tokens_local=non_pad_token, num_tokens_global=num_tokens_global, - world_size=z_world_size, ) if deepstack_visual_embeds is not None and ((idx := int(idx)) in range(len(deepstack_visual_embeds))): @@ -215,16 +216,20 @@ def _forward( output["logits"] = logits output["extra_info"] = extra_info - balancing_loss, z_loss, tokens_per_expert_global = self.aux_loss.finalize( + aux_out = self.aux_loss.finalize( balancing_ctx=balancing_ctx, z_ctx=z_ctx, non_pad_token=non_pad_token, ) - if balancing_loss is not None: - output["balancing_loss"] = balancing_loss - if z_loss is not None: - output["z_loss"] = z_loss - output["tokens_per_expert_global"] = tokens_per_expert_global + if aux_out["balancing_loss"] is not None: + output["balancing_loss"] = aux_out["balancing_loss"] + if aux_out["z_loss"] is not None: + output["z_loss"] = aux_out["z_loss"] + output["tokens_per_expert_global"] = aux_out["tokens_per_expert_global"] + if lm_loss_ctx is not None: + output["calibrated_losses"] = _build_calibrated_losses( + cast(LMHeadLossContext, lm_loss_ctx).calibrate(), balancing_ctx, z_ctx, None + ) if keep_router: # TODO: Moving router logits to CPU is costly. diff --git a/xtuner/v1/model/utils/misc.py b/xtuner/v1/model/utils/misc.py index 70fbf2d2d2..5d3b50c054 100644 --- a/xtuner/v1/model/utils/misc.py +++ b/xtuner/v1/model/utils/misc.py @@ -63,9 +63,6 @@ class ModelForwardExtraLogInfo(dict): # Tensor to store the maximum model params update ratio. # Shape: `(n_chunk, intra_layer_micro_batch, 1)` if intra_layer_micro_batch > 1 else `(n_chunk, 1)` max_ratio: torch.Tensor - # Tensor to store the ranking loss for logging. - # Shape: `(intra_layer_micro_batch, 1)` if intra_layer_micro_batch > 1 else `(1,)` - local_base_loss: torch.Tensor def __init__(self, init_dict: dict[str, Any] = {}): super().__init__() @@ -101,7 +98,6 @@ def get(self): return_dict = {} # 当增加新的字段时，需要在这里添加相应的处理逻辑 sum_keys = ( - "local_base_loss", "reduced_train_policy_ratio_abs_dev_sum", "reduced_train_policy_clip_low_count", "reduced_train_policy_clip_high_count",