From 3ce8fea8d69a9d31f42aca3620ac7efd540ff3bd Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Wed, 22 Jul 2026 09:20:29 +0000 Subject: [PATCH] [FSDP][Loss] Reduce-sum gradient reduction with per-batch aux-loss hub MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch FSDP gradient reduction from mean to pure SUM and remove the compensating ×world_size injections across CE / balancing / z-loss, so each parameter's gradient equals the global loss gradient with no compensating divides. - FSDP reduce-scatter -> SUM (set_gradient_divide_factor(1.0) + set_force_sum_reduction_for_comms(True), avoids the bf16 PreMulSum zeroing); scale_and_reduce_grad drops the expert div_(ep_size) and replicated pre-divide and SUM-reduces Replicate-placement (incl. fp32-ignored) params. - Loss: drop the CE WORLD all_reduce, balancing all_reduce_autograd, z-loss ×world_size, and the now-dead world_size plumbing. - Display: per-rank calibrate() with no cross-rank all_reduce; each rank shows its own per-token-mean loss. - Aux losses folded into a per-batch AuxLossContext hub (MoELossContextDict = {lm, aux, mtp}); AuxLossConfig.build(data, sp_mesh) conforms to the loss-config contract; per-layer accumulate takes router-output lists, applies padding removal / histc / balancing / z internally. - Drop the non-global aux-loss averaging mode (incompatible with reduce-sum). --- tests/model/test_reduce_sum_grad.py | 910 ++++++++++++++++++ xtuner/v1/engine/train_engine.py | 27 +- xtuner/v1/loss/aux_loss.py | 388 +++++--- xtuner/v1/loss/base_loss_ctx.py | 17 +- xtuner/v1/loss/ce_loss.py | 37 +- xtuner/v1/loss/moe_loss.py | 330 ++++--- xtuner/v1/model/base.py | 187 ++-- xtuner/v1/model/compose/base.py | 15 + .../compose/intern_s1/modeling_intern_s1.py | 4 + xtuner/v1/model/dense/dense.py | 6 + xtuner/v1/model/moe/moe.py | 411 ++++---- xtuner/v1/model/moe/qwen3vl_text.py | 55 +- xtuner/v1/model/utils/misc.py | 4 - 13 files changed, 1807 insertions(+), 584 deletions(-) create mode 100644 tests/model/test_reduce_sum_grad.py diff --git a/tests/model/test_reduce_sum_grad.py b/tests/model/test_reduce_sum_grad.py new file mode 100644 index 0000000000..0bfc0c7fec --- /dev/null +++ b/tests/model/test_reduce_sum_grad.py @@ -0,0 +1,910 @@ +from contextlib import contextmanager + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh +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?)" + + # balancing_loss_global is all-reduced inside the loss context, so unlike the per-rank losses + # above it MUST be identical on every rank even with distinct per-rank data. + gvals = [torch.zeros((), device="cuda") for _ in range(world)] + dist.all_gather(gvals, torch.tensor(info["logs_info"]["reduced_balancing_loss_global"], device="cuda")) + if rank == 0: + assert torch.isfinite(torch.stack(gvals)).all(), "non-finite global balancing display loss" + assert (gvals[0] - gvals[1]).abs() < 1e-5, "global balancing loss differs across ranks (not 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() + # Non-padding token count is set once at build time (set_non_pad_token), not threaded per call. + ctx.set_non_pad_token(n_tokens) + for layer in range(n_layers): + ctx.accumulate(layer_idx=layer, 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, + ) + 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() + # Token counts are set once at build time (set_token_counts), not threaded through accumulate(). + ctx.set_token_counts(num_tokens_local, num_tokens_global) + z_new = torch.zeros((), device=device) + for layer in range(n_layers): + z_new = z_new + ctx.accumulate(router_logits=logits_new[layer]) + 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) + + +class TestAuxTokenCountsBuild(_ReduceSumDDPTestCase): + @property + def world_size(self) -> int: + return 4 + + def test_build_sets_sp_local_and_global_token_counts(self): + # The aux token counts (z-loss denominators) are computed in build_loss_ctx_batch by + # sp_split-ing the padding mask and all_reduce-ing the global total, mirroring how CELoss + # builds its denominators -- so the forward runs NO aux-token collective. Under SP the per-rank + # count must be this rank's shard, and the global count must be the whole-sequence non-pad + # total, NOT multiplied by the SP size (the reduce-sum over-count guard). SP mesh == world here + # (no DP), so the global total equals the full sequence length. + self.create_pg("cuda") + device = "cuda" + world = dist.get_world_size() + sp_mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("sp",)) + + # seq_len not a multiple of world -> padded to a multiple, so the last shard is partial: this + # exercises the asymmetric-shard arithmetic, not just an even split. + seq_len = 30 + config = _tiny_moe_config(ep_size=1, balancing=True, z_loss=True) + model = MoE(config=config).to(device) + + torch.manual_seed(0) + # Draw ids in [1, 512): pad_token_id is 0, so every position is a real (non-pad) token and the + # full non-pad count is exactly seq_len. + ids = torch.randint(1, 512, (1, seq_len + 1), device=device) + seq_ctx = SequenceContext.from_input_ids(input_ids=(ids[:, :-1],), device=device) + full_non_pad = int(seq_ctx.mask.sum()) + assert full_non_pad == seq_len + + lcs = model.build_loss_ctx_batch([{"seq_ctx": seq_ctx, "shifted_labels": ids[:, 1:]}], sp_mesh=sp_mesh) + aux_ctx = lcs[0]["aux"] + z_ctx = aux_ctx.z_ctx + balancing_ctx = aux_ctx.balancing_ctx + + local = z_ctx._num_tokens_local + # This rank owns only its shard, so its count is strictly smaller than the whole sequence. + assert local < full_non_pad, f"SP-local count {local} should be a shard, not the full {full_non_pad}" + + # The SP shards tile the whole sequence: summing local counts over the SP group recovers it. + local_sum = torch.tensor(local, device=device) + dist.all_reduce(local_sum, group=sp_mesh.get_group()) + assert int(local_sum) == full_non_pad, f"SP-local counts must sum to {full_non_pad}, got {int(local_sum)}" + + # The build-side global all_reduce must equal the whole-sequence total, not full_non_pad * world + # (which is what feeding the un-split full count into the all_reduce would produce). + assert int(z_ctx._num_tokens_global) == full_non_pad, ( + f"z-loss global token count must be {full_non_pad}, got {int(z_ctx._num_tokens_global)} " + f"(x{world} would mean the SP shards were double-counted)" + ) + # Balancing shares the same SP-local count. + assert balancing_ctx._non_pad_token == local + + +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" + + +class TestDisplayCalibrateCoefficient: + """Single-process checks for the CE display coefficient (no dist / GPU). + + ``build_batches`` divides the backward loss_weight by the GLOBAL denominator, so ``calibrate()`` + multiplies back by ``global_denominator / rank_denominator`` to recover this rank's own mean. In a + single process ``global_denominator == rank_denominator``, so the correct coefficient is exactly + 1.0 for every reduction. The earlier code divided by the raw grad-token count instead of the + per-rank weight sum, which is only equal to the weight sum for ``token`` reduction -- for + ``sample`` / ``square`` it produced a wrong (non-unit) coefficient here. + """ + + def _build_coeff(self, loss_reduction: str) -> float: + from xtuner.v1.loss.ce_loss import CELossConfig, CELossContext, CELossKwargs + + cfg = CELossConfig(loss_reduction=loss_reduction) + # One sequence, two samples of unequal length (2 and 4 grad tokens) so sample / square weights + # differ from a uniform per-token weight. All labels valid (none == ignore_idx). + shifted_labels = torch.arange(6, dtype=torch.long).view(1, 6) + ctx = CELossContext(cfg, CELossKwargs(shifted_labels=shifted_labels)) + cu_seq_lens = torch.tensor([0, 2, 6], dtype=torch.int32) + CELossContext.build_batches([ctx], [cu_seq_lens], None) + return ctx._display_coeff.item() + + def test_token_reduction_coefficient_is_unit(self): + assert abs(self._build_coeff("token") - 1.0) < 1e-6 + + def test_sample_reduction_coefficient_is_unit(self): + # Old formula (global / grad_token_count) gave 2/6 here; the fix (global / weight_sum) gives 1. + assert abs(self._build_coeff("sample") - 1.0) < 1e-6 + + def test_square_reduction_coefficient_is_unit(self): + assert abs(self._build_coeff("square") - 1.0) < 1e-6 + + +class TestAuxLossContextCat: + """Single-process checks that aux contexts collapse to one whole-batch + context (no dist / GPU). + + Under intra_layer_micro_batch the per-micro-batch aux contexts are merged via ``cat`` and fed the + concatenated router statistics once, replacing the earlier fan-out to N identical contexts whose + losses were summed and divided by N. The finalize result must therefore be independent of the + number of merged chunks -- a guard against reintroducing a per-chunk ``/ batch_size`` factor. + """ + + def _finalize_one(self, n_chunks: int) -> tuple[float, float]: + from xtuner.v1.loss.moe_loss import ( + BalancingLossConfig, + BalancingLossContext, + ZLossConfig, + ZLossContext, + ) + + torch.manual_seed(0) + n_experts, top_k, non_pad = 4, 2, 3 + router_weights = torch.rand(non_pad, n_experts) + router_logits = torch.randn(non_pad, n_experts) + tokens_per_expert = torch.tensor([[2, 1, 2, 1]], dtype=torch.long) + + bal = BalancingLossContext.cat([BalancingLossConfig().build() for _ in range(n_chunks)]) + bal.set_non_pad_token(non_pad) + bal.accumulate(layer_idx=0, router_weights=router_weights) + bal_loss = bal.finalize( + tokens_per_expert_local=tokens_per_expert, + tokens_per_expert_global=tokens_per_expert, + n_routed_experts=n_experts, + num_experts_per_tok=top_k, + ) + + z = ZLossContext.cat([ZLossConfig().build() for _ in range(n_chunks)]) + z.set_token_counts(non_pad, None) + z.accumulate(router_logits=router_logits) + z_loss = z.finalize() + return bal_loss.item(), z_loss.item() + + def test_cat_result_independent_of_chunk_count(self): + one = self._finalize_one(1) + many = self._finalize_one(3) + assert abs(one[0] - many[0]) < 1e-6, "balancing loss depends on merged chunk count" + assert abs(one[1] - many[1]) < 1e-6, "z-loss depends on merged chunk count" diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 61a1f58fe0..3d32e7cc31 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -35,7 +35,6 @@ BatchForwardInfo, DataBatchInfo, ModelItem, - ModelOutputs, XTunerBaseModelConfig, ) from xtuner.v1.patch.xtuner_storage import XtunerCacheWriter, _get_async_dcp_save_timeout @@ -51,7 +50,7 @@ class TrainStepInfo(DataBatchInfo, BatchForwardInfo): - total_loss: float + pass logger = get_logger() @@ -216,7 +215,6 @@ 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) for i in range(0, len(data_batches), intra_layer_micro_batch): ProberList.set_micro_batch_iter(micro_batch_iter) @@ -238,14 +236,15 @@ def train_step(self, data_batches: list[ModelItem]) -> TrainStepInfo: micro_batch_results.append(output) - loss = self._get_total_loss(output) - loss.backward() - total_loss += loss.detach() + # The model folds every loss term (LM/policy + aux + MTP) into output.loss; the engine + # just backprops it. Per-term display values ride on output.loss_log. + assert output.loss is not None, "model forward must return a loss in training mode" + output.loss.backward() # 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(**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) @@ -558,17 +557,3 @@ def _maybe_precompute_float8_dynamic_scale_for_fsdp(self): for model in self.model.modules(): if isinstance(model, BaseModel) and model.float8_handler is not None: model.float8_handler.precompute_float8_dynamic_scale_for_fsdp(model) - - def _get_total_loss(self, model_outputs: ModelOutputs) -> torch.Tensor: - # TODO: This logic should be moved into the model layer. The model should be responsible - # for aggregating all losses (CE loss, balancing loss, z loss, etc.) and returning a - # single total_loss. The engine should only call model.forward() and use the returned - # total_loss directly, rather than iterating through fields to sum losses here. - # This would provide better separation of concerns and make the loss computation logic - # more explicit and maintainable. - loss = torch.tensor(0.0, device=DEVICE) - for key in model_outputs.model_fields: - value = getattr(model_outputs, key) - if "loss" in key and isinstance(value, torch.Tensor): - loss += value - return loss diff --git a/xtuner/v1/loss/aux_loss.py b/xtuner/v1/loss/aux_loss.py index 80d1d22264..42774e19a9 100644 --- a/xtuner/v1/loss/aux_loss.py +++ b/xtuner/v1/loss/aux_loss.py @@ -1,10 +1,26 @@ +from typing import cast + import torch import torch.nn as nn from pydantic import BaseModel, ConfigDict from torch import distributed as dist from torch.distributed._functional_collectives import all_reduce +from torch.distributed.device_mesh import DeviceMesh +from typing_extensions import TypedDict + +from xtuner.v1.loss.moe_loss import BalancingLossConfig, BalancingLossContext, ZLossConfig, ZLossContext +from xtuner.v1.loss.utils import sp_split -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): @@ -31,46 +47,116 @@ def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor class AuxLossConfig(BaseModel): - """Configuration for layer-wise split MoE auxiliary loss.""" + """Configuration and factory for the per-micro-batch MoE auxiliary-loss + hub. + + Besides the routing dimensions, it aggregates the optional balancing / z-loss sub-configs so that + :meth:`build` can turn one micro-batch's data into a fully data-bound :class:`AuxLossContext` + (non-padding indices + token-count denominators + sub-contexts) in a single call, mirroring the + ``BaseLossConfig.build(data, sp_mesh)`` contract used by CE / MTP. ``n_routed_experts`` / + ``num_experts_per_tok`` and the sub-configs are populated from the owning model config + (see ``MoEConfig``); aux losses are configured through the model config, not here. + + Args: + n_routed_experts (int | None): Number of routed experts (set from the model config). + num_experts_per_tok (int | None): Experts selected per token (set from the model config). + balancing_loss_cfg (BalancingLossConfig | None): Balancing sub-config, or ``None`` when disabled. + z_loss_cfg (ZLossConfig | None): Z-loss sub-config, or ``None`` when disabled. + """ model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) n_routed_experts: int | None = None num_experts_per_tok: int | None = None + balancing_loss_cfg: BalancingLossConfig | None = None + z_loss_cfg: ZLossConfig | None = None - def build( - self, - *, - n_routed_experts: int | None = None, - num_experts_per_tok: int | None = None, - ) -> "AuxLossContext": - """Build a layer-wise MoE auxiliary loss context.""" - resolved_n_routed_experts = n_routed_experts if n_routed_experts is not None else self.n_routed_experts - assert resolved_n_routed_experts is not None, "n_routed_experts must be provided either in config or build()." - - resolved_num_experts_per_tok = ( - num_experts_per_tok if num_experts_per_tok is not None else self.num_experts_per_tok - ) - assert resolved_num_experts_per_tok is not None, ( - "num_experts_per_tok must be provided either in config or build()." - ) + def build(self, data: dict, sp_mesh: DeviceMesh | None = None) -> "AuxLossContext": + """Build a data-bound per-micro-batch MoE auxiliary-loss hub. - return AuxLossContext( - AuxLossConfig( - n_routed_experts=resolved_n_routed_experts, - num_experts_per_tok=resolved_num_experts_per_tok, - ) - ) + Reads the padding mask from ``data["seq_ctx"]`` (SP-split when a sequence-parallel mesh is + given, mirroring CE) to derive this micro-batch's non-padding indices and the token-count + denominators, builds the optional balancing / z-loss sub-contexts from the aggregated configs, + and returns a hub already carrying all of them -- so the forward runs no aux-token collective. + With no balancing / z sub-configs the result is a tokens-only hub (still produces + ``tokens_per_expert`` for maxvio / bias update, no aux backward, no collective). + Args: + data (dict): Micro-batch data; must contain ``"seq_ctx"`` (a ``SequenceContext``). + sp_mesh (DeviceMesh | None): Sequence-parallel mesh; the mask is SP-split when its size > 1. -class AuxLossContext(nn.Module): - """Layer-wise split MoE auxiliary loss dispatcher. + Returns: + AuxLossContext: The per-micro-batch auxiliary-loss hub bound to this data. + """ + assert self.n_routed_experts is not None, "n_routed_experts must be set on AuxLossConfig before build()." + assert self.num_experts_per_tok is not None, "num_experts_per_tok must be set on AuxLossConfig before build()." + + balancing_ctx = self.balancing_loss_cfg.build() if self.balancing_loss_cfg is not None else None + z_ctx = self.z_loss_cfg.build() if self.z_loss_cfg is not None else None + + mask = data["seq_ctx"].mask + if sp_mesh is not None and sp_mesh.size() > 1: + mask = sp_split(mask, sp_mesh=sp_mesh, split_dim=1, padding_value=False) + + aux_ctx = AuxLossContext(self, balancing_ctx=balancing_ctx, z_ctx=z_ctx) + aux_ctx.set_nonpad_list([torch.nonzero(mask, as_tuple=True)[1]]) + + # Token-count denominators (SP-local count + its cross-rank total) are set here so the forward + # runs no aux collective, mirroring CELossContext.build_batches. Only the sub-contexts consume + # them, so a tokens-only hub skips this entirely. + if balancing_ctx is not None or z_ctx is not None: + num_tokens_local = int(mask.sum()) + if balancing_ctx is not None: + balancing_ctx.set_non_pad_token(num_tokens_local) + if z_ctx is not None: + z_ctx.set_token_counts( + num_tokens_local, self._global_token_count(z_ctx, num_tokens_local, mask.device) + ) + return aux_ctx + + @staticmethod + def _global_token_count( + z_ctx: ZLossContext, + num_tokens_local: int, + device: torch.device | str | int, + ) -> torch.Tensor | None: + # Cross-rank non-padding token count for z-loss global averaging. None when no process group is + # initialized (single-process reference / eval, handled as world size 1). + if not dist.is_initialized(): + return None + n = torch.tensor(num_tokens_local, device=device, dtype=torch.int64) + dist.all_reduce(n, op=dist.ReduceOp.SUM) + return n - Owns the per-layer ``tokens_per_expert`` accumulator used both by logging / bias update and by - ``BalancingLossContext.finalize``. Sub-context accumulators (router_weights_sum for balancing, - logsum / token_count for z-loss) live inside their respective contexts. + +class AuxLossContext(nn.Module): + """Per-batch MoE auxiliary-loss hub. + + Owns everything the model needs to turn a layer's raw router outputs into aux-loss backward and + routing statistics: the optional balancing / z-loss sub-contexts, the per-micro-batch non-padding + indices, and the per-layer ``tokens_per_expert`` accumulator used by logging / bias update and by + ``BalancingLossContext.finalize``. The model feeds one layer's raw router outputs for every + micro-batch at once (as lists); padding removal, histogram accumulation, balancing / z dispatch, + and z-loss injection all happen inside :meth:`accumulate`. + + One hub is built per micro-batch at ``build_loss_ctx_batch`` time, carrying that micro-batch's + non-padding indices and (via the sub-contexts) token counts. Under ``intra_layer_micro_batch`` the + per-micro-batch hubs are merged with :meth:`cat` into a single whole-batch hub whose + ``nonpad_list`` holds the indices in micro-batch order, so :meth:`accumulate` aligns each + micro-batch's nonpad with its router outputs by a single internal ``zip`` (no micro-batch index + is threaded through the model). + + Args: + loss_cfg (AuxLossConfig): Resolved aux config (``n_routed_experts`` / ``num_experts_per_tok``). + balancing_ctx (BalancingLossContext | None): Balancing sub-context, or ``None`` when disabled. + z_ctx (ZLossContext | None): Z-loss sub-context, or ``None`` when disabled. """ - def __init__(self, loss_cfg: AuxLossConfig): + def __init__( + self, + loss_cfg: AuxLossConfig, + balancing_ctx: BalancingLossContext | None = None, + z_ctx: ZLossContext | None = None, + ): super().__init__() self.loss_cfg = loss_cfg n_routed_experts = self.loss_cfg.n_routed_experts @@ -79,115 +165,169 @@ def __init__(self, loss_cfg: AuxLossConfig): assert num_experts_per_tok is not None, "num_experts_per_tok must be resolved before creating AuxLossContext." self.n_routed_experts: int = n_routed_experts self.num_experts_per_tok: int = num_experts_per_tok - self._local_load_logits_list: list[torch.Tensor] = [] + self.balancing_ctx = balancing_ctx + self.z_ctx = z_ctx + # Non-padding indices (SP-local) per micro-batch, in micro-batch order. A per-micro-batch hub + # built at build time holds a single-element list; cat() concatenates them into the whole-batch + # order that accumulate() zips against its router-output lists. + self.nonpad_list: list[torch.Tensor] = [] + # Per-layer tokens_per_expert accumulator keyed by dense MoE-layer ordinal. Same-layer + # micro-batches add into the same slot (histogram counts are additive over the token pool). + self._local_load_logits: dict[int, torch.Tensor] = {} + + @classmethod + def cat(cls, chunks: list["AuxLossContext"]) -> "AuxLossContext": + """Merge per-micro-batch aux hubs into one whole-batch hub. + + Concatenates each chunk's non-padding indices into ``nonpad_list`` (micro-batch order) and + merges the balancing / z-loss sub-contexts (which sum their per-micro-batch token counts). The + differentiable accumulators are empty at merge time -- accumulation runs per layer afterwards. + + Args: + chunks (list[AuxLossContext]): Per-micro-batch hubs (identical config). + + Returns: + AuxLossContext: A single whole-batch hub. + """ + assert len(chunks) > 0, "chunks must not be empty." + balancing_ctx = ( + BalancingLossContext.cat([cast(BalancingLossContext, c.balancing_ctx) for c in chunks]) + if chunks[0].balancing_ctx is not None + else None + ) + z_ctx = ( + ZLossContext.cat([cast(ZLossContext, c.z_ctx) for c in chunks]) if chunks[0].z_ctx is not None else None + ) + merged = cls(chunks[0].loss_cfg, balancing_ctx=balancing_ctx, z_ctx=z_ctx) + merged.set_nonpad_list([nonpad for chunk in chunks for nonpad in chunk.nonpad_list]) + return merged + + def set_nonpad_list(self, nonpad_list: list[torch.Tensor]) -> None: + """Set the per-micro-batch non-padding indices (SP-local), in micro- + batch order. + + Args: + nonpad_list (list[torch.Tensor]): One 1-D index tensor per micro-batch, ordered to match + the router-output lists passed to :meth:`accumulate`. + """ + self.nonpad_list = nonpad_list def accumulate( self, *, - selected_router_weights: torch.Tensor, - selected_router_logits: torch.Tensor, - hidden_states: torch.Tensor, - balancing_ctx: list[BalancingLossContext] | BalancingLossContext | None = None, - 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. + layer_idx: int, + router_weights_list: list[torch.Tensor], + router_logits_list: list[torch.Tensor], + hidden_states_list: list[torch.Tensor], + ) -> list[torch.Tensor]: + """Accumulate one layer's router statistics across all micro-batches + and inject z-loss into the main graph. + + For each micro-batch: strips padding with that micro-batch's nonpad indices, updates the + per-layer ``tokens_per_expert`` and balancing accumulators (all micro-batches of this layer + add into the same slot, so the result equals concatenating them and accumulating once), and + injects this micro-batch's z-loss onto its own ``hidden_states`` carrier via + :class:`AuxLossScaler`. Token counts were set on the sub-contexts at build time. Args: - selected_router_weights (torch.Tensor): Router weights with non-padding tokens already - selected. Shape: ``(non_pad, n_routed_experts)``. - selected_router_logits (torch.Tensor): Router logits with non-padding tokens already - selected. Shape: ``(non_pad, n_routed_experts)``. - hidden_states (torch.Tensor): A carrier tensor on the main forward path. Z-loss is - attached to it via :class:`AuxLossScaler` so that backward through the main loss - releases this layer's logsumexp saved tensor inline. - balancing_ctx (list[BalancingLossContext] | BalancingLossContext | None): Balancing loss - context(s) to fan-out to. ``None`` to skip. - z_ctx (list[ZLossContext] | ZLossContext | None): Z-loss context(s) to fan-out to. - ``None`` to skip. - 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``. + layer_idx (int): Dense MoE-layer ordinal (0-based); all micro-batches share this slot. + router_weights_list (list[torch.Tensor]): Raw router weights per micro-batch, each shape + ``(num_tokens, n_routed_experts)``. + router_logits_list (list[torch.Tensor]): Raw router logits per micro-batch, same shapes. + hidden_states_list (list[torch.Tensor]): Per-micro-batch carrier tensors on the main + forward path. Each micro-batch's z-loss is attached to its own carrier via + :class:`AuxLossScaler` so backward releases this layer's logsumexp saved tensor inline. Returns: - torch.Tensor: ``hidden_states`` augmented with the per-layer z-loss autograd hook. - Identical in value to the input; the caller must replace its handle so the hook is - preserved on the main forward graph. + list[torch.Tensor]: ``hidden_states_list`` with each entry augmented by the per-layer + z-loss autograd hook (identical in value); the caller must replace its handles so the + hooks are preserved on the main forward graph. """ - # tokens_per_expert is non-differentiable (topk + histc) and shared between - # logging output and BalancingLossContext.finalize. Owned here as the single source of truth. - _, selected_experts = torch.topk(selected_router_weights, self.num_experts_per_tok, dim=-1) - tokens_per_expert_l = torch.histc( - selected_experts.view(-1), - bins=self.n_routed_experts, - min=0, - max=self.n_routed_experts, - ).to(torch.long) - self._local_load_logits_list.append(tokens_per_expert_l) - - for ctx in _as_list(balancing_ctx): - ctx.accumulate(router_weights=selected_router_weights) - - for ctx in _as_list(z_ctx): - z_loss_l = ctx.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) - - return hidden_states + assert len(self.nonpad_list) == len(router_weights_list), ( + "nonpad_list and router_weights_list must have one entry per micro-batch." + ) + for i, nonpad_indices in enumerate(self.nonpad_list): + selected_router_weights = router_weights_list[i].index_select(0, nonpad_indices).contiguous().float() + + # tokens_per_expert is non-differentiable (topk + histc) and shared between logging output + # and BalancingLossContext.finalize. Owned here as the single source of truth. + _, selected_experts = torch.topk(selected_router_weights, self.num_experts_per_tok, dim=-1) + tokens_per_expert_l = torch.histc( + selected_experts.view(-1), + bins=self.n_routed_experts, + min=0, + max=self.n_routed_experts, + ).to(torch.long) + prev = self._local_load_logits.get(layer_idx) + self._local_load_logits[layer_idx] = tokens_per_expert_l if prev is None else prev + tokens_per_expert_l + + if self.balancing_ctx is not None: + self.balancing_ctx.accumulate(layer_idx=layer_idx, router_weights=selected_router_weights) + + if self.z_ctx is not None: + selected_router_logits = router_logits_list[i].index_select(0, nonpad_indices).contiguous().float() + z_loss_l = self.z_ctx.accumulate(router_logits=selected_router_logits) + hidden_states_list[i] = AuxLossScaler.apply(hidden_states_list[i], z_loss_l) + + return hidden_states_list + + def finalize(self) -> "AuxLossFinalizeOutput": + """Finalize auxiliary losses and expert counts from the accumulated + state. - def finalize( - self, - *, - 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.""" + 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 + :meth:`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 = [ - ctx.finalize( - tokens_per_expert_local=tokens_per_expert_local, - tokens_per_expert_global=tokens_per_expert_global, - n_routed_experts=self.n_routed_experts, - num_experts_per_tok=self.num_experts_per_tok, - non_pad_token=non_pad_token, - ) - for ctx in balancing_list - ] - balancing_loss = partials[0] if len(partials) == 1 else torch.stack(partials).sum(dim=0) + if self.balancing_ctx is not None: + balancing_loss = self.balancing_ctx.finalize( + tokens_per_expert_local=tokens_per_expert_local, + tokens_per_expert_global=tokens_per_expert_global, + n_routed_experts=self.n_routed_experts, + num_experts_per_tok=self.num_experts_per_tok, + ) 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) + if self.z_ctx is not None: + z_loss = self.z_ctx.finalize() + + return AuxLossFinalizeOutput( + balancing_loss=balancing_loss, + z_loss=z_loss, + tokens_per_expert_global=tokens_per_expert_global, + ) - return balancing_loss, z_loss, tokens_per_expert_global + def calibrate(self) -> dict[str, torch.Tensor]: + """Collect the auxiliary losses' detached display values for logging. + + Owns the aux side of the display pipeline so the model does not hand-assemble it. + ``balancing_loss`` / ``z_loss`` are this rank's own values; ``balancing_loss_global`` is the + cross-rank balancing loss, since expert balance is only meaningful over the whole token pool. + + Returns: + dict[str, torch.Tensor]: Aux display values keyed by loss name (only present terms). + """ + calibrated: dict[str, torch.Tensor] = {} + if self.balancing_ctx is not None: + calibrated["balancing_loss"] = self.balancing_ctx.calibrate() + calibrated["balancing_loss_global"] = self.balancing_ctx.global_calibrate() + if self.z_ctx is not None: + calibrated["z_loss"] = self.z_ctx.calibrate() + return calibrated def _cal_tokens_per_expert(self) -> tuple[torch.Tensor, torch.Tensor]: """Stack per-layer expert counts and produce both local and globally reduced views. - The local view is needed by BalancingLossContext's non-global-average branch (per-rank scaling); the global - view is what the consumer (logging / bias update) wants. + The local view is needed by BalancingLossContext's single-process branch (per-rank scaling); the global view is + what the consumer (logging / bias update) wants. """ - local_load_logits = self._local_load_logits_list - self._local_load_logits_list = [] + local_load_logits = self._local_load_logits + self._local_load_logits = {} if not local_load_logits: raise RuntimeError( @@ -195,7 +335,9 @@ def _cal_tokens_per_expert(self) -> tuple[torch.Tensor, torch.Tensor]: "This usually means the model has no MoE layers or finalize() was called " "without a preceding accumulate()." ) - tokens_per_expert_local = torch.stack(local_load_logits, dim=0) + # Stack in ascending layer-ordinal order; consumers (bias update / maxvio) index rows by the + # dense MoE-layer ordinal, so the row order must follow the ordinal, not insertion order. + tokens_per_expert_local = torch.stack([local_load_logits[k] for k in sorted(local_load_logits)], dim=0) if dist.is_initialized(): group = dist.group.WORLD assert group is not None @@ -203,13 +345,3 @@ def _cal_tokens_per_expert(self) -> tuple[torch.Tensor, torch.Tensor]: else: tokens_per_expert_global = tokens_per_expert_local return tokens_per_expert_local, tokens_per_expert_global - - -def _as_list( - ctx: list | object | None, -) -> list: - if ctx is None: - return [] - if isinstance(ctx, list): - return ctx - return [ctx] 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..63a23f056c 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 / + # rank_denom) 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) @@ -171,17 +175,28 @@ def build_batches( # type: ignore[override] loss_ctx.loss_kwargs.loss_weight = loss_weight loss_weight_list.append(loss_weight) - # Compute the denominator used in the global calibration of the loss + # Denominators for calibration. rank_denominator is this rank's local weight sum; clone before + # the in-place all_reduce so global_denominator does not overwrite it (same tensor otherwise). rank_denominator = sum(loss_weight.sum() for loss_weight in loss_weight_list) rank_denominator = cast(torch.Tensor, rank_denominator) - global_denominator = rank_denominator + global_denominator = rank_denominator.clone() if dist.is_initialized(): dist.all_reduce(global_denominator, op=dist.ReduceOp.SUM) + # Coefficient to turn each rank's backward loss (local weighted sum / global_denominator) back + # into this rank's own reduction-correct mean for display: multiply by global_denominator, + # divide by this rank's own weight sum. The weight sum is the right per-rank denominator for + # every reduction -- token (weight==1, so it equals the grad-token count), sample, and square + # alike -- whereas a raw grad-token count is only correct for token reduction. Summing + # `calibrate()` over the micro-batches yields this rank's mean with no cross-rank all_reduce, + # equal to the global loss when world size is 1 (global_denominator == rank_denominator). + display_coeff = (global_denominator / rank_denominator.clamp_min(1e-12)).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..f2f3ae5e0a 100644 --- a/xtuner/v1/loss/moe_loss.py +++ b/xtuner/v1/loss/moe_loss.py @@ -1,11 +1,10 @@ -from typing import Annotated, Literal +from typing import Annotated, Literal, cast import torch import torch.nn as nn 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" ) @@ -80,43 +58,72 @@ def __init__(self, loss_cfg: BalancingLossConfig, loss_kwargs: BalancingLossKwar super().__init__() self.loss_cfg = loss_cfg self.loss_kwargs = loss_kwargs - self._batch_size = 1 - # 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] = [] + # Per-layer differentiable accumulator keyed by dense MoE-layer ordinal. Under + # intra_layer_micro_batch a layer is accumulated once per micro-batch; the per-micro-batch + # router-weight sums add into the same slot (the sum is linear over the token pool), so the + # result equals concatenating the micro-batches and accumulating once. tokens_per_expert is + # owned by AuxLossContext and passed in at finalize() time to avoid duplicate storage / + # duplicate all_reduce. + self.routing_weights_sum: dict[int, torch.Tensor] = {} + # Detached display values set by finalize(). `_calibrated` is this rank's local balancing loss + # (calibrate()); `_calibrated_global` is the cross-rank balancing loss over the whole token + # pool (global_calibrate()), the semantically meaningful load-balance signal. + self._calibrated: torch.Tensor | None = None + self._calibrated_global: torch.Tensor | None = None + # This rank's non-padding token count, set at build time by set_non_pad_token() + # (see MoE._build_aux_ctx) and summed across micro-batches in cat(). + self._non_pad_token: int = 1 + + def set_non_pad_token(self, non_pad_token: int) -> None: + """Set this rank's non-padding token count for the current batch. + + Args: + non_pad_token (int): Non-padding token count on this rank. + """ + self._non_pad_token = non_pad_token @staticmethod - def build_batches( - loss_ctx_list: list["BalancingLossContext"], - ) -> list["BalancingLossContext"]: - """Build batches for balancing loss contexts. + def cat(chunks: list["BalancingLossContext"]) -> "BalancingLossContext": + """Merge per-micro-batch balancing contexts into one whole-batch + context. - For balancing loss, we set the batch size for proper gradient accumulation. + Balancing loss is a whole-batch quantity (its expert-load statistics are bilinear in the + token pool, so a per-micro-batch sum is not the batch value). Under + ``intra_layer_micro_batch`` the micro-batches are concatenated into one forward and their + router statistics are fed to a single context. The differentiable accumulators are still empty + at merge time -- accumulation runs per layer afterwards -- but the build-time non-padding token + count is per-micro-batch, so it is summed here to recover the whole-batch count. Args: - loss_ctx_list (list[BalancingLossContext]): List of loss contexts. + chunks (list[BalancingLossContext]): Per-micro-batch contexts (identical config). Returns: - list[BalancingLossContext]: The same list with batch_size set. + BalancingLossContext: A single whole-batch context. """ - for loss_ctx in loss_ctx_list: - loss_ctx._batch_size = len(loss_ctx_list) - return loss_ctx_list + assert len(chunks) > 0, "chunks must not be empty." + merged = BalancingLossContext(chunks[0].loss_cfg, BalancingLossKwargs()) + merged._non_pad_token = sum(chunk._non_pad_token for chunk in chunks) + return merged def accumulate( self, *, + layer_idx: int, router_weights: torch.Tensor, ) -> None: """Update the per-layer differentiable accumulator for balancing loss. Args: + layer_idx (int): Dense MoE-layer ordinal (0-based). Same-layer micro-batches add into the + same slot, so the accumulated sum matches concatenating them and accumulating once. router_weights (torch.Tensor): Router weights with non-padding tokens already selected. Shape: ``(non_pad, n_routed_experts)``. """ # router_weights.sum(dim=0) is [n_routed_experts]; sum's backward does not save the input # tensor, so the [non_pad, n_routed_experts] activation is not pinned by this accumulator. - self.routing_weights_sum_list.append(router_weights.sum(dim=0)) + layer_sum = router_weights.sum(dim=0) + prev = self.routing_weights_sum.get(layer_idx) + self.routing_weights_sum[layer_idx] = layer_sum if prev is None else prev + layer_sum def finalize( self, @@ -125,52 +132,124 @@ def finalize( tokens_per_expert_global: torch.Tensor, n_routed_experts: int, num_experts_per_tok: int, - non_pad_token: int, ) -> torch.Tensor: """Finalize balancing loss from accumulators. + The non-padding token count was set at build time via ``set_non_pad_token()``. + Args: tokens_per_expert_local (torch.Tensor): Per-layer expert token counts on this rank, - ``(num_layers, n_routed_experts)``. Used by the non-global-average branch. + ``(num_layers, n_routed_experts)``. Used by the single-process branch. tokens_per_expert_global (torch.Tensor): All-reduced ``tokens_per_expert``, ``(num_layers, n_routed_experts)``. Used by the global-average branch. n_routed_experts (int): Number of routed experts. num_experts_per_tok (int): Number of experts selected per token. - 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) + routing_weights_sum = self.routing_weights_sum + self.routing_weights_sum = {} + if self.loss_cfg.balancing_loss_alpha == 0 or not routing_weights_sum: + zero = torch.tensor(0.0, device=tokens_per_expert_local.device, dtype=torch.float32) + self._calibrated = zero + self._calibrated_global = zero + return zero - local_gating_sum = torch.stack(routing_weights_sum_list, dim=0) + # Stack in ascending layer-ordinal order so the rows align with tokens_per_expert_local / + # tokens_per_expert_global (both keyed by the same dense MoE-layer ordinal). + local_gating_sum = torch.stack([routing_weights_sum[k] for k in sorted(routing_weights_sum)], 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: - valid_tokens = max(non_pad_token, 1) + # Single-process path (no process group). Numerically identical to the distributed branch + # at world size 1: tokens_per_expert_global == tokens_per_expert_local, tokens_global == + # valid_tokens * num_experts_per_tok, and seqlen_global == valid_tokens, so scale / mean / + # loss match term for term. Kept so reference / eval (dist uninitialized) still works. + valid_tokens = max(self._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 + routing_weights_mean = local_gating_sum / valid_tokens + loss_vec = scale_global * (tokens_per_expert_local * routing_weights_mean).sum(-1) + + loss = loss_vec.sum() * alpha + # Display values (detached). `_calibrated`: this rank's balancing loss from LOCAL statistics + # (its own tokens_per_expert / seqlen), a readable per-rank number. `_calibrated_global`: the + # balancing loss over the whole cross-rank token pool (one all_reduce of the gating sum), which + # is the load-balance signal that actually matters -- a per-rank number cannot show whether + # experts are balanced across EP ranks. Both are display-only, NOT backward modes; they do 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 + ) + self._calibrated_global = self._global_balancing_loss( + local_gating_sum, tokens_per_expert_global, n_routed_experts, num_experts_per_tok + ) + 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, + ) -> torch.Tensor: + valid_tokens = max(self._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).detach() - @property - def batch_size(self) -> int: - return self._batch_size + def calibrate(self) -> torch.Tensor: + """This rank's balancing loss (from local statistics) for display + (detached, no all_reduce). + + Returns: + torch.Tensor: This rank's local balancing loss. + """ + assert self._calibrated is not None, "finalize() must be called before calibrate()" + return self._calibrated + + def global_calibrate(self) -> torch.Tensor: + """The balancing loss over the whole cross-rank token pool for display + (detached). + + Reflects whether experts are balanced across all ranks, which the per-rank ``calibrate()`` + value cannot show. Costs one all_reduce of the gating sum, done at ``finalize()`` time. + + Returns: + torch.Tensor: The global balancing loss. + """ + assert self._calibrated_global is not None, "finalize() must be called before global_calibrate()" + return self._calibrated_global + + def _global_balancing_loss( + self, + local_gating_sum: torch.Tensor, + tokens_per_expert_global: torch.Tensor, + n_routed_experts: int, + num_experts_per_tok: int, + ) -> torch.Tensor: + # Mirror finalize()'s global-average branch but on the ALL-REDUCED gating sum, so the value is + # the true global balancing loss rather than this rank's backward component. + global_gating_sum = local_gating_sum.detach().clone() + if dist.is_initialized(): + dist.all_reduce(global_gating_sum, op=dist.ReduceOp.SUM) + tokens_global = tokens_per_expert_global.sum(-1).clamp_min(1) + seqlen_global = (tokens_global // num_experts_per_tok).clamp_min(1) + scale_global = n_routed_experts / tokens_global + routing_weights_mean = global_gating_sum / seqlen_global.unsqueeze(-1) + loss = scale_global * (tokens_per_expert_global * routing_weights_mean).sum(-1) + return (loss.sum() * self.loss_cfg.balancing_loss_alpha).detach() class ZLossConfig(BaseModel): @@ -178,13 +257,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. @@ -214,38 +290,65 @@ def __init__(self, loss_cfg: ZLossConfig, loss_kwargs: ZLossKwargs): super().__init__() 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 + # Token counts set at build time by set_token_counts() (see MoE._build_aux_ctx) and summed + # across micro-batches in cat(). Used as the z-loss normalization denominators. + self._num_tokens_local: int = 1 + self._num_tokens_global: torch.Tensor | None = None @staticmethod - def build_batches( - loss_ctx_list: list["ZLossContext"], - ) -> list["ZLossContext"]: - """Build batches for z-loss contexts. + def cat(chunks: list["ZLossContext"]) -> "ZLossContext": + """Merge per-micro-batch z-loss contexts into one whole-batch context. - For z-loss, we set the batch size for proper gradient accumulation. + Mirrors :meth:`BalancingLossContext.cat`: under ``intra_layer_micro_batch`` the concatenated + router logits are fed to a single context, so the per-micro-batch contexts collapse to one. + The differentiable accumulators are empty at merge time, but the build-time token counts are + per-micro-batch, so both the local count and the (linear) cross-rank count are summed here to + recover the whole-batch denominators. Args: - loss_ctx_list (list[ZLossContext]): List of loss contexts. + chunks (list[ZLossContext]): Per-micro-batch contexts (identical config). Returns: - list[ZLossContext]: The same list with batch_size set. + ZLossContext: A single whole-batch context. + """ + assert len(chunks) > 0, "chunks must not be empty." + merged = ZLossContext(chunks[0].loss_cfg, ZLossKwargs()) + merged._num_tokens_local = sum(chunk._num_tokens_local for chunk in chunks) + # num_tokens_global is None iff no process group is initialized -- identical across chunks + # (shared config). all_reduce is linear, so summing per-chunk globals == all_reduce of the + # summed local counts, i.e. the whole-batch global token total. + if chunks[0]._num_tokens_global is None: + merged._num_tokens_global = None + else: + merged._num_tokens_global = torch.stack( + [cast(torch.Tensor, chunk._num_tokens_global) for chunk in chunks] + ).sum() + return merged + + def set_token_counts(self, num_tokens_local: int, num_tokens_global: torch.Tensor | None) -> None: + """Set this rank's non-padding token counts (local and cross-rank) for + the current batch. + + Args: + num_tokens_local (int): Non-padding token count on this rank. + num_tokens_global (torch.Tensor | None): All-reduced non-padding token count across ranks + (int64 scalar), or ``None`` when no process group is initialized (single-process + reference / eval). """ - for loss_ctx in loss_ctx_list: - loss_ctx._batch_size = len(loss_ctx_list) - return loss_ctx_list + self._num_tokens_local = num_tokens_local + self._num_tokens_global = num_tokens_global def accumulate( self, *, 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. @@ -253,60 +356,65 @@ def accumulate( The caller is expected to inject the returned scalar back into the main forward graph via :class:`AuxLossScaler`; the autograd graph behind this scalar (logsumexp -> square -> sum -> scaling) is then released as soon as the corresponding layer is reached during - backward, instead of being pinned until a global ``finalize()``. + backward, instead of being pinned until a global ``finalize()``. Token counts were set at + build time via ``set_token_counts()``. Args: router_logits (torch.Tensor): Router logits with non-padding tokens already selected. Shape ``(non_pad, n_routed_experts)``. - 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``. Returns: torch.Tensor: Per-layer z-loss as a 0-d tensor with autograd graph back to ``router_logits``. """ + num_tokens_local = self._num_tokens_local + num_tokens_global = self._num_tokens_global if self.loss_cfg.z_loss_alpha == 0: zero = torch.tensor(0.0, device=router_logits.device, dtype=torch.float32) self._update_running(zero) 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 + + # Single-process (dist uninitialized): num_tokens_global is None, so the loss stays `base`, + # which equals the distributed branch at world size 1 (denom_global == num_tokens_local). + loss = base + if num_tokens_global is not None: + # Under reduce-sum gradients 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 this local component + # over ranks reproduces the global z-loss. 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()) + loss = loss * self.loss_cfg.z_loss_alpha + # 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).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: self._running_loss_for_log = value.clone() else: self._running_loss_for_log = self._running_loss_for_log + value - - @property - def batch_size(self) -> int: - return self._batch_size diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index c02c4dc9dd..2a63aeb965 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -85,6 +85,7 @@ class DataBatchInfo(TypedDict): class BatchForwardInfo(TypedDict): + total_loss: float logs_info: dict[str, float] extra_info: ModelForwardExtraLogInfo @@ -402,8 +403,13 @@ class ModelOutputs(PydanticBaseModel): model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") hidden_states: list[torch.Tensor] | None = None logits: torch.Tensor | None = None + # The single total loss the engine backprops (LM/policy + aux + MTP, folded by the model). 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. + loss_log: dict[str, torch.Tensor] | None = None def free_nongrad_feature(self): """Release large intermediate tensors not needed for backward or @@ -576,7 +582,41 @@ 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. + # + # Invariant: SUM over the replicate group is correct only when each rank's grad is that rank's + # OWN local contribution (partition), so summing partitions reconstructs the global grad. + # TODO(reduce-sum): the replicate group includes the SP dim. Under a TPSP scheme where SP is + # realized by an all-gather before attention, the attention block (including any QK norm) runs + # over the full gathered sequence, so a replicated param's grad on each SP rank may be the + # FULL-sequence grad rather than a per-rank partition. Whether this SUM then over-counts such a + # grad by the SP size is an open point that has not been specifically handled here; revisit and + # add a per-param SP-scaling guard if a TPSP path is confirmed to differentiate these params + # over the full sequence. + 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 +625,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 +696,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( @@ -1223,48 +1302,6 @@ def build_loss_ctx_batch( return res - def _add_auxiliary_loss( - self, - loss_name: str, - loss_cfg: Any, - data_batch: list[dict], - res: list[dict], - ) -> None: - """Add auxiliary loss contexts to result. - - This helper builds loss contexts, calibrates them across the batch, - and adds them to the result dictionary. If loss_cfg is None, does nothing. - - Args: - loss_name (str): Name of the loss (e.g., "balancing", "z_loss"). - loss_cfg (Any): Loss configuration with a build() method. If None, skipped. - data_batch (list[dict]): Batch data. - res (list[dict]): Result dictionary to populate. Modified in-place. - - Example: - def build_loss_ctx_batch(self, data_batch, sp_mesh): - res = super().build_loss_ctx_batch(data_batch, sp_mesh) - - # One line per auxiliary loss - self._add_auxiliary_loss("balancing", self.config.balancing_loss_cfg, data_batch, res) - self._add_auxiliary_loss("z_loss", self.config.z_loss_cfg, data_batch, res) - - return res - """ - if loss_cfg is None: - return - - # Build loss contexts for all microbatches - ctx_list = [loss_cfg.build() for _ in data_batch] - - # Calibrate across batch - ctx_cls = ctx_list[0].__class__ - ctx_list = ctx_cls.build_batches(ctx_list) - - # Add to result - for i, ctx in enumerate(ctx_list): - res[i][loss_name] = ctx # type: ignore - def pre_micro_batch_forward(self, data_batches: Sequence[ModelItem]) -> DataBatchInfo: step_consumed_tokens = torch.tensor(0, device=DEVICE) step_consumed_img_tokens = torch.tensor(0.0, device=DEVICE) @@ -1305,41 +1342,41 @@ 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) + # This rank's displayed total loss: sum of the per-rank loss terms. `*_global` keys are + # cross-rank monitoring curves (e.g. balancing_loss_global), NOT additive loss terms, so they + # are excluded to avoid double-counting the term they mirror. + total_loss = sum(v for k, v in logs_info.items() if not k.endswith("_global")) 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_copy: - extra_info = output["extra_info"] - train_engine_extra_info.append(extra_info) - - 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() + if "extra_info" in output: + train_engine_extra_info.append(output["extra_info"]) + return BatchForwardInfo(total_loss=total_loss, logs_info=logs_info, extra_info=train_engine_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. + + Naming note: "reduce" / the ``reduced_`` key prefix here means reduced ACROSS THE STEP'S + MICRO-BATCHES on this rank, following the project-wide ``reduced_`` logging convention. + It is NOT a cross-rank reduction: 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), + stored on ``output.loss_log``, and here we only sum those over the micro-batches. So + every rank logs ITS OWN loss (equal to the global loss at world size 1). The one cross-rank + term is ``balancing_loss_global`` (all-reduced inside the loss context, identical on every + rank). 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.loss_log 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..8a8a4b2e61 100644 --- a/xtuner/v1/model/dense/dense.py +++ b/xtuner/v1/model/dense/dense.py @@ -119,6 +119,8 @@ def forward( output["loss"] = loss output["logits"] = logits output["extra_info"] = extra_info + # Dense has no aux loss, so the total loss is the LM loss and loss_log carries only it. + output["loss_log"] = {"llm_loss": cast(CELossContext, loss_ctx["lm"]).calibrate().detach()} return ModelOutputs(**output) @@ -300,6 +302,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..a8f9539ced 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -8,9 +8,8 @@ import torch.distributed as dist import torch.nn.functional as F from cyclopts import Parameter -from pydantic import ConfigDict +from pydantic import ConfigDict, model_validator from torch import nn -from torch.distributed._functional_collectives import all_reduce from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl from torch.distributed.device_mesh import DeviceMesh, init_device_mesh from torch.distributed.distributed_c10d import ReduceOp @@ -29,12 +28,10 @@ AuxLossConfig, AuxLossContext, BalancingLossConfig, - BalancingLossContext, BaseLossContext, LMHeadLossContext, MTPLossContext, ZLossConfig, - ZLossContext, ) from xtuner.v1.loss.mtp_loss import MTPLossConfig from xtuner.v1.model.base import ( @@ -96,13 +93,66 @@ MOE_EP_COMPILE_CFG.pop("xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward") +def _total_backward_loss( + llm_loss: torch.Tensor, + balancing_loss: torch.Tensor | None, + mtp_loss: torch.Tensor | None, +) -> torch.Tensor: + """Sum the graph-carrying loss terms into the single scalar the engine + backprops. + + Z-loss is not added here: it is injected into the main graph via ``AuxLossScaler`` at + accumulate() time, so its gradient already flows through ``llm_loss``. Only the terms whose + backward graph is NOT already on the main path (balancing, MTP) are summed in. + + Args: + llm_loss (torch.Tensor): The LM (or RL policy) loss. + balancing_loss (torch.Tensor | None): Balancing loss carrying its graph, or ``None``. + mtp_loss (torch.Tensor | None): MTP loss carrying its graph, or ``None``. + + Returns: + torch.Tensor: The total backward loss. + """ + total = llm_loss + if balancing_loss is not None: + total = total + balancing_loss + if mtp_loss is not None: + total = total + mtp_loss + return total + + +def _build_loss_log( + lm_display: torch.Tensor, + aux_calibrated: dict[str, torch.Tensor], + mtp_display: torch.Tensor | None, +) -> dict[str, torch.Tensor]: + """Assemble the per-loss display dict from the LM, aux, and MTP terms. + + Stored on ``output.loss_log`` and aggregated over micro-batches by the display pipeline. The aux + terms (``balancing_loss`` / ``balancing_loss_global`` / ``z_loss``) are produced by + :meth:`AuxLossContext.calibrate`; this only prepends ``llm_loss`` and appends ``mtp_loss``. + + Args: + lm_display (torch.Tensor): This rank's LM loss display value. + aux_calibrated (dict[str, torch.Tensor]): Aux display values from ``AuxLossContext.calibrate``. + mtp_display (torch.Tensor | None): MTP loss display value, or ``None``. + + Returns: + dict[str, torch.Tensor]: Display values keyed by loss name. + """ + loss_log: dict[str, torch.Tensor] = {"llm_loss": lm_display.detach()} + loss_log |= aux_calibrated + if mtp_display is not None: + loss_log["mtp_loss"] = mtp_display.detach() + return loss_log + + class MoEModelOutputs(ModelOutputs): router_logits: dict[str, torch.Tensor] | None = None router_weights: dict[str, torch.Tensor] | None = None - balancing_loss: torch.Tensor | None = None - z_loss: torch.Tensor | None = None + # Aux losses (balancing / z / MTP) are folded into `loss` for backward and into `loss_log` for + # display; only the routing statistics needed by metrics / bias update are surfaced separately. tokens_per_expert_global: torch.Tensor - mtp_loss: torch.Tensor | None = None def free_nongrad_feature(self): """Release large intermediate tensors not needed for backward or @@ -127,8 +177,7 @@ class MoEBatchForwardInfo(BatchForwardInfo): class MoELossContextDict(TypedDict): lm: BaseLossContext - balancing: BalancingLossContext | None - z_loss: ZLossContext | None + aux: AuxLossContext | None mtp: list[BaseLossContext] | None @@ -161,6 +210,22 @@ class MoEConfig(TransformerConfig): # keeping it unsharded after forward to avoid repeated all-gathers. embed_reshard_after_forward: bool = True + @model_validator(mode="after") + def _populate_aux_loss_cfg(self) -> "MoEConfig": + # aux_loss_cfg is the internal aggregate AuxLossConfig.build(data, sp_mesh) reads: mirror the + # routing dims and the user-facing balancing / z-loss switches into it so build() is + # self-contained. Users configure aux losses via balancing_loss_cfg / z_loss_cfg on the model + # config; this keeps aux_loss_cfg the single, in-sync source build() consumes. + self.aux_loss_cfg = self.aux_loss_cfg.model_copy( + update={ + "n_routed_experts": self.n_routed_experts, + "num_experts_per_tok": self.num_experts_per_tok, + "balancing_loss_cfg": self.balancing_loss_cfg, + "z_loss_cfg": self.z_loss_cfg, + } + ) + return self + def build(self) -> "MoE": from xtuner.v1.model.moe.moe import MoE @@ -212,66 +277,32 @@ def __init__(self, config: MoEConfig): self._maybe_enable_compile(self.compile_cfg) self.offload_stream = torch.cuda.Stream() - self.aux_loss: AuxLossContext = self.config.aux_loss_cfg.build( - n_routed_experts=self.config.n_routed_experts, - num_experts_per_tok=self.config.num_experts_per_tok, - ) def _maybe_offload_router(self, tensor: torch.Tensor) -> torch.Tensor: if self.config.router_async_offload: return async_offload_to_cpu(tensor, self.offload_stream) return tensor - def _z_loss_dist_token_count( - self, - z_ctx: list[ZLossContext] | ZLossContext | None, - num_tokens_local: int, - device: torch.device | str | int, - ) -> tuple[torch.Tensor | None, int]: - """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. - """ - 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 - 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() - - def _extract_aux_loss_ctx( + def _resolve_aux_ctx( self, - loss_ctx: list[MoELossContextDict] | MoELossContextDict | None, - ) -> tuple[ - list[BalancingLossContext] | BalancingLossContext | None, - list[ZLossContext] | ZLossContext | None, - ]: - if loss_ctx is None: - return None, None - - if isinstance(loss_ctx, list): - balancing_ctx: list[BalancingLossContext] = [] - z_ctx: list[ZLossContext] = [] - for ctx in loss_ctx: - ctx_bal = ctx.get("balancing") - if ctx_bal is not None: - balancing_ctx.append(ctx_bal) - ctx_z = ctx.get("z_loss") - if ctx_z is not None: - z_ctx.append(ctx_z) - # Collapse empty fan-out lists to None so downstream guards - # (`if ctx is None`, `_z_loss_dist_token_count`, AuxLoss.accumulate fan-out) - # can treat "no context across any micro-batch" as the no-op case. - return (balancing_ctx or None), (z_ctx or None) - - return loss_ctx.get("balancing"), loss_ctx.get("z_loss") + loss_ctx: list[MoELossContextDict | None] | MoELossContextDict | None, + seq_ctx: list[SequenceContext] | SequenceContext, + ) -> AuxLossContext: + # The whole-batch aux hub for this forward. Aux stats are whole-batch quantities: under + # intra_layer_micro_batch the per-micro-batch hubs built at build_loss_ctx_batch time are + # merged via cat(), which concatenates their nonpad indices into nonpad_list. + loss_ctx_list = loss_ctx if isinstance(loss_ctx, list) else [loss_ctx] + aux_chunks = [lc.get("aux") if lc is not None else None for lc in loss_ctx_list] + if aux_chunks and all(aux is not None for aux in aux_chunks): + return AuxLossContext.cat([cast(AuxLossContext, aux) for aux in aux_chunks]) + # Fallback: the loss context carries no aux hub (forward_only / reference model, which pass a + # bare {"lm": ...} or None). Build a tokens-only hub -- balancing / z sub-configs stripped, so + # no aux backward and no aux collective -- so tokens_per_expert is still produced for maxvio / + # bias update. The forward seq_ctx mask is already SP-local, so build runs with no sp_split. + seq_ctx_list = seq_ctx if isinstance(seq_ctx, list) else [seq_ctx] + tokens_only_cfg = self.config.aux_loss_cfg.model_copy(update={"balancing_loss_cfg": None, "z_loss_cfg": None}) + chunks = [tokens_only_cfg.build({"seq_ctx": sc}) for sc in seq_ctx_list] + return chunks[0] if len(chunks) == 1 else AuxLossContext.cat(chunks) @torch.no_grad() def update_bias(self, total_expert_counts_pre_iter, expected_loads): @@ -317,24 +348,18 @@ def build_loss_ctx_batch( # type: ignore[override] list[dict]: Loss context dict for each microbatch. Each dict contains: - "lm": LM loss context - - "balancing": Balancing loss context (if configured) - - "z_loss": Z-loss context (if configured) + - "aux": Per-micro-batch auxiliary-loss hub (routing stats + optional balancing / z) - "mtp": MTP loss contexts (if configured) - - Note: - Auxiliary loss contexts are built without parameters. - All data is passed to forward() at runtime: - - balancing_ctx(router_weights, n_routed_experts, num_experts_per_tok) - - z_loss_ctx(router_logits) """ # Build LM loss context _data_batch: list[dict] = data_batch # type: ignore[assignment] res: list[dict] = super().build_loss_ctx_batch(_data_batch, sp_mesh) cu_seq_lens_list = [data["seq_ctx"].cu_seq_lens_k for data in data_batch] - # Add auxiliary losses - self._add_auxiliary_loss("balancing", self.config.balancing_loss_cfg, _data_batch, res) - self._add_auxiliary_loss("z_loss", self.config.z_loss_cfg, _data_batch, res) + # Build the per-micro-batch aux hub (routing stats + optional balancing / z sub-contexts), + # data-bound like the lm / mtp contexts: one AuxLossConfig.build(data, sp_mesh) per micro-batch. + for loss_ctx_dict, data in zip(res, _data_batch): + loss_ctx_dict["aux"] = self.config.aux_loss_cfg.build(data, sp_mesh) # Add MTP loss contexts if MTP is enabled if self.config.mtp_config is not None: @@ -450,12 +475,6 @@ def _micro_batch_forward( cat_position_embeddings[1].chunk(len(seq_ctx_list), dim=1), ) ) - cat_mask = torch.cat([ctx.mask for ctx in seq_ctx_list], dim=1) - # Hoisted out of the per-layer accumulate path: mask is constant across layers, - # so the non-pad index lookup runs once per forward instead of once per (layer, ctx). - nonpad_indices = torch.nonzero(cat_mask, as_tuple=True)[1] - non_pad_token = nonpad_indices.numel() - # Initialize output containers output: dict = {} @@ -465,14 +484,19 @@ def _micro_batch_forward( router_logits_list: list[dict[str, torch.Tensor]] = ( [{} 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) + # Whole-batch aux hub: padding removal / stat accumulation / z-loss injection all happen + # inside it; the model feeds one layer's raw router outputs for every micro-batch as lists and + # the hub zips them against its own per-micro-batch nonpad_list (built at build_loss_ctx_batch). + aux_ctx = self._resolve_aux_ctx(loss_ctx_list, seq_ctx_list) # type: ignore[arg-type] # Process through layers cat_seq_ctx: SequenceContext | None = None hidden_states_list: list[torch.Tensor] = [] moe_forward = False + # Dense 0-based ordinal of the MoE layers actually processed (immune to PP layer subsets and + # to the absolute dict index); it keys the aux hub's per-layer slots and continues into MTP. + moe_layer_ordinal = 0 for seq_ctx in seq_ctx_list: self._mark_dynamic(seq_ctx) @@ -523,36 +547,32 @@ def _micro_batch_forward( position_embeddings=position_embeddings_list, seq_ctx=seq_ctx_list, ) - hidden_states = layer_results[: len(hidden_states_list)] + layer_hidden = layer_results[: len(hidden_states_list)] router_logits = layer_results[len(hidden_states_list) : len(hidden_states_list) * 2] router_weights = layer_results[len(hidden_states_list) * 2 :] - # Update hidden states and (optionally) collect router logits. - # router_weights are only consumed by aux_loss.accumulate below, so we - # never stash them per-MB the way we do for logits. - for i, hidden_states in enumerate(hidden_states): - hidden_states_list[i] = hidden_states + # Per micro-batch: stash router logits if requested, then hand this layer's raw router + # outputs to the aux hub. The hub strips padding with the micro-batch's own nonpad + # indices and adds the stats into this layer's slot -- summing over micro-batches is + # equal to concatenating them and accumulating once. The z-loss rides on each + # micro-batch's own carrier; all carriers converge into the same total_loss backward, + # so each aux node fires exactly once. + for i in range(len(layer_hidden)): + hidden_states_list[i] = layer_hidden[i] if keep_router: router_logits_list[i][f"layer{idx}"] = self._maybe_offload_router(router_logits[i]) - - cat_router_weights = torch.cat(router_weights, dim=0) - cat_router_logits = torch.cat(router_logits, dim=0) - # Pin the per-layer z-loss to MB0's hidden_states stream. With multiple MBs, only - # one carrier may be chosen — all MBs converge into the same total_loss backward, - # so MB0's path traverses every aux-loss node exactly once. - hidden_states_list[0] = self.aux_loss.accumulate( - selected_router_weights=cat_router_weights.index_select(0, nonpad_indices).contiguous().float(), - selected_router_logits=cat_router_logits.index_select(0, nonpad_indices).contiguous().float(), - hidden_states=hidden_states_list[0], - balancing_ctx=balancing_ctx, - z_ctx=z_ctx, - num_tokens_local=non_pad_token, - num_tokens_global=num_tokens_global, - world_size=z_world_size, + hidden_states_list = aux_ctx.accumulate( + layer_idx=moe_layer_ordinal, + router_weights_list=list(router_weights), + router_logits_list=list(router_logits), + hidden_states_list=hidden_states_list, ) + moe_layer_ordinal += 1 assert hidden_states_list, "XTuner Internal Error, found empty hidden states for domino EP" + mtp_loss: torch.Tensor | None = None + mtp_calibrated: torch.Tensor | None = None if self.mtp_block is not None: assert self.config.mtp_config is not None @@ -579,6 +599,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,54 +608,41 @@ 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: # MTP routed experts feed the same balancing / z aux loss as the main MoE layers # (mirrors the single-microbatch path in `_forward`); without this they are silently - # excluded from the aux loss and from `tokens_per_expert` in this path. Unlike the LM - # loss above, balancing cannot be summed per micro-batch: it must combine all micro- - # batches into one row per MTP depth, because `finalize` multiplies tokens_per_expert - # by the router-weight mean per row and that product over a combined token pool - # differs from the sum of per-microbatch products (it would also drift with - # intra_layer_micro_batch, a perf knob). So we concatenate the per-microbatch router - # results per depth and accumulate once, exactly like the main layer loop above does - # per layer. MTP shares the main micro-batch masks (mtp_seq_ctx_list is copied from - # seq_ctx_list), so the main nonpad indices and token counts apply directly. The - # z-loss carrier is hidden_states_list[0], the same main-loss path the per-layer aux - # loss already rides on, so backward traverses each MTP aux node exactly once. + # excluded from the aux loss and from `tokens_per_expert` in this path. Each MTP depth + # is one aux slot, accumulated per micro-batch into that slot (summing over micro- + # batches equals concatenating them and accumulating once). MTP shares the main micro- + # batch masks (mtp_seq_ctx_list is copied from seq_ctx_list), so the same per-mb nonpad + # indices apply. The depth slots continue the dense MoE-layer ordinal. The z-loss rides + # on each micro-batch's own carrier -- the same main-loss path the per-layer aux loss + # already rides on -- so backward traverses each MTP aux node exactly once. + n_moe_layers = moe_layer_ordinal for mtp_idx in range(self.config.mtp_config.num_layers): - cat_mtp_router_weights = torch.cat( - [mb_outputs[mtp_idx][2] for mb_outputs in mtp_outputs_per_mb], dim=0 - ) - cat_mtp_router_logits = torch.cat( - [mb_outputs[mtp_idx][1] for mb_outputs in mtp_outputs_per_mb], dim=0 - ) - hidden_states_list[0] = self.aux_loss.accumulate( - selected_router_weights=cat_mtp_router_weights.index_select(0, nonpad_indices) - .contiguous() - .float(), - selected_router_logits=cat_mtp_router_logits.index_select(0, nonpad_indices) - .contiguous() - .float(), - hidden_states=hidden_states_list[0], - balancing_ctx=balancing_ctx, - z_ctx=z_ctx, - num_tokens_local=non_pad_token, - num_tokens_global=num_tokens_global, - world_size=z_world_size, + hidden_states_list = aux_ctx.accumulate( + layer_idx=n_moe_layers + mtp_idx, + router_weights_list=[mb_outputs[mtp_idx][2] for mb_outputs in mtp_outputs_per_mb], + router_logits_list=[mb_outputs[mtp_idx][1] for mb_outputs in mtp_outputs_per_mb], + hidden_states_list=hidden_states_list, ) - output["mtp_loss"] = mtp_losses * self.config.mtp_config.loss_scaling_factor + 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,26 +652,25 @@ 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) - output["loss"] = loss.sum() moe_extra_info = ModelForwardExtraLogInfo() if extra_info: moe_extra_info.append(extra_info) output["extra_info"] = moe_extra_info - split_aux_output = self.aux_loss.finalize( - balancing_ctx=balancing_ctx, - z_ctx=z_ctx, - non_pad_token=non_pad_token, + aux_out = aux_ctx.finalize() + # output.loss is the single total the engine backprops; per-term values are display-only. + output["loss"] = _total_backward_loss(loss.sum(), aux_out["balancing_loss"], mtp_loss) + output["tokens_per_expert_global"] = aux_out["tokens_per_expert_global"] + output["loss_log"] = _build_loss_log( + cast(LMHeadLossContext, cat_loss_ctx).calibrate(), + aux_ctx.calibrate(), + mtp_calibrated, ) - 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 keep_router: # TODO: Returning router logits is costly. @@ -720,11 +728,13 @@ def _forward( output["router_logits"] = None output["router_weights"] = None self._mark_dynamic(seq_ctx) - balancing_ctx, z_ctx = self._extract_aux_loss_ctx(loss_ctx) - # 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) + # Whole-batch aux hub (single micro-batch here): padding removal / stat accumulation / z-loss + # injection all happen inside it; the model feeds each layer's raw router outputs as + # single-element lists and the hub zips them against its own nonpad_list. + aux_ctx = self._resolve_aux_ctx(loss_ctx, seq_ctx) + # Dense 0-based ordinal of the MoE layers actually processed; keys the aux hub's per-layer + # slots and continues into MTP. Immune to PP layer subsets / the absolute dict index. + moe_layer_ordinal = 0 for idx, decoder_layer in self.layers.items(): if int(idx) < self.config.first_k_dense_replace: @@ -758,16 +768,13 @@ def _forward( if keep_router: output["router_logits"][f"layer{idx}"] = self._maybe_offload_router(router_results) output["router_weights"][f"layer{idx}"] = self._maybe_offload_router(router_weights) - hidden_states = self.aux_loss.accumulate( - selected_router_weights=router_weights.index_select(0, nonpad_indices).contiguous().float(), - selected_router_logits=router_results.index_select(0, nonpad_indices).contiguous().float(), - hidden_states=hidden_states, - balancing_ctx=balancing_ctx, - z_ctx=z_ctx, - num_tokens_local=non_pad_token, - num_tokens_global=num_tokens_global, - world_size=z_world_size, - ) + hidden_states = aux_ctx.accumulate( + layer_idx=moe_layer_ordinal, + router_weights_list=[router_weights], + router_logits_list=[router_results], + hidden_states_list=[hidden_states], + )[0] + moe_layer_ordinal += 1 if self.config.return_hidden_states: output["hidden_states"].append(hidden_states) @@ -778,11 +785,12 @@ def _forward( # Get LM loss context from dict lm_loss_ctx = loss_ctx["lm"] if loss_ctx is not None else None loss, (logits, extra_info) = self.lm_head(hidden_states, lm_loss_ctx) # type: ignore - output["loss"] = loss output["logits"] = logits output["extra_info"] = extra_info # MTP forward pass and loss computation + mtp_loss_total: torch.Tensor | None = None + mtp_calibrated: torch.Tensor | None = None if ( self.mtp_block is not None and loss_ctx is not None @@ -793,12 +801,9 @@ def _forward( position_ids=position_ids.clone(), inputs_embeds=seq_ctx.inputs_embeds.clone() if seq_ctx.inputs_embeds is not None else None, ) - # 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 shares the main padding mask (mtp_seq_ctx is a copy that does not override the + # cu_seq_lens / num_padding the mask derives from), so the main non-pad indices and the + # token counts already set on the aux contexts apply directly. # Forward through MTP block mtp_outputs = self.mtp_block( @@ -810,47 +815,47 @@ def _forward( # Compute MTP losses for each depth mtp_losses = torch.tensor(0.0, device=DEVICE) - for idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)): + # Per-rank display value (from each depth's calibrate()), mirroring `mtp_losses`. + mtp_display = torch.tensor(0.0, device=DEVICE) + # MTP depth slots continue the dense MoE-layer ordinal; MTP shares the main mask, so the + # main nonpad indices apply directly. + n_moe_layers = moe_layer_ordinal + for depth, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)): mtp_hidden_states, mtp_router_results, mtp_router_weights = mtp_hidden if keep_router: - output["router_logits"][f"mtp_layer{idx}"] = mtp_router_results - output["router_weights"][f"mtp_layer{idx}"] = mtp_router_weights + output["router_logits"][f"mtp_layer{depth}"] = mtp_router_results + output["router_weights"][f"mtp_layer{depth}"] = mtp_router_weights # Inject this MTP layer's z-loss before lm_head so backward through mtp_loss # traverses the AuxLossScaler node and releases this layer's logsumexp activations. - mtp_hidden_states = self.aux_loss.accumulate( - selected_router_weights=mtp_router_weights.index_select(0, mtp_nonpad_indices) - .contiguous() - .float(), - selected_router_logits=mtp_router_results.index_select(0, mtp_nonpad_indices).contiguous().float(), - hidden_states=mtp_hidden_states, - balancing_ctx=balancing_ctx, - 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_hidden_states = aux_ctx.accumulate( + layer_idx=n_moe_layers + depth, + router_weights_list=[mtp_router_weights], + router_logits_list=[mtp_router_results], + hidden_states_list=[mtp_hidden_states], + )[0] 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) scaled_mtp_loss = mtp_losses * self.config.mtp_config.loss_scaling_factor # type: ignore - # Add to total loss - output["mtp_loss"] = scaled_mtp_loss - - split_aux_output = 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 + mtp_loss_total = scaled_mtp_loss + mtp_calibrated = mtp_display / len(mtp_loss_ctx_list) * self.config.mtp_config.loss_scaling_factor # type: ignore + + aux_out = aux_ctx.finalize() + output["tokens_per_expert_global"] = aux_out["tokens_per_expert_global"] + if lm_loss_ctx is not None: + assert loss is not None + # output.loss is the single total the engine backprops; per-term values are display-only. + output["loss"] = _total_backward_loss(loss, aux_out["balancing_loss"], mtp_loss_total) + output["loss_log"] = _build_loss_log( + cast(LMHeadLossContext, lm_loss_ctx).calibrate(), + aux_ctx.calibrate(), + mtp_calibrated, + ) if keep_router: # TODO: Moving router logits to CPU is costly. @@ -1157,6 +1162,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 +1195,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 +1225,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..bb73772118 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_loss_log, _total_backward_loss from .qwen3 import Qwen3MoE, Qwen3MoE30BA3Config, Qwen3MoE235BA22Config @@ -144,11 +146,13 @@ def _forward( output["router_weights"] = None self._mark_dynamic(seq_ctx) - balancing_ctx, z_ctx = self._extract_aux_loss_ctx(loss_ctx) - # 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) + # Whole-batch aux hub (single micro-batch here): padding removal / stat accumulation / z-loss + # injection all happen inside it; the model feeds each layer's raw router outputs as + # single-element lists and the hub zips them against its own nonpad_list. + aux_ctx = self._resolve_aux_ctx(loss_ctx, seq_ctx) + # Dense 0-based ordinal of the MoE layers actually processed; keys the aux hub's per-layer + # slots. Immune to PP layer subsets / the absolute dict index. + moe_layer_ordinal = 0 # ===================================================== deepstack_visual_embeds = seq_ctx.deepstack_visual_embeds @@ -188,16 +192,13 @@ def _forward( if keep_router: output["router_logits"][f"layer{idx}"] = router_results output["router_weights"][f"layer{idx}"] = router_weights - hidden_states = self.aux_loss.accumulate( - selected_router_weights=router_weights.index_select(0, nonpad_indices).contiguous().float(), - selected_router_logits=router_results.index_select(0, nonpad_indices).contiguous().float(), - hidden_states=hidden_states, - balancing_ctx=balancing_ctx, - z_ctx=z_ctx, - num_tokens_local=non_pad_token, - num_tokens_global=num_tokens_global, - world_size=z_world_size, - ) + hidden_states = aux_ctx.accumulate( + layer_idx=moe_layer_ordinal, + router_weights_list=[router_weights], + router_logits_list=[router_results], + hidden_states_list=[hidden_states], + )[0] + moe_layer_ordinal += 1 if deepstack_visual_embeds is not None and ((idx := int(idx)) in range(len(deepstack_visual_embeds))): assert visual_pos_masks is not None @@ -211,20 +212,20 @@ def _forward( # Get LM loss context from dict lm_loss_ctx = loss_ctx["lm"] if loss_ctx is not None else None loss, (logits, extra_info) = self.lm_head(hidden_states, lm_loss_ctx) # type: ignore - output["loss"] = loss output["logits"] = logits output["extra_info"] = extra_info - balancing_loss, z_loss, tokens_per_expert_global = 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 + aux_out = aux_ctx.finalize() + output["tokens_per_expert_global"] = aux_out["tokens_per_expert_global"] + if lm_loss_ctx is not None: + assert loss is not None + # output.loss is the single total the engine backprops; per-term values are display-only. + output["loss"] = _total_backward_loss(loss, aux_out["balancing_loss"], None) + output["loss_log"] = _build_loss_log( + cast(LMHeadLossContext, lm_loss_ctx).calibrate(), + aux_ctx.calibrate(), + 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",