markbaseengine/NAN_FIX_SUMMARY.md

NaN Bug Fix Summary

Problem

MarkBaseServer forward pass produced NaN in all model outputs, preventing successful inference.

Root Cause Analysis

Investigation Chain

1. Layer 0 DownProj → NaN output
2. DownProj input (gate buffer) → NaN at position 7782+
3. Gate buffer NaN source → fusedGateUp kernel
4. Kernel NaN origin → Out-of-bounds scales/biases access
5. Buffer size mismatch → Scales/biases loaded as BF16 (2 bytes) instead of Float32 (4 bytes)

Critical Discovery

Safetensors stores scales/biases as BF16 (2 bytes per element), but code loaded them as raw bytes into Metal buffer without conversion.

Expected vs Actual:

• Expected scales size: 15360 × 60 = 921,600 floats = 3,686,400 bytes
• Actual buffer size: 1,843,200 bytes = 460,800 floats (half-size!)

Kernel Impact: For output position 7782:

• Expected scales index: 7782 × 60 = 466,920
• Buffer capacity: 460,800 floats
• Access beyond bounds → garbage/NaN values

Fixes Applied

1. BF16→Float32 Conversion (CRITICAL FIX)

File: Sources/MarkBase/Model.swift:559-597

// Convert scales from BF16 to Float32 (safetensors stores as BF16)
let sBuf: MTLBuffer?
if sDesc?.dtype == .bf16 {
    let sFloats = SafeTensorsReader.bf16ToFloat32(sData)
    sBuf = engine.device.makeBuffer(
        bytes: sFloats, length: sFloats.count * MemoryLayout<Float>.stride,
        options: .storageModeShared
    )
} else {
    sBuf = sData.withUnsafeBytes { ptr in
        engine.device.makeBuffer(bytes: ptr.baseAddress!, length: sData.count, options: .storageModeShared)
    }
}

// Same conversion for biases

Before:

• Scales buffer: 1,843,200 bytes = 460,800 floats

After:

• Scales buffer: 3,686,400 bytes = 921,600 floats ✅

2. groupSize Calculation Fix

File: Sources/MarkBase/Model.swift:610

// FIX: groupSize = inDim / sShape[1], NOT sShape[1] directly
// scales shape is [outDim, inDim/groupSize], so sShape[1] = inDim/groupSize
let groupSize = (sShape.count > 1 && sShape[1] > 0) ? inDim / sShape[1] : 64

Before: groupSize = sShape[1] (wrong interpretation) After: groupSize = inDim / sShape[1] (correct calculation)

3. Fallback Kernel groupSize Parameter

File: Sources/MarkBase/Layers/Layer.swift:374

// Fallback to original
let pso = try engine.pipeline(named: "quantized_matmul")
let enc = cmdBuf.makeComputeCommandEncoder()!
enc.setComputePipelineState(pso)
enc.setBuffer(input, offset: 0, index: 0)
enc.setBuffer(weights.weight, offset: 0, index: 1)
enc.setBuffer(weights.scales, offset: 0, index: 2)
enc.setBuffer(weights.biases, offset: 0, index: 3)
enc.setBuffer(output, offset: 0, index: 4)
var inDim = UInt32(weights.inDim)
enc.setBytes(&inDim, length: MemoryLayout<UInt32>.size, index: 5)
var outDim = UInt32(weights.outDim)
enc.setBytes(&outDim, length: MemoryLayout<UInt32>.size, index: 6)
var groupSize = UInt32(weights.groupSize)  // FIX: Add groupSize!
enc.setBytes(&groupSize, length: MemoryLayout<UInt32>.size, index: 7)

Before: Missing groupSize parameter (index 7) After: Correctly passes groupSize to kernel ✅

Test Results

Before Fix

Layer 0:
  Gate buffer: [7782]=nan, [7800]=10.0
  DownProj: h=[nan, nan, nan, nan, nan]
  NaN count: 262,144/262,144

After Fix

Layer 0:
  Gate buffer: [7782]=0.0815, [7800]=0.0763 (valid!)
  DownProj: h=[1.07, 1.04, 8.47, -1.77, -1.82] (valid!)
  
All layers:
  NaN count: 0/262,144 ✅
  Has NaN: false ✅
  
Final logits:
  Max: 30.0, Min: -29.99 ✅
  Top tokens generated successfully ✅

Technical Details

Safetensors Storage Format

• Dtype: BF16 (bfloat16)
• Size: 2 bytes per element
• Range: Same as Float32 but reduced precision
• Use case: Saves memory/storage space

Metal Kernel Requirements

• All buffer inputs must be Float32 (4 bytes)
• Buffer sizes must match kernel expectations
• Out-of-bounds access → undefined behavior/NaN

Conversion Method

SafeTensorsReader.bf16ToFloat32() implementation:

public static func bf16ToFloat32(_ data: Data) -> [Float] {
    data.withUnsafeBytes { ptr in
        let bf16 = ptr.assumingMemoryBound(to: UInt16.self)
        return (0..<data.count / 2).map { i in
            Float(bitPattern: UInt32(bf16[i]) << 16)
        }
    }
}

Impact

Models Fixed

• ✅ E4B-MarkBase (4.4GB)
• ✅ E4B-12B (6.3GB)
• ✅ E4B-26B-Standard (15GB)
• ✅ E4B-31B (17GB)

Performance

• No performance impact (conversion happens during model loading)
• Correct inference (all layers produce valid output)
• Target performance: <100ms/token (previously achieved 21-27ms)

Files Modified

1. Sources/MarkBase/Model.swift

   • Lines 559-597: BF16→Float32 conversion
   • Line 610: groupSize calculation fix

2. Sources/MarkBase/Layers/Layer.swift

   • Line 374: Fallback kernel groupSize parameter

Deployment

1. Build:

   cd ~/MarkBaseEngine
   swift build -c release --product MarkBaseServer
   

2. Test:

   .build/release/MarkBaseServer
   

3. Deploy to M5Max48:

   • Copy binary to target machine
   • Test with all models
   • Monitor for NaN in logs

Verification Checklist

• ✅ Scales/biases dtype check (BF16)
• ✅ Buffer size verification (2× original)
• ✅ Forward pass NaN check (0 NaN)
• ✅ Logit range check ([-30, 30])
• ✅ Token generation test (valid output)

Future Considerations

1. ** Dtype detection** - Check all tensor dtypes during loading
2. ** Automatic conversion** - Handle BF16, FP16, other formats
3. ** Kernel robustness** - Add bounds checking in Metal shaders
4. ** Testing framework** - Automated NaN detection tests

---

Date: 2025-06-23
Status: ✅ FIXED
Impact: Critical fix enabling all model inference