Merge 5422d3a132 into c7ba252142

dev/LOG.md

@@ -749,7 +749,7 @@ See the branch `fp8_attempt_fail` for:
 ### Open Questions
 
 - Why does the custom op approach use more memory than vanilla BF16?
-- Why is the bump in tok_per_sec so low? We should see ~1.6X speedup in both the forward pass and also (twice) in backward pass for the gradients. Granted, Ahmdal's law is part of the solution because our vocab_size is only 32K so the final layer isn't a huge part of the profile but the expected speedup is still not fully realized.
+- Why is the bump in tok_per_sec so low? We should see ~1.6X speedup in both the forward pass and also (twice) in backward pass for the gradients. Granted, Amdahl's law is part of the solution because our vocab_size is only 32K so the final layer isn't a huge part of the profile but the expected speedup is still not fully realized.
 
 **Conclusion:** Negative result for now. The implementation works correctly but provides marginal speedup with *increased* memory usage. I'm not understanding the torch.compile interaction here. The complexity of FP8 custom ops isn't justified for lm_head alone. TODO to study in more detail the way this is implemented in other libraries, e.g. torchao.
 
@@ -913,7 +913,7 @@ Cherry-picked improvements from NorMuon (modded-nanogpt) into our simpler Muon i
 - Now defaults to ON for Muon via the `weight_decay` param. AdamW still has no weight decay and is hardcoded to 0 weight decay, might try to re-tune this later.
 
 **4. Weight decay schedule**
-- Added a linear schedule to weight decay that is default on from 1.0 to 0.0 (i.e. start with max weight decay in the beginning of training, them ramp to 0 by the end). Worked better than a static setting in experiments. (modded-nanogpt has the same schedule but it is imlpemented in a more confusing way by multiplying twice by the learning rate, which is already wired up to a decay schedule).
+- Added a linear schedule to weight decay that is default on from 1.0 to 0.0 (i.e. start with max weight decay in the beginning of training, then ramp to 0 by the end). Worked better than a static setting in experiments. (modded-nanogpt has the same schedule but it is implemented in a more confusing way by multiplying twice by the learning rate, which is already wired up to a decay schedule).
 
 ### Weight Decay Scaling Experiments
 
@@ -957,6 +957,6 @@ Muon was changed to use Polar Express, added NorMuon variance reduction, and cau
 
 **Bug Found:** Original implementation clipped local gradients before sync. Since this codebase doesn't use DDP (gradient sync is in the optimizers), each rank was clipping based on its own local norm. Fixed on the branch with proper distributed all-reduce.
 
-**Observartion:** modded-nanogpt does not appear to clip either right now.
+**Observation:** modded-nanogpt does not appear to clip either right now.
 
 **Summary:** Deleted all grad-clip code paths. The code naturally produces well-behaved gradients. This improves a bit of MFU because we don't have to calculate and sync grad norms.

nanochat/gpt.py

@@ -24,6 +24,7 @@ from nanochat.optim import MuonAdamW, DistMuonAdamW
 
 # Our custom Flash Attention module that automatically uses FA3 on Hopper+ and SDPA fallback elsewhere
 from nanochat.flash_attention import flash_attn
+from nanochat.moe import MoE
 
 @dataclass
 class GPTConfig:
@@ -33,6 +34,9 @@ class GPTConfig:
     n_head: int = 6 # number of query heads
     n_kv_head: int = 6 # number of key/value heads (GQA)
     n_embd: int = 768
+    num_experts: int = 8  # MoE: number of routed expert MLPs
+    top_k: int = 2  # MoE: number of active routed experts per token
+    num_shared_experts: int = 1  # MoE: number of shared (always-active) experts
     # Sliding window attention pattern string, tiled across layers. Final layer always L.
     # Characters: L=long (full context), S=short (half context)
     # Examples: "L"=all full context, "SL"=alternating, "SSL"=two short then one long
@@ -118,28 +122,15 @@ class CausalSelfAttention(nn.Module):
         return y
 
 
