markbaseengine/ROUTER_SCALE_FIX_APPLIED.md

Router Scale Normalization Fix Applied

Fix Date

2026-06-20 22:16

Problem

routerScale = 31.25 (raw value, too large)

• Causes softmax overflow in MoE router computation
• Similar to 26B-Standard scales issue

Solution

Normalize routerScale by hiddenSize

// Model.swift:516-519 (modified)
let rawRouterScale = rsFloats.first ?? 1.0
routerScale = rawRouterScale / Float(hiddenSize)

Effect:

• Before: routerScale = 31.25
• After: routerScale = 31.25 / 2816 = 0.01105
• Result: Stable softmax, no overflow

Why This Works

Similar to 26B-Standard fix:

26B-Standard scales:
  - Raw: ~120
  - Fix: Divide by hiddenSize (120/2816 = 0.0426)
  - Result: Fixed NaN, model works

26B-A4B routerScale:
  - Raw: 31.25
  - Fix: Divide by hiddenSize (31.25/2816 = 0.01105)
  - Expected: Fix generation hanging

Router computation flow:

1. Router logits (raw): [numExperts]
2. Scale logits: logits * routerScale
3. Softmax: exp(scaled_logits) / sum

If routerScale too large:

• scaled_logits = logits * 31.25
• exp(scaled_logits) can overflow
• NaN in softmax
• Generation hangs

If routerScale normalized:

• scaled_logits = logits * 0.01105
• exp(scaled_logits) stable
• Softmax works
• Generation succeeds

Code Changes

File: /Users/accusys/MarkBase12B/Sources/G12B/Model.swift

Lines: 516-519

Before:

let rsData = try rsReader.read(tensor: rsDesc)
let rsFloats = SafeTensorsReader.bf16ToFloat32(rsData)
routerScale = rsFloats.first ?? 1.0  // Raw value

After:

let rsData = try rsReader.read(tensor: rsDesc)
let rsFloats = SafeTensorsReader.bf16ToFloat32(rsData)
let rawRouterScale = rsFloats.first ?? 1.0
// Normalize router scale by hidden_size (similar to scales normalization for 26B-Standard)
// This prevents softmax overflow in MoE router computation
routerScale = rawRouterScale / Float(hiddenSize)

Testing

Next step: Test generation with normalized routerScale

Expected:

• ✅ Generation works (no hang)
• ✅ No NaN in router computation
• ✅ Stable softmax
• ✅ Valid token generation

Test command:

swift test --filter MoEDebugTests/test26BA4BSimpleGenerationDebug

Pattern Recognition

Normalization pattern discovered:

1. 26B-Standard scales: Divide by hiddenSize
2. 26B-A4B routerScale: Divide by hiddenSize
3. Pattern: Raw scale values need normalization by hiddenSize

General rule:

If scale value is loaded from tensor and seems large (>10):
  → Normalize by dividing by hiddenSize
  → Prevents numerical overflow
  → Matches model training normalization

Confidence

Confidence level: ⭐⭐⭐⭐⭐ (Very High)

Reasons:

• ✅ Same pattern as 26B-Standard fix (proven to work)
• ✅ Router scale purpose is to scale logits before softmax
• ✅ Large scale values cause overflow
• ✅ Normalization prevents overflow

If Fix Works

Implications:

• ✅ 26B-A4B MoE model will work
• ✅ First MoE model successfully running
• ✅ MoE implementation validated
• ✅ Pattern for fixing MoE numerical issues

Comparison:

26B-Standard: 40 tok/s (Dense, already works)
26B-A4B MoE: Expected 20-30 tok/s (MoE, should work after fix)
31B-IT: 11.7 tok/s (Dense, already works)

If Fix Doesn't Work

Next debugging steps:

1. Check expert scales normalization
2. Add NaN checks in router computation
3. Test router forward pass separately
4. Check Metal kernels

But: High confidence this fix will work

Summary

✅ Fix applied: Router scale normalization

📊 Expected result: Generation works (no hang)

🔧 Pattern: Normalize scale values by hiddenSize

⏱️ Next: Test generation with fix

---

Status: Fix applied, ready for testing