Initial commit: E4B-MarkBase model integration with passing tests

- 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
2026-06-23 18:12:35 +08:00
commit ac75faa0cc
301 changed files with 63426 additions and 0 deletions
@@ -0,0 +1,122 @@
+import XCTest
+@testable import MarkBase
+
+class A4BDiagnosticTest: XCTestCase {
+    
+    func testEmbeddingScalesNaNCheck() throws {
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  26B-A4B Embedding Scales NaN Diagnostic")
+        print("═══════════════════════════════════════════════════════════════════\n")
+        
+        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
+        
+        guard FileManager.default.fileExists(atPath: modelPath) else {
+            print("⚠ Model not found")
+            return
+        }
+        
+        // Load SafeTensors readers
+        let index = try SafeTensorsIndex(modelDir: modelPath)
+        
+        var readers: [String: SafeTensorsReader] = [:]
+        for shardFile in index.weightMap.values {
+            if readers[shardFile] == nil {
+                readers[shardFile] = try SafeTensorsReader(path: "\(modelPath)/\(shardFile)")
+            }
+        }
+        
+        print("Loaded \(readers.count) shard readers")
+        
+        // Find embed_tokens tensors
+        let allTensors = readers.values.flatMap { $0.allTensors }
+        
+        let embedWeightTensor = allTensors.first { $0.name.contains("embed_tokens.weight") }
+        let embedScalesTensor = allTensors.first { $0.name.contains("embed_tokens.scales") }
+        let embedBiasesTensor = allTensors.first { $0.name.contains("embed_tokens.biases") }
+        
+        print("\nEmbedding tensors:")
+        if let wt = embedWeightTensor {
+            print("  weight: shape=\(wt.shape), dtype=\(wt.dtype), size=\(wt.dataSize)")
+        }
+        if let st = embedScalesTensor {
+            print("  scales: shape=\(st.shape), dtype=\(st.dtype), size=\(st.dataSize)")
+        }
+        if let bt = embedBiasesTensor {
+            print("  biases: shape=\(bt.shape), dtype=\(bt.dtype), size=\(bt.dataSize)")
+        }
+        
+        // Read scales and check for NaN
+        if let st = embedScalesTensor {
+            print("\nChecking scales for NaN...")
+            let reader = readers[index.weightMap[st.name] ?? "model-00001-of-00003.safetensors"]!
+            let scalesData = try reader.read(tensor: st)
+            
+            // Parse scales as Float array
+            let scales = scalesData.withUnsafeBytes { ptr in
+                Array(ptr.assumingMemoryBound(to: Float.self))
+            }
+            
+            let nanCount = scales.filter { $0.isNaN }.count
+            let infCount = scales.filter { $0.isInfinite }.count
+            
+            print("  Total scales: \(scales.count)")
+            print("  NaN count: \(nanCount)")
+            print("  Inf count: \(infCount)")
+            
+            if nanCount > 0 {
+                // Find NaN positions
+                let nanIndices = scales.enumerated().filter { $0.element.isNaN }.map { $0.offset }
+                print("  NaN positions (first 20): \(nanIndices.prefix(20))")
+                
+                // Calculate which tokens these belong to
+                // Assuming scales shape = [vocabSize, groupSize] or [vocabSize]
+                let vocabSize = 262144
+                let groupSize = st.shape.count > 1 ? st.shape[1] : 1
+                
+                print("\n  Token mapping:")
+                for idx in nanIndices.prefix(10) {
+                    let tokenId = idx / groupSize
+                    let groupId = idx % groupSize
+                    print("    scales[\(idx)] = NaN → token \(tokenId), group \(groupId)")
+                }
+            }
+            
+            // Check scales distribution
+            let validScales = scales.filter { !$0.isNaN && !$0.isInfinite }
+            if !validScales.isEmpty {
+                let avgScale = validScales.reduce(0, +) / Float(validScales.count)
+                let minScale = validScales.min() ?? 0
+                let maxScale = validScales.max() ?? 0
+                print("\n  Valid scales distribution:")
+                print("    min=\(minScale), max=\(maxScale), avg=\(avgScale)")
+            }
+        }
+        
+        // Read biases and check for NaN
+        if let bt = embedBiasesTensor {
+            print("\nChecking biases for NaN...")
+            let reader = readers[index.weightMap[bt.name] ?? "model-00001-of-00003.safetensors"]!
+            let biasesData = try reader.read(tensor: bt)
+            
+            let biases = biasesData.withUnsafeBytes { ptr in
+                Array(ptr.assumingMemoryBound(to: Float.self))
+            }
+            
+            let nanCount = biases.filter { $0.isNaN }.count
+            let infCount = biases.filter { $0.isInfinite }.count
+            
+            print("  Total biases: \(biases.count)")
+            print("  NaN count: \(nanCount)")
+            print("  Inf count: \(infCount)")
+            
+            if nanCount > 0 {
+                let nanIndices = biases.enumerated().filter { $0.element.isNaN }.map { $0.offset }
+                print("  NaN positions (first 20): \(nanIndices.prefix(20))")
+            }
+        }
+        
+        print("\n═══════════════════════════════════════════════════════════════════")
+        print("  Diagnosis Complete")
+        print("═══════════════════════════════════════════════════════════════════\n")
+    }
+}