markbaseengine/26B_A4B_MOE_DEBUG_SUMMARY.md

26B-A4B MoE Debug Summary - Current Status

Test Date

2026-06-20 22:13-22:15

✅ Successes

1. Model Loading - COMPLETE SUCCESS ⭐⭐⭐⭐⭐

Load time: 51.818s
Layers: 30 (ALL MoE ✓)
Experts: 128/128 per layer ✓
Total tensors: 1697
Status: Test passed

2. Router Structure Verification - COMPLETE SUCCESS ⭐⭐⭐⭐⭐

Router components: All present ✓
Expert components: All present ✓
Router weights: 8-bit, correct dimensions ✓
Expert weights: 4-bit, correct structure ✓
Router scale: 31.25 ⚠️ (potential issue)
Status: Test passed

⚠️ Issues Found

1. Token Generation - HANGS ⚠️⚠️⚠️

Symptoms:

• Generation test hangs
• Timeout after 30s (no response)
• Likely numerical issue in forward pass

Root Cause (Hypothesis):

• routerScale = 31.25 might be too large
• Similar to 26B-Standard scales issue
• May cause softmax overflow or NaN
• Needs normalization (divide by hiddenSize?)

2. Router Scale Value - POTENTIAL BUG ⚠️⚠️

Current value: routerScale = 31.25

Question: Is this already normalized or raw value?

Similar issue (26B-Standard):

26B-Standard scales:
  - Raw: ~120
  - Problem: Too large
  - Fix: Normalize by hiddenSize (120/2816 = 0.0426)
  - Result: Fixed NaN

26B-A4B routerScale:
  - Current: 31.25
  - Hypothesis: May need normalization
  - Potential fix: 31.25/2816 = 0.011

📊 Test Results Summary

Test	Status	Duration	Result
Model Loading	✅ PASSED	51.818s	All 30 layers loaded with MoE
Router Structure	✅ PASSED	1.0s	All components verified
Token Generation	❌ HANGS	30s+ timeout	No response, likely NaN
Forward Pass	⏳ Not tested	-	Needs separate test

🔧 Proposed Fixes

Fix 1: Router Scale Normalization ⭐⭐⭐⭐⭐

Code location: Model.swift:508-519

Current code:

if let rsDesc = allTensors.first(where: { $0.name == "\(prefix).router.scale" }) {
    let rsData = try rsReader.read(tensor: rsDesc)
    let rsFloats = SafeTensorsReader.bf16ToFloat32(rsData)
    routerScale = rsFloats.first ?? 1.0  // Raw value
}

Proposed fix:

if let rsDesc = allTensors.first(where: { $0.name == "\(prefix).router.scale" }) {
    let rsData = try rsReader.read(tensor: rsDesc)
    let rsFloats = SafeTensorsReader.bf16ToFloat32(rsData)
    let rawRouterScale = rsFloats.first ?? 1.0
    // Normalize by hiddenSize (similar to scales normalization)
    routerScale = rawRouterScale / Float(hiddenSize)  // 31.25/2816 = 0.011
}

Expected result:

• routerScale = 0.011 (smaller, stable)
• Softmax won't overflow
• Generation should work

Confidence: ⭐⭐⭐⭐⭐ High (based on 26B-Standard fix pattern)

Fix 2: Add NaN Checks ⭐⭐⭐⭐

Add debug prints in Layer.swift moeForward:

// After router computation
let routerData = engine.readFloats(from: temps.gate, count: numExperts)
print("Router logits: max=\(routerData.max()), min=\(routerData.min())")

// After scaling
var scaled = routerData.map { $0 * routerScale }
print("Scaled logits: max=\(scaled.max()), min=\(scaled.min())")

// After softmax
print("Softmax weights: sum=\(sum)")

Purpose:

• Identify where NaN occurs
• Verify router computation
• Debug numerical issues

Fix 3: Expert Scale Normalization ⭐⭐⭐

