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,162 @@
|
||||
# 26B-A4B Router Scale Analysis - Potential Issue Found
|
||||
|
||||
## Discovery Date
|
||||
2026-06-20 22:13
|
||||
|
||||
## ✅ Router Structure Test: PASSED
|
||||
|
||||
### Router Components Verified
|
||||
```
|
||||
Layer 0 Router:
|
||||
✓ routerProj: present (8-bit, inDim=2816, outDim=128)
|
||||
✓ routerScale: 31.25 ⚠️ POTENTIAL ISSUE
|
||||
✓ perExpertScale: present [128 values]
|
||||
✓ topK: 8
|
||||
|
||||
Expert Components:
|
||||
✓ expertGate: present (128 experts, 704 output, 2816 input, 4-bit)
|
||||
✓ expertUp: present (same structure)
|
||||
✓ expertDown: present (same structure)
|
||||
```
|
||||
|
||||
### ⚠️ Key Finding: routerScale = 31.25
|
||||
|
||||
**Potential Issue**: Router scale value is 31.25, which might need normalization
|
||||
|
||||
**Comparison with 26B-Standard**:
|
||||
```
|
||||
26B-Standard scales issue:
|
||||
- Original: scales ~120
|
||||
- Problem: Too large, caused numerical issues
|
||||
- Fix: Normalize by hidden_size (120/2816 = 0.0426)
|
||||
- Result: Fixed NaN issues
|
||||
|
||||
26B-A4B router scale:
|
||||
- Current: routerScale = 31.25
|
||||
- Question: Is this already normalized? Or needs normalization?
|
||||
- Potential fix: Divide by hidden_size? (31.25/2816 = 0.011)
|
||||
```
|
||||
|
||||
### Router Scale Purpose
|
||||
|
||||
In MoE models, router scale is used to scale router logits before softmax:
|
||||
```swift
|
||||
// Layer.swift:837 (moeForward)
|
||||
var scaled = routerData.map { $0 * routerScale }
|
||||
```
|
||||
|
||||
**Effect**:
|
||||
- If routerScale is too large → softmax overflow
|
||||
- If routerScale is too small → softmax underflow
|
||||
- Both cause numerical instability or NaN
|
||||
|
||||
### Analysis
|
||||
|
||||
**Router computation flow**:
|
||||
1. Router proj: input [hidden_size] → output [num_experts]
|
||||
2. Raw logits: ~some range
|
||||
3. Scale logits: logits * routerScale
|
||||
4. Softmax: exp(scaled_logits) / sum
|
||||
|
||||
**If routerScale=31.25 is too large**:
|
||||
- scaled_logits could overflow exp() function
|
||||
- NaN in softmax computation
|
||||
- Generation hangs or crashes
|
||||
|
||||
### Hypothesis
|
||||
|
||||
**routerScale might need normalization**:
|
||||
```swift
|
||||
// Possible fix in Model.swift
|
||||
let routerScale = rsFloats.first ?? 1.0
|
||||
let normalizedRouterScale = routerScale / Float(hiddenSize)
|
||||
|
||||
// Use normalizedRouterScale in Layer
|
||||
```
|
||||
|
||||
**Or**: routerScale is already correct and issue is elsewhere
|
||||
|
||||
### Testing Required
|
||||
|
||||
1. **Check router computation values**:
|
||||
- What are raw router logits?
|
||||
- What are scaled logits?
|
||||
- Do they overflow?
|
||||
|
||||
2. **Try normalization**:
|
||||
- Divide routerScale by hidden_size
|
||||
- Test if generation works
|
||||
|
||||
3. **Check softmax implementation**:
|
||||
- Is it handling overflow correctly?
|
||||
- Are there NaN checks?
|
||||
|
||||
### Related Code
|
||||
|
||||
**Router scale loading** (Model.swift:508-519):
|
||||
```swift
|
||||
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 // Gets first value
|
||||
}
|
||||
```
|
||||
|
||||
**Router scale usage** (Layer.swift:837):
|
||||
```swift
|
||||
var scaled = routerData.map { $0 * routerScale }
|
||||
```
|
||||
|
||||
### Comparison with Other Models
|
||||
|
||||
| Model | MoE | routerScale | Notes |
|
||||
|-------|-----|-------------|-------|
|
||||
| 26B-Standard | No | N/A | Uses scales normalization (120/2816) |
|
||||
| 31B-IT | No | N/A | Dense, no router |
|
||||
| **26B-A4B** | Yes | **31.25** | Needs investigation |
|
||||
|
||||
### Next Steps
|
||||
|
||||
**Immediate**:
|
||||
1. ✅ Run generation test (currently in progress)
|
||||
2. If hangs → try router scale normalization
|
||||
3. Test with routerScale / hiddenSize
|
||||
|
||||
**If normalization fixes**:
|
||||
- Add normalization to Model.swift
|
||||
- Similar to scales normalization fix
|
||||
- Document in validation report
|
||||
|
||||
**If normalization doesn't fix**:
|
||||
- Check other potential issues
|
||||
- Expert selection logic
|
||||
- Metal kernels
|
||||
- Forward pass sequence
|
||||
|
||||
### Files
|
||||
|
||||
**Test code**:
|
||||
- `/Users/accusys/MarkBase12B/Tests/G12BTests/MoEDebugTests.swift`
|
||||
|
||||
**Test output**:
|
||||
- `/Users/accusys/MarkBase12B/MOE_ROUTER_STRUCTURE_TEST.log`
|
||||
|
||||
**Model**:
|
||||
- `/Users/accusys/MarkBase12B/models/gemma-4-26b-a4b-it-4bit/`
|
||||
|
||||
**Router scale tensor**:
|
||||
- `language_model.model.layers.0.router.scale`
|
||||
- Shape: [2816] bf16
|
||||
- Value: 31.25 (first element)
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
**✅ Router structure is correct and complete**
|
||||
|
||||
**⚠️ Potential issue**: routerScale=31.25 might need normalization
|
||||
|
||||
**🔧 Possible fix**: Divide by hiddenSize (31.25/2816 = 0.011)
|
||||
|
||||
**📊 Test result**: Router structure test passed, generation test in progress
|
||||
Reference in New Issue
Block a user