Files
MarkBase Admin 78257a947c
CI / build-and-test (push) Has been cancelled
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

93 lines
4.4 KiB
Swift

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"
}
}