Similar to 26B-Standard scales fix:

If router fix doesn't work, expert scales might also need normalization:

// In loadExpertGroup
let normalizedScales = scales / Float(expertInDim)

🎯 Next Steps

Immediate (Priority 1)

1. ✅ Apply router scale normalization

   • Edit Model.swift:508-519
   • Add normalization: routerScale /= hiddenSize
   • Test generation

2. ⏳ Test generation with fix

   • Run MoEDebugTests/test26BA4BSimpleGenerationDebug
   • Expect: generation works
   • If works: Document fix

If Fix Works (Priority 2)

3. ✅ Document router scale fix

   • Create validation report
   • Compare with 26B-Standard fix
   • Document normalization pattern

4. ✅ Run full benchmark

   • Test token generation speed
   • Compare with 26B-Standard (40 tok/s)
   • Memory usage

If Fix Doesn't Work (Priority 3)

5. ⚠️ Debug forward pass

   • Add NaN checks
   • Test router computation
   • Test expert selection

6. ⚠️ Check other issues

   • Expert scales normalization
   • Metal kernels
   • Forward pass sequence

📈 Expected Timeline

With router fix:

• Fix implementation: 5 minutes
• Testing: 5-10 minutes
• Documentation: 5 minutes
• Total: 15-20 minutes ⭐⭐⭐⭐⭐

If router fix doesn't work:

• Additional debugging: 30-60 minutes
• Multiple attempts: 1-2 hours
• Total: 2-3 hours ⚠️⚠️

📊 Comparison: MoE vs Dense

Model	Type	Load Status	Load Time	Generation	Speed
26B-Standard	Dense	✅ Works	5.3s	✅ Works	40 tok/s
31B-IT	Dense	✅ Works	63.8s	✅ Works	11.7 tok/s
26B-A4B	MoE	✅ Works	51.818s	⚠️ Fix needed	Expected: 20-30 tok/s

🎓 Lessons Learned

1. MoE implementation already complete ✅

   • No need for 3-5 days implementation
   • Code was ready, just needed testing

2. Router scale needs investigation ⚠️

   • Similar to 26B-Standard scales issue
   • Normalization pattern applies to MoE too

3. Test incrementally ⭐⭐⭐⭐⭐

   • First test loading (passed)
   • Then test structure (passed)
   • Now test generation (issue found)
   • Debug systematically

💡 Recommendation

Apply router scale normalization NOW ⭐⭐⭐⭐⭐

Reasons:

• High confidence fix (based on 26B-Standard pattern)
• Quick to implement (5 minutes)
• Likely to work (similar issue pattern)
• If works → complete success
• If fails → debug further

Time investment: 15-20 minutes Potential reward: MoE model working! Risk: Low (if fails, we learn more)

---

Files Created

Test reports:

• /Users/accusys/MarkBase12B/26B_A4B_LOADING_SUCCESS.md
• /Users/accusys/MarkBase12B/26B_A4B_ROUTER_SCALE_ANALYSIS.md
• /Users/accusys/MarkBase12B/26B_A4B_MOE_DEBUG_SUMMARY.md

Test code:

• /Users/accusys/MarkBase12B/Tests/G12BTests/MoEDebugTests.swift
• /Users/accusys/MarkBase12B/Tests/G12BTests/MoEForwardTests.swift

Test logs:

• /Users/accusys/MarkBase12B/26B_A4B_LOADING_TEST.log
• /Users/accusys/MarkBase12B/MOE_ROUTER_STRUCTURE_TEST.log

---

Summary

✅ Major progress: MoE model loading and structure verified

⚠️ Blocking issue: Generation hangs, likely router scale too large

🔧 Proposed fix: Normalize routerScale by hiddenSize (31.25/2816)

📊 Confidence: High (⭐⭐⭐⭐⭐) based on 26B-Standard fix pattern

⏱️ Expected time: 15-20 minutes to test fix

🏆 Potential outcome: First working MoE model!