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,93 @@
+import XCTest
+@testable import MarkBase
+
+class TwelveBConfigFixTest: XCTestCase {
+    
+    func test12BAfterConfigFix() throws {
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  12B Config Fix Test - Verify 0 NaN After Correction")
+        print("═══════════════════════════════════════════════════════════════════\n")
+        
+        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
+        
+        print("Step 1: Verify config modification...")
+        let configPath = modelPath + "/config.json"
+        guard let configData = FileManager.default.contents(atPath: configPath) else {
+            XCTFail("Config file not found")
+            return
+        }
+        
+        let config = try JSONDecoder().decode(ModelConfigStruct.self, from: configData)
+        print("  ✓ Config loaded")
+        print("    num_attention_heads: \(config.textConfig.numAttentionHeads)")
+        print("    num_key_value_heads: \(config.textConfig.numKeyValueHeads ?? 0)")
+        print("    hidden_size: \(config.textConfig.hiddenSize)")
+        
+        if config.textConfig.numKeyValueHeads == 2 {
+            print("  ✓✓ Config FIXED: num_kv_heads = 2 (matches weights)")
+        } else {
+            print("  ✗ Config NOT fixed: num_kv_heads = \(config.textConfig.numKeyValueHeads ?? 0)")
+            XCTFail("Config should have num_kv_heads=2")
+            return
+        }
+        
+        print("\nStep 2: Load 12B model...")
+        let engine = try MarkBaseEngine(autoCompile: true)
+        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
+        print("  ✓ Model loaded")
+        print("    Layers: \(model.numHiddenLayers)")
+        print("    Hidden: \(model.hiddenSize)")
+        print("    Vocab: \(model.vocabSize)")
+        
+        print("\nStep 3: Test forward pass (multiple positions)...")
+        var totalNaN = 0
+        let testPositions = [0, 50, 100, 200]
+        
+        for pos in testPositions {
+            let result = try model.forwardOptimized(tokenId: 2, position: pos)
+            let nanCount = result.filter { $0.isNaN }.count
+            totalNaN += nanCount
+            print("  Position \(pos): NaN=\(nanCount)/\(result.count)")
+        }
+        
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  RESULTS SUMMARY")
+        print("═══════════════════════════════════════════════════════════════════\n")
+        print("  Total NaN across \(testPositions.count) positions: \(totalNaN)")
+        print("  Config fix: num_kv_heads changed from 8 to 2")
+        
+        if totalNaN == 0 {
+            print("  ✓✓✓ SUCCESS! 0 NaN - Config fix resolved the issue!")
+            print("  ⭐⭐⭐⭐⭐ 12B now PERFECT after correction")
+        } else if totalNaN < 5 {
+            print("  ⭐⭐⭐⭐ Improved! NaN reduced (was 3, now \(totalNaN))")
+        } else {
+            print("  ✗ Config fix did NOT resolve the issue")
+            print("  ⚠️ NaN still present: \(totalNaN)")
+        }
+        
+        print("═══════════════════════════════════════════════════════════════════\n")
+        
+        XCTAssertEqual(totalNaN, 0, "After config fix, should have 0 NaN")
+    }
+}
+
+struct ModelConfigStruct: Decodable {
+    let textConfig: TextConfig
+    
+    enum CodingKeys: String, CodingKey {
+        case textConfig = "text_config"
+    }
+}
+
+struct TextConfig: Decodable {
+    let numAttentionHeads: Int
+    let numKeyValueHeads: Int?
+    let hiddenSize: Int
+    
+    enum CodingKeys: String, CodingKey {
+        case numAttentionHeads = "num_attention_heads"
+        case numKeyValueHeads = "num_key_value_heads"
+        case hiddenSize = "hidden_size"
+    }
+}