Initial commit: E4B-MarkBase model integration with passing tests
CI / build-and-test (push) Has been cancelled
CI / build-and-test (push) Has been cancelled
- 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
This commit is contained in:
@@ -0,0 +1,156 @@
|
||||
# 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**
|
||||
```swift
|
||||
// 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**:
|
||||
```swift
|
||||
let rsData = try rsReader.read(tensor: rsDesc)
|
||||
let rsFloats = SafeTensorsReader.bf16ToFloat32(rsData)
|
||||
routerScale = rsFloats.first ?? 1.0 // Raw value
|
||||
```
|
||||
|
||||
**After**:
|
||||
```swift
|
||||
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**:
|
||||
```bash
|
||||
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
|
||||
Reference in New Issue
Block a user