-class MLP(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
-        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
-
-    def forward(self, x):
-        x = self.c_fc(x)
-        x = F.relu(x).square()
-        x = self.c_proj(x)
-        return x
-
-
 class Block(nn.Module):
     def __init__(self, config, layer_idx):
         super().__init__()
         self.attn = CausalSelfAttention(config, layer_idx)
-        self.mlp = MLP(config)
+        self.moe = MoE(config)
 
     def forward(self, x, ve, cos_sin, window_size, kv_cache):
         x = x + self.attn(norm(x), ve, cos_sin, window_size, kv_cache)
-        x = x + self.mlp(norm(x))
+        x = x + self.moe(norm(x))
         return x
 
 
@@ -197,8 +188,11 @@ class GPT(nn.Module):
             attn.c_k:        uniform, std=1/sqrt(n_embd)
             attn.c_v:        uniform, std=1/sqrt(n_embd)
             attn.c_proj:     zeros
-            mlp.c_fc:        uniform, std=1/sqrt(n_embd)
-            mlp.c_proj:      zeros
+            moe.router.gate:     uniform, std=1/sqrt(n_embd)
+            moe.experts.w_up:           uniform, std=1/sqrt(n_embd)
+            moe.experts.w_down:          zeros
+            moe.shared_expert.w_up:      uniform, std=1/sqrt(n_embd)
+            moe.shared_expert.w_down:    zeros
         """
 
         # Embedding and unembedding
@@ -213,8 +207,16 @@ class GPT(nn.Module):
             torch.nn.init.uniform_(block.attn.c_k.weight, -s, s)
             torch.nn.init.uniform_(block.attn.c_v.weight, -s, s)
             torch.nn.init.zeros_(block.attn.c_proj.weight) # projections are zero
-            torch.nn.init.uniform_(block.mlp.c_fc.weight, -s, s)
-            torch.nn.init.zeros_(block.mlp.c_proj.weight)
+            # MoE: router gate and expert up-projections get uniform, down-projections get zero
+            torch.nn.init.uniform_(block.moe.router.gate.weight, -s, s)
+            torch.nn.init.uniform_(block.moe.experts.w_up, -s, s)
+            torch.nn.init.zeros_(block.moe.experts.w_down)
+            if block.moe.shared_expert is not None:
+                torch.nn.init.uniform_(block.moe.shared_expert.w_up.weight, -s, s)
+                torch.nn.init.zeros_(block.moe.shared_expert.w_down.weight)
+            # MoE load balancing buffers (zero after to_empty from meta device)
+            block.moe.router.expert_bias.zero_()
+            block.moe.router.tokens_per_expert_counter.zero_()
 
         # Per-layer scalars
         self.resid_lambdas.fill_(1.0)   # 1.0 => typical residual connections at init
@@ -306,6 +308,12 @@ class GPT(nn.Module):
         value_embeds_numel = sum(ve.weight.numel() for ve in self.value_embeds.values())
         nparams_exclude = (self.transformer.wte.weight.numel() + value_embeds_numel +
                           self.resid_lambdas.numel() + self.x0_lambdas.numel())
+        # MoE: only top_k/num_experts fraction of routed expert params active per token
+        # Shared expert is always active so its params stay in the active count
+        expert_hidden = self.transformer.h[0].moe.expert_hidden_dim
+        routed_params_per_layer = self.config.num_experts * 2 * self.config.n_embd * expert_hidden
+        inactive_per_layer = routed_params_per_layer * (self.config.num_experts - self.config.top_k) // self.config.num_experts
+        nparams_exclude += inactive_per_layer * self.config.n_layer
         h, q, t = self.config.n_head, self.config.n_embd // self.config.n_head, self.config.sequence_len
         # Sum attention FLOPs per layer, accounting for sliding window
         attn_flops = 0
@@ -327,6 +335,10 @@ class GPT(nn.Module):
 
         Returns a dict with counts for each parameter group, so downstream analysis
         can experiment with which combination gives the cleanest scaling laws.
+
+        For MoE, 'active_*' fields count only the parameters active per token
+        (top_k out of num_experts routed experts, plus shared experts).
+        Following DeepSeek convention of reporting both total and active params.
         """
         # Count each group separately (mirrors the grouping in setup_optimizers)
         wte = sum(p.numel() for p in self.transformer.wte.parameters())
@@ -336,13 +348,24 @@ class GPT(nn.Module):
         scalars = self.resid_lambdas.numel() + self.x0_lambdas.numel()
         total = wte + value_embeds + lm_head + transformer_matrices + scalars
         assert total == sum(p.numel() for p in self.parameters()), "Parameter count mismatch"
+        # MoE: only top_k/num_experts fraction of routed expert params active per token
+        # Shared expert is always active so its params stay in the active count
+        expert_hidden = self.transformer.h[0].moe.expert_hidden_dim
+        routed_params_per_layer = self.config.num_experts * 2 * self.config.n_embd * expert_hidden
+        inactive_per_layer = routed_params_per_layer * (self.config.num_experts - self.config.top_k) // self.config.num_experts
+        moe_inactive = inactive_per_layer * self.config.n_layer
+        active_transformer_matrices = transformer_matrices - moe_inactive
+        active_total = total - moe_inactive
         return {
             'wte': wte,
             'value_embeds': value_embeds,
             'lm_head': lm_head,
             'transformer_matrices': transformer_matrices,
+            'active_transformer_matrices': active_transformer_matrices,
             'scalars': scalars,
+            'moe_inactive': moe_inactive,
             'total': total,
+            'active_total': active_total,
         }
 
     def setup_optimizer(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, weight_decay=0.0, adam_betas=(0.8, 0.95), scalar_lr=0.5):
@@ -385,6 +408,31 @@ class GPT(nn.Module):
             group["initial_lr"] = group["lr"]
         return optimizer
 
+    def update_moe_balancing(self, coeff=1e-3):
+        """Update expert routing bias for load balancing. Call before optimizer.step()."""
+        for block in self.transformer.h:
+            block.moe.router.update_expert_bias(coeff)
+
+    def get_moe_stats(self):
+        """Collect MoE routing statistics for logging. Call BEFORE update_moe_balancing (which resets counters)."""
+        all_counts = []
+        all_biases = []
+        for block in self.transformer.h:
+            router = block.moe.router
+            all_counts.append(router.tokens_per_expert_counter)
+            all_biases.append(router.expert_bias)
+        counts = torch.stack(all_counts).float()    # (n_layer, num_experts)
+        biases = torch.stack(all_biases).float()    # (n_layer, num_experts)
+        # Load imbalance: coefficient of variation (std/mean) per layer, averaged
+        counts_mean = counts.mean(dim=-1).clamp(min=1)
+        counts_std = counts.std(dim=-1)
+        load_imbalance = (counts_std / counts_mean).mean().item()
+        return {
+            "moe/load_imbalance": load_imbalance,
+            "moe/expert_bias_std": biases.std().item(),
+            "moe/expert_bias_max": biases.abs().max().item(),
+        }
+
     def forward(self, idx, targets=None, kv_cache=None, loss_reduction='mean'):
         B, T = idx.size()
 

nanochat/moe.py

@@ -0,0 +1,241 @@
+"""
+Mixture of Experts (MoE) layer for nanochat.
+
+Replaces the standard dense MLP in each transformer block. Each token picks its
+top-K experts via a learned sigmoid router, so total parameters scale with
+num_experts but per-token FLOPs remain constant (iso-FLOP with the dense MLP).
+
+Expert hidden dim = 4 * dim / (top_k + num_shared), rounded to 128, ensures
+approximately iso-FLOP with the dense MLP:
+  Dense:            2 * dim * (4*dim) = 8*dim²
+  MoE per token:    (top_k + num_shared) * 2 * dim * H ≈ 8*dim²
+
+Expert weights are 3D tensors of shape (num_experts, hidden, dim). Muon's Polar
+Express orthogonalization operates on the last two dims, so the expert dimension
+acts as a batch dim and each expert is independently orthogonalized.
+
+At forward time, torch._grouped_mm dispatches tokens to experts via cumulative
+offsets — a single kernel per projection instead of a Python for-loop.
+"""
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class TopKRouter(nn.Module):
+    """Sigmoid-gated top-K router. Each token independently picks K experts."""
+
+    def __init__(self, dim, num_experts, top_k):
+        super().__init__()
+        self.gate = nn.Linear(dim, num_experts, bias=False)
+        self.num_experts = num_experts
+        self.top_k = top_k
+        # Auxiliary-loss-free load balancing (DeepSeekV3)
+        self.register_buffer('expert_bias', torch.zeros(num_experts))
+        self.register_buffer('tokens_per_expert_counter', torch.zeros(num_experts))
+
+    def forward(self, x):
+        """
+        Args:
+            x: (T, dim) flattened token representations
+        Returns:
+            top_scores:             (T, top_k)    routing weights for selected experts
+            selected_experts:       (T, top_k)    which experts each token chose
+            num_tokens_per_expert:  (num_experts,) how many tokens each expert received
+        """
+        scores = self.gate(x)                       # (T, num_experts)
+        scores = torch.sigmoid(scores.float())      # values in (0, 1)
+        # Bias affects expert SELECTION but not gating weights (DeepSeekV3)
+        biased_scores = scores + self.expert_bias
+        _, selected_experts = torch.topk(biased_scores, k=self.top_k, dim=-1, sorted=False)
+        top_scores = scores.gather(dim=-1, index=selected_experts)
+        num_tokens_per_expert = torch.histc(
+            selected_experts.float().view(-1),
+            bins=self.num_experts, min=0, max=self.num_experts,
+        )
+        # Accumulate token counts for load balancing updates
+        self.tokens_per_expert_counter += num_tokens_per_expert
+        return top_scores, selected_experts, num_tokens_per_expert
+
+    def update_expert_bias(self, coeff=1e-3):
+        """Auxiliary-loss-free bias update (DeepSeekV3). Call before optimizer.step()."""
+        counts = self.tokens_per_expert_counter
+        # Sync token counts across GPUs if distributed
+        if dist.is_initialized():
+            dist.all_reduce(counts)
+        if counts.sum() == 0:
+            return
+        mean_count = counts.mean()
+        # Nudge underloaded experts up, overloaded experts down
+        self.expert_bias += coeff * torch.sign(mean_count - counts)
+        self.expert_bias -= self.expert_bias.mean()  # center to prevent drift
+        self.tokens_per_expert_counter.zero_()
+
+
+def _run_experts_grouped_mm(w_up, w_down, x, num_tokens_per_expert):
+    """Run all experts via grouped matmul — single kernel per projection.
+
+    torch._grouped_mm handles variable tokens-per-expert internally via
+    cumulative offsets, so no Python for-loop or .tolist() device sync needed.
+    All tensor shapes are static (the dynamic token distribution is encoded
+    in the offsets, not in tensor dimensions).
+
+    Args:
+        w_up:   (num_experts, expert_hidden_dim, dim) - stacked up-projections
+        w_down: (num_experts, dim, expert_hidden_dim) - stacked down-projections
+        x:      (total_tokens, dim) - tokens sorted by expert assignment
+        num_tokens_per_expert: (num_experts,) - count per expert
+    Returns:
+        output: (total_tokens, dim)
+    """
+    offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32)
+    # Cast everything to bf16 upfront (weights are fp32 for Muon, need bf16 for grouped_mm)
+    x_bf16 = x.bfloat16()
+    w_up_bf16 = w_up.bfloat16().transpose(-2, -1)
+    w_down_bf16 = w_down.bfloat16().transpose(-2, -1)
+    # Up-project all experts at once: (total_tokens, dim) → (total_tokens, expert_hidden_dim)
+    h = torch._grouped_mm(x_bf16, w_up_bf16, offs=offsets)
+    h = F.relu(h).square()  # ReLU² activation
+    # Down-project all experts at once: (total_tokens, expert_hidden_dim) → (total_tokens, dim)
+    out = torch._grouped_mm(h.bfloat16(), w_down_bf16, offs=offsets)
+    return out.type_as(x)
+
+
+@torch.compiler.disable
+def _run_experts_for_loop(w_up, w_down, x, num_tokens_per_expert):
+    """Fallback for-loop implementation for CPU/MPS where grouped_mm isn't available.
+
+    Decorated with @torch.compiler.disable because .tolist() causes a device-host
+    sync that torch.compile can't handle. Only used on non-CUDA devices.
+    """
+    token_counts = num_tokens_per_expert.tolist()
+    chunks = torch.split(x, [int(c) for c in token_counts], dim=0)
+    outputs = []
+    for i, chunk in enumerate(chunks):
+        # No empty-chunk skip: matmul with (0, dim) tensors is valid and produces
+        # zero gradients (vs None), which the optimizer needs for stacking.
+        h = chunk @ w_up[i].T
+        h = F.relu(h).square()
+        h = h @ w_down[i].T
+        outputs.append(h)
+    return torch.cat(outputs, dim=0)
+
+
+class SharedExpert(nn.Module):
+    """Dense MLP shared expert — processes ALL tokens (no routing).
+
+    Same architecture as each routed expert (up → ReLU² → down) but uses
+    standard nn.Linear layers (2D weights, regular matmul) since there's
+    no need for the grouped_mm dispatch machinery.
+    """
+
+    def __init__(self, dim, expert_hidden_dim):
+        super().__init__()
+        self.w_up = nn.Linear(dim, expert_hidden_dim, bias=False)
+        self.w_down = nn.Linear(expert_hidden_dim, dim, bias=False)
+
+    def forward(self, x):
+        h = F.relu(self.w_up(x)).square()
+        return self.w_down(h)
+
+
+class ExpertGroup(nn.Module):
+    """
+    N independent expert MLPs stored as 3D weight tensors.
+    Shape (num_experts, hidden, dim) — Muon's Polar Express operates on the
+    last two dims, so each expert matrix is independently orthogonalized.
+    """
+
+    def __init__(self, dim, expert_hidden_dim, num_experts):
+        super().__init__()
+        self.num_experts = num_experts
+        self.w_up = nn.Parameter(torch.empty(num_experts, expert_hidden_dim, dim))
+        self.w_down = nn.Parameter(torch.empty(num_experts, dim, expert_hidden_dim))
+
+    def forward(self, x, num_tokens_per_expert):
+        """
+        Args:
+            x:                      (T*K, dim)     tokens sorted by expert assignment
+            num_tokens_per_expert:  (num_experts,) count per expert
+        Returns:
+            output:                 (T*K, dim)
+        """
+        if x.is_cuda:
+            return _run_experts_grouped_mm(self.w_up, self.w_down, x, num_tokens_per_expert)
+        return _run_experts_for_loop(self.w_up, self.w_down, x, num_tokens_per_expert)
+
+
+class MoE(nn.Module):
+    """
+    Mixture of Experts layer — approximately iso-FLOP replacement for the dense MLP.
+
+    For each token:
+    1. Shared expert processes all tokens via standard dense matmul
+    2. Router scores all routed experts via sigmoid(gate(x))
+    3. Top-K routed experts are selected
+    4. Token is dispatched to those experts (weighted by routing score)
+    5. Routed + shared expert outputs are summed together
+
+    Total active experts per token = top_k + num_shared_experts.
+    Expert hidden dim is sized so total active FLOPs ≈ dense MLP FLOPs.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        dim = config.n_embd
+        num_experts = config.num_experts
+        top_k = config.top_k
+        num_shared = config.num_shared_experts
+        self.top_k = top_k
+        # Iso-FLOP sizing: total active experts per token = top_k + num_shared
+        # Round to nearest 128 for tensor core alignment
+        active_experts = top_k + num_shared
+        expert_hidden_dim = round(4 * dim / active_experts / 128) * 128
+        self.expert_hidden_dim = expert_hidden_dim
+        self.router = TopKRouter(dim, num_experts, top_k)
+        self.experts = ExpertGroup(dim, expert_hidden_dim, num_experts)
+        self.shared_expert = SharedExpert(dim, expert_hidden_dim) if num_shared > 0 else None
+
+    def forward(self, x):
+        """
+        Args:  x: (bs, slen, dim)
+        Returns: output: (bs, slen, dim) — same shape, drop-in MLP replacement
+        """
+        bs, slen, dim = x.shape
+        x_flat = x.view(-1, dim)                                        # (T, dim)
+
+        # Step 1: Route — each token picks its top-K experts
+        top_scores, selected_experts, num_tokens_per_expert = self.router(x_flat)
+
+        # Step 2: Sort tokens by expert assignment for contiguous expert processing
+        # argsort groups all assignments to expert 0 first, then expert 1, etc.
+        token_indices_sorted = torch.argsort(selected_experts.view(-1), stable=True)
+        scores_sorted = top_scores.view(-1)[token_indices_sorted]       # (T*K,)
+        token_ids = token_indices_sorted // self.top_k                  # map back to original token
+        routed_input = x_flat[token_ids]                                # (T*K, dim)
+
+        # Step 3: Pre-multiply by routing scores (score_before_experts strategy)
+        routed_input = (routed_input.float() * scores_sorted.unsqueeze(-1)).to(x.dtype)
+
+        # Step 4: Shared expert — runs on ALL tokens via standard dense matmul
+        # Launched before routed experts so compute can overlap (no data dependency)
+        shared_output = self.shared_expert(x_flat) if self.shared_expert is not None else None
+
+        # Step 5: Run routed experts on their assigned token blocks
+        routed_output = self.experts(routed_input, num_tokens_per_expert)
+
+        # Step 6: Scatter outputs back to original positions and sum over top-K
+        combined = torch.zeros(
+            bs * slen * self.top_k, dim,
+            dtype=routed_output.dtype, device=routed_output.device,
+        )
+        combined[token_indices_sorted] = routed_output
+        output = combined.view(bs * slen, self.top_k, dim).sum(dim=1)   # (T, dim)
+
+        # Step 7: Add shared expert output
+        if shared_output is not None:
+            output = output + shared_output
+
+        return output.view(bs, slen, dim)

nanochat/optim.py

@@ -246,9 +246,13 @@ class MuonAdamW(torch.optim.Optimizer):
             state["momentum_buffer"] = torch.zeros(num_params, *shape, dtype=dtype, device=device)
         momentum_buffer = state["momentum_buffer"]
 
-        # Second momentum buffer is factored, either per-row or per-column
+        # Second momentum buffer is factored, either per-row or per-column.
+        # Uses *shape[:-1] / *shape[:-2] to preserve leading dims (e.g. expert dim for 3D MoE params).
         if "second_momentum_buffer" not in state:
-            state_shape = (num_params, shape[-2], 1) if shape[-2] >= shape[-1] else (num_params, 1, shape[-1])
+            if shape[-2] >= shape[-1]:
+                state_shape = (num_params, *shape[:-1], 1)
+            else:
+                state_shape = (num_params, *shape[:-2], 1, shape[-1])
             state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device)
         second_momentum_buffer = state["second_momentum_buffer"]
         red_dim = -1 if shape[-2] >= shape[-1] else -2
@@ -463,8 +467,12 @@ class DistMuonAdamW(torch.optim.Optimizer):
         state = self.state[p]
         if "momentum_buffer" not in state:
             state["momentum_buffer"] = torch.zeros(chunk_size, *shape, dtype=dtype, device=device)
+        # Second momentum buffer: preserve leading dims for 3D MoE params
         if "second_momentum_buffer" not in state:
-            state_shape = (chunk_size, shape[-2], 1) if shape[-2] >= shape[-1] else (chunk_size, 1, shape[-1])
+            if shape[-2] >= shape[-1]:
+                state_shape = (chunk_size, *shape[:-1], 1)
+            else:
+                state_shape = (chunk_size, *shape[:-2], 1, shape[-1])
             state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device)
         red_dim = -1 if shape[-2] >= shape[-1] else -2
 

scripts/base_train.py

@@ -237,6 +237,10 @@ def disable_fp8(model):
 # -----------------------------------------------------------------------------
 # Compile the model
 
+# MoE uses torch._grouped_mm with cumulative offsets — dynamo needs this to
+# trace through scalar tensor operations that arise from cumsum/histc in routing
+torch._dynamo.config.capture_scalar_outputs = True
+
 orig_model = model # original, uncompiled model, for saving raw model state_dict and for inference/evaluation (because the shapes may change shape)
 model = torch.compile(model, dynamic=False) # the inputs to model will never change shape so dynamic=False is safe
 
@@ -257,8 +261,9 @@ print0(f"Estimated FLOPs per token: {num_flops_per_token:e}")
 # We've already initialized the model so we have Params. Optimal Tokens is now simply target-param-data-ratio * Params
 def get_scaling_params(m):
     # As for which params to use exactly, transformer matrices + lm_head gives cleanest scaling laws (see dev/LOG.md Jan 27, 2026)
+    # For MoE, use active params (only top_k routed experts + shared, not all experts)
     params_counts = m.num_scaling_params()
-    scaling_params = params_counts['transformer_matrices'] + params_counts['lm_head']
+    scaling_params = params_counts['active_transformer_matrices'] + params_counts['lm_head']
     return scaling_params
 num_scaling_params = get_scaling_params(model)
 target_tokens = int(args.target_param_data_ratio * num_scaling_params) # optimal tokens for the model we are about to train
@@ -506,6 +511,8 @@ while True:
         if group['kind'] == 'muon':
             group["momentum"] = muon_momentum
             group["weight_decay"] = muon_weight_decay
+    moe_stats = orig_model.get_moe_stats() if step % 100 == 0 else {}
+    model.update_moe_balancing()
     optimizer.step()
     model.zero_grad(set_to_none=True)
     train_loss_f = train_loss.item() # .item() is a CPU-GPU sync point
@@ -547,6 +554,7 @@ while True:
             "train/mfu": mfu,
             "train/epoch": epoch,
         }
+        log_data.update(moe_stats)
         wandb_run.log(log_data)
 
     # state update

scripts/chat_rl.py

@@ -305,6 +305,7 @@ for step in range(num_steps):
     lrm = get_lr_multiplier(step)
     for group in optimizer.param_groups:
         group["lr"] = group["initial_lr"] * lrm
+    model.update_moe_balancing()
     optimizer.step()
     model.zero_grad(set_to_none=True)
     wandb_run.log({

scripts/chat_sft.py

@@ -118,6 +118,8 @@ for name, fallback, source in [
         print0(f"Using {name}={arg_val}")
 
 orig_model = model
+# MoE uses torch._grouped_mm — dynamo needs this for scalar tensor tracing
+torch._dynamo.config.capture_scalar_outputs = True
 model = torch.compile(model, dynamic=False)
 depth = model.config.n_layer
 num_flops_per_token = model.estimate_flops()
@@ -430,6 +432,7 @@ while True:
         group["lr"] = group["initial_lr"] * lrm
         if group['kind'] == 'muon':
             group["momentum"] = muon_momentum
+    model.update_moe_balancing()
     optimizer.step()
     model.zero_grad(set_to_none=True)
     synchronize()