Merge pull request #120 from KempnerInstitute/docs/moe-router-zloss

docs: document moe_router_z_loss_weight (router z-loss)

Author: amazloumi

CHANGELOG.md

@@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- **MoE router z-loss** (`moe_router_z_loss_weight`, ST-MoE style). An optional penalty on the router's pre-softmax logits — per MoE layer `mean_token(logsumexp(router_logits))²` — summed across layers and added to the training loss as `moe_router_z_loss_weight × z_loss`. It keeps router logits from growing without bound, targeting *logit-growth stability* (not load balance — that's the aux loss). Default `0.0` is off: the term is never added, so training, outputs, and gradients are unchanged. `z_loss` is a plain attribute like `aux_loss` (not a buffer/parameter), so it never enters `state_dict` — checkpoint-safe.
+  - `kempnerforge/config/model.py`: `moe_router_z_loss_weight: float = 0.0` (with a non-negativity check).
+  - `kempnerforge/model/router.py`: both `SoftmaxTopKRouter` and `SigmoidTopKRouter` set `self.z_loss = (logsumexp(logits, dim=-1) ** 2).mean()`.
+  - `kempnerforge/model/moe.py` / `transformer.py`: `MoEMLP.z_loss` exposes the per-layer value; `Transformer.get_moe_router_z_loss()` sums it across MoE layers (mirrors `get_moe_aux_loss()`).
+  - `scripts/train.py`: adds `moe_router_z_loss_weight × z_loss` to the loss on both forward paths (gated on `weight > 0`) and logs `moe/router_z_loss`.
+  - Tests: `tests/unit/test_config.py` (default, rejects negative) and `tests/unit/test_moe.py` (z-loss computed; scales with the coefficient).
 - **Fine-grained MoE experts** (`moe_expert_ffn_multiplier`). Decouples each expert's FFN hidden width from the dense FFN: the per-expert hidden dim is `computed_ffn_hidden_dim × moe_expert_ffn_multiplier`, rounded to a multiple of 16. The default `1.0` is a no-op (each expert is a full dense FFN, zero behavior change); set `0.5` for fine-grained experts so top-2 routing matches the dense FFN's activated FLOPs (`2 × F/2 = F`) while adding total capacity — the DeepSeekMoE recipe. Applies to routed and shared experts wherever they are built (`build_moe` and MoMa's `ExpertChoiceMoE`).
   - `kempnerforge/config/model.py`: `moe_expert_ffn_multiplier: float = 1.0` (with a positivity check) and a `computed_expert_ffn_hidden_dim` property; `num_params_estimate` accounts for the smaller experts.
   - `kempnerforge/model/{moe,moma,mot,transformer}.py`: experts are built at `computed_expert_ffn_hidden_dim` instead of the dense FFN width.

README.md

@@ -27,7 +27,7 @@ PyTorch-native framework for fault-tolerant distributed training of foundation m
 
 **Architecture**
 - Decoder-only Transformer: RoPE, GQA, SwiGLU, RMSNorm, `torch.compile`
-- Mixture-of-Experts: softmax top-k and DeepSeek-V3 sigmoid routers, shared experts, fine-grained experts, configurable frequency
+- Mixture-of-Experts: softmax top-k and DeepSeek-V3 sigmoid routers, shared experts, fine-grained experts, router z-loss, configurable frequency
 
 **Parallelism**
 - FSDP2, Tensor, Expert, and Pipeline Parallelism
@@ -212,7 +212,7 @@ Core MoE, Expert Parallelism, DeepSeekMoE, grouped GEMM, FSDP2 compatibility, an
 
 | Feature | Status | Impact |
 |---------|:------:|--------|
-| Router improvements (sequence aux loss, gradient scaling, adaptive bias) | **Done** | Training quality |
+| Router improvements (sequence aux loss, gradient scaling, adaptive bias, z-loss) | **Done** | Training quality |
 | FP8 mixed precision | **Done** | 2x compute throughput |
 | Pipeline parallelism for MoE (PP+EP+TP composition) | Blocked | Required for 100B+ |
 | Communication-computation overlap (async EP dispatch) | Planned | 15-30% throughput |

docs/moe/aux-loss-and-balancing.md

@@ -153,6 +153,41 @@ Leave `moe_aux_loss_weight` at its default 0.01 — it's the
 multiplier on `aux_loss` downstream, so the effective weight on the
 balance loss is `moe_aux_loss_weight × sequence_aux_loss_weight`.
 
+## Router z-loss (logit-growth stability)
+
+Opt-in via `moe_router_z_loss_weight > 0` (default `0.0` = off). The
+ST-MoE *z-loss* penalizes the router's pre-softmax logits so they don't
+grow without bound. It is orthogonal to the balancing losses above — it
+targets *logit-growth stability*, not load balance:
+
+```python
+# kempnerforge/model/router.py — both routers, after computing logits
+self.z_loss = (torch.logsumexp(logits, dim=-1) ** 2).mean()
+```
+
+Both `SoftmaxTopKRouter` and `SigmoidTopKRouter` set it. `MoEMLP.z_loss`
+surfaces the per-layer value, and `Transformer.get_moe_router_z_loss()`
+sums it across MoE layers — mirroring `aux_loss` / `get_moe_aux_loss()`.
+The training loop adds it only when the weight is positive:
+
+```python
+# scripts/train.py
+if mc.moe_router_z_loss_weight > 0:
+    z = inner_transformer(model).get_moe_router_z_loss()
+    loss = loss + mc.moe_router_z_loss_weight * z
+```
+
+Logged as `moe/router_z_loss`. `z_loss` is a plain attribute (like
+`aux_loss`), not a buffer or parameter, so it never enters `state_dict`
+(checkpoint-safe); with the default `0.0` the term is never added, so
+training, outputs, and gradients are unchanged.
+
+**When to use**: long runs where router logits drift upward. A small
+weight such as `1e-3` keeps the summed router z-loss bounded (~2–3, vs.
+~14→24 when off) at identical LM loss — stabilizing the router at
+negligible cost. Distinct from the output-logit `z_loss_weight`
+(a cross-entropy regularizer); this one acts on the *router* logits.
+
 ## Per-expert gradient scaling
 
 Opt-in via `moe_gradient_scale = true`. Normalizes each expert's

docs/moe/index.md

@@ -38,6 +38,7 @@ the distributed mechanics (all-to-all, expert parallelism) live under
 | `moe_sequence_aux_loss_weight` | `0.0` | Sigmoid-only: sequence-level balance penalty |
 | `moe_bias_schedule` | `"constant"` | Sigmoid-only: bias update rate schedule |
 | `moe_expert_ffn_multiplier` | `1.0` | Per-expert FFN hidden = `multiplier × dense FFN`; `0.5` = fine-grained, so top-2 matches the dense FFN's activated FLOPs |
+| `moe_router_z_loss_weight` | `0.0` | `>0` adds the ST-MoE router z-loss `(logsumexp logits)²` to stabilize router logits (see [Aux loss and balancing](aux-loss-and-balancing.md)) |
 
 All fields live in
 [`kempnerforge/config/model.py`](https://github.com/KempnerInstitute/KempnerForge/blob/main/kempnerforge/config/model.py)