- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
5.9 KiB
Router Scale Fix Result - Needs Further Investigation
Test Date
2026-06-20 22:17-22:19
❌ Router Scale Normalization Fix Did NOT Solve Generation Hanging
Fix Applied
// Model.swift:518
routerScale = rawRouterScale / Float(hiddenSize)
// Before: 31.25
// After: 31.25/2816 = 0.01105
Test Result
Generation test: STILL HANGS (timeout after 120s)
No improvement: Router scale normalization alone did not fix the issue
⚠️ New Findings
Issue Complexity
Not just router scale: Multiple normalization issues possible
Potential additional problems:
-
Expert scales normalization
- Expert gate/up/down scales might need normalization
- Similar to 26B-Standard scales fix
-
Router proj weights normalization
- Router projection output might need scaling
-
Expert intermediate computation
- Expert fusion computation might overflow
-
Top-k expert selection
- Expert selection logic might hang
Next Steps Required
Immediate debugging:
- ✅ Add debug prints to MoE forward pass
- ✅ Check router computation step by step
- ✅ Check expert scales values
- ✅ Check expert selection process
Additional normalization fixes:
- ⏳ Expert scales normalization (divide by expertInDim?)
- ⏳ Router proj output normalization
- ⏳ Expert intermediate normalization
Comparison: What Worked for 26B-Standard
26B-Standard had multiple fixes:
Fix 1: Scales normalization (divide by hiddenSize)
Fix 2: Logits scaling (multiply by 0.00486)
Fix 3: Remove softcapping
Fix 4: Sampler temperature fix
26B-A4B might need similar multiple fixes:
Fix 1: Router scale normalization (applied, but not enough)
Fix 2: Expert scales normalization (not yet applied)
Fix 3: Router output normalization (not yet applied)
Fix 4: Expert intermediate normalization (not yet applied)
🔍 Debugging Strategy
Step 1: Add Debug Prints
Add to Layer.swift moeForward:
// After router computation
let routerData = engine.readFloats(from: temps.gate, count: numExperts)
print("Router logits: \(routerData[0..<10])")
print("Router max/min: \(routerData.max()), \(routerData.min())")
// After scaling
var scaled = routerData.map { $0 * routerScale }
print("Scaled logits: \(scaled[0..<10])")
print("Scaled max/min: \(scaled.max()), \(scaled.min())")
// After softmax
print("Softmax weights: \(scaled[0..<10])")
Step 2: Check Expert Scales
Add to Model.swift loadExpertGroup:
// After loading expert scales
print("Expert scales first 10: \(scalesData[0..<10])")
let expertScalesMax = scalesData.max()
print("Expert scales max: \(expertScalesMax)")
// If large (>100), need normalization
Step 3: Test Router Forward Pass
Create minimal router test:
- Test router computation only (no expert)
- Check if router works with normalized scale
- Verify softmax is stable
📊 Current Status
| Component | Status | Issue |
|---|---|---|
| Model loading | ✅ Works | All 30 layers, 3840 experts |
| Router structure | ✅ Works | All components present |
| Router scale fix | ⚠️ Applied | Normalized (31.25→0.01105) |
| Token generation | ❌ Hangs | Timeout 120s, no response |
| Expert computation | ⏳ Unknown | Needs testing |
💡 Revised Assessment
Router Scale Fix Confidence
Previous confidence: ⭐⭐⭐⭐⭐ (5/5) Actual result: ❌ Did not fix
Lesson: MoE models have more complex normalization requirements than Dense models
New Hypothesis
MoE normalization complexity:
- Router scale normalization (tried, not enough)
- Expert scales normalization (not tried yet)
- Multiple normalization steps needed
Similar to 26B-Standard: Multiple fixes required MoE adds: More components need normalization (router + experts)
🎯 Next Action Plan
Option A: Add Debug Prints (Recommended) ⭐⭐⭐⭐⭐
Reason: Need to see where it hangs Time: 10-15 minutes Benefit: Identify exact problem location
Steps:
- Add debug prints to moeForward
- Run test with prints
- Identify where it hangs
- Fix specific issue
Option B: Try Expert Scales Fix ⭐⭐⭐⭐
Reason: Expert scales might be too large Time: 5-10 minutes Benefit: Additional normalization
Steps:
- Add expert scales normalization
- Divide by expertInDim (2816)
- Test generation
Option C: Multiple Fixes ⭐⭐⭐
Reason: Combine router + expert fixes Time: 15-20 minutes Benefit: Comprehensive fix
Steps:
- Router scale fix (already applied)
- Expert scales fix
- Router output fix
- Test generation
📈 Timeline Estimate
Option A (Debug prints):
- Add prints: 10 minutes
- Run test: 2-5 minutes
- Analyze: 5-10 minutes
- Fix issue: 10-30 minutes
- Total: 30-60 minutes ⭐⭐⭐⭐⭐
Option B (Expert fix):
- Apply fix: 5 minutes
- Test: 2-5 minutes
- Total: 7-10 minutes ⭐⭐⭐⭐
Option C (Multiple fixes):
- Apply multiple fixes: 15-20 minutes
- Test: 2-5 minutes
- Total: 20-25 minutes ⭐⭐⭐
Recommendation
Use Option A (Debug prints) ⭐⭐⭐⭐⭐
Reasons:
- Router scale fix didn't work → need to see where hangs
- Debug prints give visibility
- Identify exact problem
- Fix specific issue
Alternative: Combine A + B (add debug prints + expert scales fix)
Files Updated
Fix applied:
/Users/accusys/MarkBase12B/Sources/G12B/Model.swift(lines 516-519)
Documentation:
/Users/accusys/MarkBase12B/ROUTER_SCALE_FIX_APPLIED.md/Users/accusys/MarkBase12B/26B_A4B_ROUTER_FIX_FAILED_ANALYSIS.md
Summary
✅ Router scale fix applied: 31.25 → 0.01105 (normalized)
❌ Generation still hangs: Router fix not sufficient
⏳ Next: Add debug prints to identify exact hang location
📊 Lesson: MoE needs multiple normalization fixes, similar to 26B-Standard
💡 Recommendation: Add debug prints to moeForward, identify where it hangs