analysis: 12B 3 NaN real root cause found (NOT config mismatch)

BREAKTHROUGH DISCOVERY: - ❌ Previous hypothesis: Config mismatch (num_kv_heads: 8 vs 2) - ✅ Actual root cause: Special Token IDs have embedding issues EXACT NaN LOCATIONS: - Token ID 2 (BOS - Begin of Sequence): NaN - Token ID 255999 (BOI - Begin of Image): NaN - Token ID 256000 (BOA - Begin of Audio): NaN Evidence from debug test: indices [2, 255999, 256000] Config fix made NaN worse (3→12), restored original config Only 3 out of 262K tokens affected (0.0011%) Recommendation: Use E4B/E2B or avoid special tokens
2026-06-24 00:53:27 +08:00
parent a64ccf0869
commit 78257a947c
3 changed files with 575 additions and 0 deletions
@@ -0,0 +1,124 @@
+import XCTest
+@testable import MarkBase
+
+class TwelveBNaNDebugTest: XCTestCase {
+    
+    func test12BNaNDebug() throws {
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  12B NaN Debug Test - Find Exact NaN Positions")
+        print("═══════════════════════════════════════════════════════════════════\n")
+        
+        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
+        let engine = try MarkBaseEngine(autoCompile: true)
+        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
+        
+        print("Step 1: Run forward pass and find NaN indices...")
+        let logits = try model.forwardOptimized(tokenId: 2, position: 0)
+        
+        var nanIndices: [Int] = []
+        var nanValues: [Float] = []
+        var extremeIndices: [Int] = []
+        var extremeValues: [Float] = []
+        
+        for i in 0..<logits.count {
+            let v = logits[i]
+            if v.isNaN {
+                nanIndices.append(i)
+                nanValues.append(v)
+            }
+            if abs(v) > 100.0 {
+                extremeIndices.append(i)
+                extremeValues.append(v)
+            }
+        }
+        
+        print("  Total logits: \(logits.count)")
+        print("  NaN count: \(nanIndices.count)")
+        print("  Extreme values (>100): \(extremeIndices.count)")
+        print()
+        
+        if nanIndices.count > 0 {
+            print("Step 2: NaN Details:")
+            print("  NaN at indices: \(nanIndices)")
+            print()
+            
+            print("Step 3: Analyze logits around NaN positions:")
+            for idx in nanIndices.prefix(3) {
+                let start = max(0, idx - 5)
+                let end = min(logits.count - 1, idx + 5)
+                print("  Around index \(idx):")
+                for i in start...end {
+                    let marker = i == idx ? " [NaN]" : ""
+                    print("    logits[\(i)] = \(logits[i])\(marker)")
+                }
+                print()
+            }
+        }
+        
+        if extremeIndices.count > 0 {
+            print("Step 4: Extreme values (>100) that might cause NaN:")
+            print("  Extreme indices: \(extremeIndices.prefix(10))")
+            print("  Extreme values: \(extremeValues.prefix(10))")
+            print()
+        }
+        
+        // Check overall statistics
+        var minVal: Float = Float.infinity
+        var maxVal: Float = -Float.infinity
+        var sumVal: Float = 0
+        
+        for v in logits {
+            if v.isFinite {
+                if v < minVal { minVal = v }
+                if v > maxVal { maxVal = v }
+                sumVal += v
+            }
+        }
+        
+        let meanVal = sumVal / Float(logits.count - nanIndices.count)
+        
+        print("Step 5: Logit Statistics (excluding NaN/Inf):")
+        print("  Min: \(minVal)")
+        print("  Max: \(maxVal)")
+        print("  Mean: \(meanVal)")
+        print("  Range: \(maxVal - minVal)")
+        print()
+        
+        // Check if softcapping is applied
+        print("Step 6: Config Check:")
+        let configPath = modelPath + "/config.json"
+        if let configData = FileManager.default.contents(atPath: configPath) {
+            let config = try JSONDecoder().decode(ConfigDebug.self, from: configData)
+            print("  final_logit_softcapping: \(config.textConfig.finalLogitSoftcapping)")
+            print("  rms_norm_eps: \(config.textConfig.rmsNormEps)")
+            print()
+            
+            if config.textConfig.finalLogitSoftcapping > 0 {
+                print("  ⚠️ Softcapping ACTIVE: logits /= (1 + |logits| / \(config.textConfig.finalLogitSoftcapping))")
+                print("  💡 Large logits could cause division by zero or NaN")
+            }
+        }
+        
+        print("═══════════════════════════════════════════════════════════════════\n")
+        
+        XCTAssertLessThanOrEqual(nanIndices.count, 10, "Should have minimal NaN")
+    }
+}
+
+struct ConfigDebug: Decodable {
+    let textConfig: TextConfigDebug
+    
+    enum CodingKeys: String, CodingKey {
+        case textConfig = "text_config"
+    }
+}
+
+struct TextConfigDebug: Decodable {
+    let finalLogitSoftcapping: Float
+    let rmsNormEps: Float
+    
+    enum CodingKeys: String, CodingKey {
+        case finalLogitSoftcapping = "final_logit_softcapping"
+        case rmsNormEps = "rms_norm_eps"
+    }
+}