diff --git a/Sources/SwiftLM/ModelProfiler.swift b/Sources/SwiftLM/ModelProfiler.swift index 2914de5..ba04abb 100644 --- a/Sources/SwiftLM/ModelProfiler.swift +++ b/Sources/SwiftLM/ModelProfiler.swift @@ -204,6 +204,8 @@ enum ModelProfiler { let headDim: Int? let intermediateSize: Int? let vocabSize: Int? + let numExperts: Int? + let numExpertsPerTok: Int? enum CodingKeys: String, CodingKey { case numHiddenLayers = "num_hidden_layers" @@ -213,6 +215,8 @@ enum ModelProfiler { case headDim = "head_dim" case intermediateSize = "intermediate_size" case vocabSize = "vocab_size" + case numExperts = "num_experts" + case numExpertsPerTok = "num_experts_per_tok" } } @@ -252,9 +256,9 @@ enum ModelProfiler { let quantBits = config.quantizationConfig?.bits ?? detectQuantBits(modelId: modelId) // Detect MoE - let isMoE = config.numExperts != nil && (config.numExperts ?? 0) > 1 - let numExperts = config.numExperts - let numActiveExperts = config.numExpertsPerTok + let numExperts = config.numExperts ?? config.textConfig?.numExperts + let numActiveExperts = config.numExpertsPerTok ?? config.textConfig?.numExpertsPerTok + let isMoE = (numExperts ?? 0) > 1 // Measure weight file sizes on disk (only for MoE to avoid slow walks on dense models) let weightSize = isMoE ? measureWeightFiles(directory: modelDirectory) : 0 diff --git a/tests/SwiftLMTests/ModelProfilerTests.swift b/tests/SwiftLMTests/ModelProfilerTests.swift new file mode 100644 index 0000000..872f344 --- /dev/null +++ b/tests/SwiftLMTests/ModelProfilerTests.swift @@ -0,0 +1,45 @@ +import Foundation +import XCTest +@testable import SwiftLM + +final class ModelProfilerTests: XCTestCase { + func testDetectsNestedQwenMoEConfiguration() throws { + let profile = try profile(json: """ + { + "model_type": "qwen3_5_moe", + "text_config": { + "num_experts": 256, + "num_experts_per_tok": 8 + } + } + """) + + XCTAssertTrue(profile.isMoE) + XCTAssertEqual(profile.numExperts, 256) + XCTAssertEqual(profile.numActiveExperts, 8) + } + + func testKeepsDetectingTopLevelLocalExperts() throws { + let profile = try profile(json: """ + { + "model_type": "legacy_moe", + "num_local_experts": 64, + "num_experts_per_tok": 4 + } + """) + + XCTAssertTrue(profile.isMoE) + XCTAssertEqual(profile.numExperts, 64) + XCTAssertEqual(profile.numActiveExperts, 4) + } + + private func profile(json: String) throws -> ModelProfile { + let directory = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: directory) } + + try Data(json.utf8).write(to: directory.appendingPathComponent("config.json")) + return try XCTUnwrap(ModelProfiler.profile(modelDirectory: directory, modelId: "test-model")) + } +}