From 9e4d3c15e4f984159894f7adabb9a9d9aacf6f0d Mon Sep 17 00:00:00 2001 From: Kefeng Zhou Date: Sun, 10 May 2026 13:51:43 +0700 Subject: [PATCH] =?UTF-8?q?feat(lora):=20PEFT=20=E2=86=92=20mlx-swift-lm?= =?UTF-8?q?=20adapter=20format=20converter?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second of three v0.5 LoRA PRs. Pure-Swift in-process converter so HuggingFace PEFT-format adapters (the de-facto standard on the Hub) work as drop-in inputs to mlx-swift-lm's LoRAContainer.from(directory:), which expects mlx's own format. Three translations: Config schema: PEFT r / lora_alpha / target_modules / peft_type mlx lora_parameters.{rank, scale, keys} + num_layers + fine_tune_type scale = lora_alpha / r num_layers inferred from deepest model.layers. seen in keys Weight key naming: PEFT base_model.model..lora_A.weight base_model.model..lora_B.weight mlx .lora_a / .lora_b Also collapses adjacent model.model.* runs. Tensor shape (transpose): PEFT lora_A is [rank, in], mlx lora_a is [in, rank] PEFT lora_B is [out, rank], mlx lora_b is [rank, out] forward = x @ a @ b ≡ x @ A.T @ B.T API surface: LoRAAdapterConverter.convertPEFTAdapter(source:destination: numLayersOverride:) writes a fresh mlx-format directory; source files are not modified. Errors typed via LoRAAdapterConverter.Error. Tests: 8 pure-Swift unit tests in LoRAAdapterConverterTests cover key rewrite (4 cases incl. throw on unrecognised suffix), layer-index extraction (2), config translation (2). 6 MLX-backed XCTest tests in LoRAAdapterConverterMLXTests cover weight transpose / dropping non- LoRA keys / tracking deepest layer / end-to-end on disk / two error paths — all gated on requireMetalOrSkip so SPM 'swift test' skips cleanly (no default.metallib) while xcodebuild executes. 133/133 Core tests green under swift test (+ 6 skipped MLX cases). Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 24 +- .../Adapters/LoRAAdapterConverter.swift | 239 ++++++++++++++++++ .../LoRAAdapterConverterMLXTests.swift | 165 ++++++++++++ .../Adapters/LoRAAdapterConverterTests.swift | 89 +++++++ 4 files changed, 514 insertions(+), 3 deletions(-) create mode 100644 MacMLXCore/Sources/MacMLXCore/Adapters/LoRAAdapterConverter.swift create mode 100644 MacMLXCore/Tests/MacMLXCoreTests/Adapters/LoRAAdapterConverterMLXTests.swift create mode 100644 MacMLXCore/Tests/MacMLXCoreTests/Adapters/LoRAAdapterConverterTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index 999d888..b3371ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,24 @@ Versioning follows [Semantic Versioning](https://semver.org/). ## [Unreleased] ### Added -- **LoRA Foundation** (v0.5, part 1 of 2). Pure-Swift Core layer +- **LoRA PEFT → mlx Converter** (v0.5, part 2 of 3). Pure-Swift + in-process converter that turns a HuggingFace PEFT-format adapter + directory into the mlx-swift-lm native format that + `MLXLMCommon.LoRAContainer.from(directory:)` expects. Three + translations: PEFT `r` / `lora_alpha` / `target_modules` → + mlx `lora_parameters.{rank, scale, keys}` (scale = alpha/rank); + PEFT keys `base_model.model..lora_A.weight` → mlx + `.lora_a` (drop wrapper, lowercase, drop `.weight`); + PEFT tensor shapes `[r, in]` / `[out, r]` → mlx `[in, r]` / + `[r, out]` (transpose). `num_layers` auto-inferred from the + deepest `model.layers.` index in the input keys. 8 pure-Swift + unit tests cover key rewrite + layer-index extraction + config + translation; 6 MLX-backed XCTest tests cover end-to-end + filesystem round-trip (gated on Metal — skip cleanly under + `swift test`, run under `xcodebuild`). Engine application of the + converted adapter (parameters-inspector picker + Settings + adapters section) lands in v0.5 part 3. +- **LoRA Foundation** (v0.5, part 1 of 3). Pure-Swift Core layer for HuggingFace LoRA adapter discovery — no engine integration yet, no UI. `LocalAdapter` value type holds adapter metadata parsed from PEFT's `adapter_config.json` (`base_model_name_or_path`, @@ -19,8 +36,9 @@ Versioning follows [Semantic Versioning](https://semver.org/). contain both `adapter_config.json` and `adapter_model.safetensors` — best-effort, malformed configs / missing weights silently drop. 10 new unit tests cover decode + scan paths. Engine application - of adapter weights via `LoRATrain.convert` + `loadLoRAWeights` - and the parameters-inspector picker land in v0.5 part 2. + via `MLXLMCommon.LoRAContainer.from(directory:)` (after PEFT → + mlx conversion) and the parameters-inspector picker land in + v0.5 part 3. - **Prompt cache tiering** (v0.4.0 engine parity, part 1 of 3). Successive chat turns on the same model now reuse the KV cache when the new prompt extends the previous one — the shared prefix diff --git a/MacMLXCore/Sources/MacMLXCore/Adapters/LoRAAdapterConverter.swift b/MacMLXCore/Sources/MacMLXCore/Adapters/LoRAAdapterConverter.swift new file mode 100644 index 0000000..a6e8b5d --- /dev/null +++ b/MacMLXCore/Sources/MacMLXCore/Adapters/LoRAAdapterConverter.swift @@ -0,0 +1,239 @@ +import Foundation +import MLX + +/// Converts a HuggingFace **PEFT**-format LoRA adapter directory into +/// the **mlx-swift-lm** native format that +/// `MLXLMCommon.LoRAContainer.from(directory:)` expects. +/// +/// The two formats differ in three ways: +/// +/// 1. **Config schema.** PEFT writes `r`, `lora_alpha`, `target_modules`, +/// `peft_type`. mlx writes `lora_parameters.{rank, scale, keys}` plus +/// `num_layers` and `fine_tune_type`. `scale = lora_alpha / r`. +/// +/// 2. **Weight key naming.** PEFT keys are +/// `base_model.model..lora_A.weight` and +/// `…lora_B.weight`. mlx keys drop the `base_model.model.` prefix, +/// drop the trailing `.weight`, and lowercase the `A`/`B` suffix +/// (`lora_a` / `lora_b`). +/// +/// 3. **Tensor shape.** PEFT stores `lora_A` as `[rank, in]` and +/// `lora_B` as `[out, rank]` (so `forward = x @ A.T @ B.T`). +/// mlx stores `lora_a` as `[in, rank]` and `lora_b` as +/// `[rank, out]` (so `forward = x @ a @ b`). Each weight is the +/// transpose of its PEFT counterpart. +/// +/// The converter writes the destination as a new directory containing +/// `adapter_config.json` (mlx schema) and `adapters.safetensors`. +/// Source files are not modified — callers that want an in-place +/// conversion should write to a sibling directory and rename. +public enum LoRAAdapterConverter { + + public enum Error: Swift.Error, CustomStringConvertible, Equatable { + case missingPEFTConfig(URL) + case missingPEFTWeights(URL) + case malformedPEFTConfig(String) + case noLoRAWeightsFound + case unrecognisedKeyFormat(String) + + public var description: String { + switch self { + case .missingPEFTConfig(let url): + return "PEFT adapter_config.json not found at \(url.path)" + case .missingPEFTWeights(let url): + return "PEFT adapter_model.safetensors not found at \(url.path)" + case .malformedPEFTConfig(let reason): + return "PEFT adapter_config.json could not be parsed: \(reason)" + case .noLoRAWeightsFound: + return "PEFT adapter_model.safetensors contained no recognisable LoRA weights" + case .unrecognisedKeyFormat(let key): + return "Unrecognised PEFT weight key shape: \(key)" + } + } + } + + /// Convert one PEFT-format adapter directory into a freshly-written + /// mlx-format adapter directory. + /// + /// - Parameters: + /// - source: directory containing `adapter_config.json` + + /// `adapter_model.safetensors` (PEFT). + /// - destination: directory to write `adapter_config.json` (mlx) + /// + `adapters.safetensors`. Should not equal `source` — + /// write to a sibling and rename if you want an in-place feel. + /// - numLayersOverride: explicit value for the mlx config's + /// `num_layers`. Pass `nil` to auto-infer from the deepest + /// `model.layers.` index seen in PEFT weight keys. + public static func convertPEFTAdapter( + source: URL, + destination: URL, + numLayersOverride: Int? = nil + ) throws { + let configURL = source.appendingPathComponent("adapter_config.json") + let weightsURL = source.appendingPathComponent("adapter_model.safetensors") + + guard FileManager.default.fileExists(atPath: configURL.path) else { + throw Error.missingPEFTConfig(configURL) + } + guard FileManager.default.fileExists(atPath: weightsURL.path) else { + throw Error.missingPEFTWeights(weightsURL) + } + + let peftConfig: LocalAdapter.PEFTConfig + do { + let data = try Data(contentsOf: configURL) + peftConfig = try JSONDecoder().decode(LocalAdapter.PEFTConfig.self, from: data) + } catch { + throw Error.malformedPEFTConfig(error.localizedDescription) + } + + let peftArrays = try MLX.loadArrays(url: weightsURL) + let (mlxArrays, deepestLayer) = try translateWeights(peftArrays) + guard !mlxArrays.isEmpty else { throw Error.noLoRAWeightsFound } + + let inferredNumLayers = (deepestLayer + 1) + let mlxConfig = mlxConfiguration( + from: peftConfig, + numLayers: numLayersOverride ?? inferredNumLayers + ) + + try FileManager.default.createDirectory( + at: destination, withIntermediateDirectories: true) + + // Write mlx config. + let outConfig = destination.appendingPathComponent("adapter_config.json") + let encoder = JSONEncoder() + encoder.outputFormatting = [.prettyPrinted, .sortedKeys] + try encoder.encode(mlxConfig).write(to: outConfig, options: .atomic) + + // Write mlx safetensors. MLX.save flushes the lazy compute + // graph internally before serialising — we don't need to + // manually realise the transposed arrays. + let outWeights = destination.appendingPathComponent("adapters.safetensors") + try MLX.save(arrays: mlxArrays, url: outWeights) + } + + // MARK: - Schema translation + + /// Internal mirror of `MLXLMCommon.LoRAConfiguration` so the + /// converter compiles in targets that don't import MLXLMCommon. + /// JSON shape is identical so the file mlx-swift-lm reads back + /// is byte-equivalent to one it would write itself. + struct MLXAdapterConfig: Codable, Equatable { + let numLayers: Int + let fineTuneType: String + let loraParameters: LoRAParameters + + struct LoRAParameters: Codable, Equatable { + let rank: Int + let scale: Float + let keys: [String]? + + private enum CodingKeys: String, CodingKey { + case rank, scale, keys + } + } + + private enum CodingKeys: String, CodingKey { + case numLayers = "num_layers" + case fineTuneType = "fine_tune_type" + case loraParameters = "lora_parameters" + } + } + + static func mlxConfiguration( + from peft: LocalAdapter.PEFTConfig, + numLayers: Int + ) -> MLXAdapterConfig { + let rank = peft.r ?? 8 + let alpha = Float(peft.loraAlpha ?? rank) + let scale = alpha / Float(rank) + return MLXAdapterConfig( + numLayers: numLayers, + fineTuneType: "lora", + loraParameters: .init( + rank: rank, + scale: scale, + keys: peft.targetModules + ) + ) + } + + // MARK: - Weight translation + + /// Translate PEFT-shaped weights into mlx-shaped weights. + /// + /// Returns the new dictionary plus the deepest `model.layers.` + /// index seen in the input keys (used to auto-infer `num_layers` + /// when the caller doesn't override it). + static func translateWeights( + _ peftArrays: [String: MLXArray] + ) throws -> (arrays: [String: MLXArray], deepestLayer: Int) { + var out: [String: MLXArray] = [:] + var deepest = -1 + + for (peftKey, peftArray) in peftArrays { + // Only translate keys ending in `.lora_A.weight` or + // `.lora_B.weight`. Other keys (e.g. PEFT's + // `…modules_to_save…`) are silently dropped — mlx-swift-lm's + // runtime cares only about the LoRA pair. + guard peftKey.hasSuffix(".lora_A.weight") || peftKey.hasSuffix(".lora_B.weight") else { + continue + } + + let mlxKey = try mlxKey(forPEFTKey: peftKey) + // PEFT stores transposed wrt mlx; `.T` materialises lazily, + // MLX.save flushes the graph below. + out[mlxKey] = peftArray.T + + if let layerIdx = layerIndex(in: mlxKey) { + deepest = max(deepest, layerIdx) + } + } + + return (out, deepest) + } + + /// Map one PEFT key to the mlx-equivalent key. + static func mlxKey(forPEFTKey peftKey: String) throws -> String { + // Drop the `base_model.` prefix(es). PEFT can wrap the model + // once (`base_model.model.<…>`) or twice for some causal-LM + // setups (`base_model.model.model.<…>`). + var key = peftKey + while key.hasPrefix("base_model.") { + key = String(key.dropFirst("base_model.".count)) + } + // Collapse adjacent `model.` runs so paths like + // `model.model.layers.0.…` become `model.layers.0.…` to match + // mlx's module hierarchy. + while key.hasPrefix("model.model.") { + key = String(key.dropFirst("model.".count)) + } + + // Suffix rewrite — drop `.weight`, lowercase the A/B side. + if key.hasSuffix(".lora_A.weight") { + key = key.dropLast(".lora_A.weight".count) + ".lora_a" + } else if key.hasSuffix(".lora_B.weight") { + key = key.dropLast(".lora_B.weight".count) + ".lora_b" + } else { + throw Error.unrecognisedKeyFormat(peftKey) + } + return key + } + + /// Extract the integer layer index from `…model.layers..…` + /// keys. Returns nil for keys outside the per-layer hierarchy + /// (embedding adapters etc.). + static func layerIndex(in mlxKey: String) -> Int? { + guard let range = mlxKey.range(of: ".layers.") else { return nil } + let tail = mlxKey[range.upperBound...] + let segment = tail.prefix { $0.isNumber } + return Int(segment) + } +} + +// SubSequence + String concat helper used in `mlxKey(forPEFTKey:)` +// to keep the suffix-rewrite line readable. +private func + (lhs: String.SubSequence, rhs: String) -> String { + String(lhs) + rhs +} diff --git a/MacMLXCore/Tests/MacMLXCoreTests/Adapters/LoRAAdapterConverterMLXTests.swift b/MacMLXCore/Tests/MacMLXCoreTests/Adapters/LoRAAdapterConverterMLXTests.swift new file mode 100644 index 0000000..c5092db --- /dev/null +++ b/MacMLXCore/Tests/MacMLXCoreTests/Adapters/LoRAAdapterConverterMLXTests.swift @@ -0,0 +1,165 @@ +import XCTest +import MLX +@testable import MacMLXCore + +/// MLX-backed converter tests — `MLXArray` allocation pulls in the +/// Metal stack, which the SPM test binary does not always bundle +/// (`default.metallib` ships with Xcode-archive builds only). These +/// tests skip cleanly when run under `swift test` and execute under +/// `xcodebuild`. +final class LoRAAdapterConverterMLXTests: XCTestCase { + + /// Mirror of the helper used in `PromptCacheStoreTests` — same + /// rationale, same skip behaviour. + private func requireMetalOrSkip() throws { + let bundle = Bundle(identifier: "mlx-swift_Cmlx.resources") + ?? Bundle.allBundles.first(where: { $0.bundlePath.contains("Cmlx") }) + let metallib = bundle?.url(forResource: "default", withExtension: "metallib") + if metallib == nil { + throw XCTSkip("Requires default.metallib (SPM test binaries often lack it — run under xcodebuild)") + } + } + + func testTranslateWeightsTransposesAndRenamesLoRAPair() throws { + try requireMetalOrSkip() + + // PEFT shape: lora_A is [rank, in], lora_B is [out, rank]. + let peftA = MLXArray([Float](repeating: 0.1, count: 8 * 4), [8, 4]) // r=8, in=4 + let peftB = MLXArray([Float](repeating: 0.2, count: 16 * 8), [16, 8]) // out=16, r=8 + + let inputs: [String: MLXArray] = [ + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": peftA, + "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": peftB, + ] + let result = try LoRAAdapterConverter.translateWeights(inputs) + + XCTAssertEqual(result.arrays.count, 2) + XCTAssertEqual(result.deepestLayer, 0) + + let aKey = "model.layers.0.self_attn.q_proj.lora_a" + let bKey = "model.layers.0.self_attn.q_proj.lora_b" + let outA = try XCTUnwrap(result.arrays[aKey]) + let outB = try XCTUnwrap(result.arrays[bKey]) + + // Shapes are transposed. + XCTAssertEqual(outA.shape, [4, 8], "lora_a should be [in, rank]") + XCTAssertEqual(outB.shape, [8, 16], "lora_b should be [rank, out]") + } + + func testTranslateWeightsDropsNonLoRAKeys() throws { + try requireMetalOrSkip() + + let lora = MLXArray([Float](repeating: 0.1, count: 8 * 4), [8, 4]) + let stray = MLXArray([Float](repeating: 0.5, count: 16), [16]) + let inputs: [String: MLXArray] = [ + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": lora, + "base_model.model.modules_to_save.score.weight": stray, + ] + let result = try LoRAAdapterConverter.translateWeights(inputs) + XCTAssertEqual(result.arrays.count, 1) + XCTAssertTrue(result.arrays.keys.contains("model.layers.0.self_attn.q_proj.lora_a")) + } + + func testTranslateWeightsTracksDeepestLayerAcrossPair() throws { + try requireMetalOrSkip() + + let p0 = MLXArray([Float](repeating: 0.1, count: 8 * 4), [8, 4]) + let p11 = MLXArray([Float](repeating: 0.1, count: 8 * 4), [8, 4]) + let p23 = MLXArray([Float](repeating: 0.1, count: 8 * 4), [8, 4]) + let inputs: [String: MLXArray] = [ + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": p0, + "base_model.model.model.layers.11.self_attn.q_proj.lora_A.weight": p11, + "base_model.model.model.layers.23.self_attn.q_proj.lora_A.weight": p23, + ] + let result = try LoRAAdapterConverter.translateWeights(inputs) + XCTAssertEqual(result.deepestLayer, 23) + } + + func testConvertPEFTAdapterWritesMLXConfigAndSafetensors() throws { + try requireMetalOrSkip() + + let temp = try TempDir() + let source = temp.url.appendingPathComponent("peft-source", isDirectory: true) + let dest = temp.url.appendingPathComponent("mlx-output", isDirectory: true) + try FileManager.default.createDirectory(at: source, withIntermediateDirectories: true) + + let peftCfg = """ + { + "base_model_name_or_path": "Qwen3-8B-4bit", + "r": 4, + "lora_alpha": 8, + "target_modules": ["q_proj", "v_proj"], + "peft_type": "LORA" + } + """ + try Data(peftCfg.utf8).write(to: source.appendingPathComponent("adapter_config.json")) + + let peftA = MLXArray([Float](repeating: 0.1, count: 4 * 4), [4, 4]) + let peftB = MLXArray([Float](repeating: 0.2, count: 8 * 4), [8, 4]) + let weights: [String: MLXArray] = [ + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": peftA, + "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": peftB, + ] + try MLX.save(arrays: weights, url: source.appendingPathComponent("adapter_model.safetensors")) + + try LoRAAdapterConverter.convertPEFTAdapter(source: source, destination: dest) + + // Output config is mlx schema. + let outCfgURL = dest.appendingPathComponent("adapter_config.json") + let outCfg = try JSONDecoder().decode( + LoRAAdapterConverter.MLXAdapterConfig.self, + from: Data(contentsOf: outCfgURL) + ) + XCTAssertEqual(outCfg.fineTuneType, "lora") + XCTAssertEqual(outCfg.loraParameters.rank, 4) + XCTAssertEqual(outCfg.loraParameters.scale, 2.0) // 8 / 4 + XCTAssertEqual(outCfg.loraParameters.keys, ["q_proj", "v_proj"]) + XCTAssertEqual(outCfg.numLayers, 1) // inferred from layer 0 + + // Output safetensors round-trips. + let outArrays = try MLX.loadArrays( + url: dest.appendingPathComponent("adapters.safetensors") + ) + XCTAssertEqual(outArrays.count, 2) + let outA = try XCTUnwrap(outArrays["model.layers.0.self_attn.q_proj.lora_a"]) + let outB = try XCTUnwrap(outArrays["model.layers.0.self_attn.q_proj.lora_b"]) + XCTAssertEqual(outA.shape, [4, 4]) // transposed from [4, 4] (square) + XCTAssertEqual(outB.shape, [4, 8]) // transposed from [8, 4] + } + + func testConvertThrowsForMissingConfig() throws { + try requireMetalOrSkip() + + let temp = try TempDir() + let source = temp.url.appendingPathComponent("empty", isDirectory: true) + try FileManager.default.createDirectory(at: source, withIntermediateDirectories: true) + let dest = temp.url.appendingPathComponent("dest", isDirectory: true) + XCTAssertThrowsError( + try LoRAAdapterConverter.convertPEFTAdapter(source: source, destination: dest) + ) + } + + func testConvertThrowsForMissingWeights() throws { + try requireMetalOrSkip() + + let temp = try TempDir() + let source = temp.url.appendingPathComponent("config-only", isDirectory: true) + try FileManager.default.createDirectory(at: source, withIntermediateDirectories: true) + try Data(#"{"r":4}"#.utf8) + .write(to: source.appendingPathComponent("adapter_config.json")) + let dest = temp.url.appendingPathComponent("dest", isDirectory: true) + XCTAssertThrowsError( + try LoRAAdapterConverter.convertPEFTAdapter(source: source, destination: dest) + ) + } +} + +private struct TempDir { + let url: URL + init() throws { + let base = FileManager.default.temporaryDirectory + .appendingPathComponent("macmlx-lora-converter-tests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: base, withIntermediateDirectories: true) + self.url = base + } +} diff --git a/MacMLXCore/Tests/MacMLXCoreTests/Adapters/LoRAAdapterConverterTests.swift b/MacMLXCore/Tests/MacMLXCoreTests/Adapters/LoRAAdapterConverterTests.swift new file mode 100644 index 0000000..c2ecf14 --- /dev/null +++ b/MacMLXCore/Tests/MacMLXCoreTests/Adapters/LoRAAdapterConverterTests.swift @@ -0,0 +1,89 @@ +import Testing +import Foundation +@testable import MacMLXCore + +@Suite("LoRAAdapterConverter — pure Swift translation") +struct LoRAAdapterConverterUnitTests { + + // MARK: Key rewrite + + @Test + func mlxKeyStripsBaseModelPrefixAndRewritesSuffixA() throws { + let key = try LoRAAdapterConverter.mlxKey( + forPEFTKey: "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" + ) + #expect(key == "model.layers.0.self_attn.q_proj.lora_a") + } + + @Test + func mlxKeyStripsBaseModelPrefixAndRewritesSuffixB() throws { + let key = try LoRAAdapterConverter.mlxKey( + forPEFTKey: "base_model.model.model.layers.7.self_attn.v_proj.lora_B.weight" + ) + #expect(key == "model.layers.7.self_attn.v_proj.lora_b") + } + + @Test + func mlxKeyHandlesSingleBaseModelWrapper() throws { + let key = try LoRAAdapterConverter.mlxKey( + forPEFTKey: "base_model.model.layers.3.mlp.gate_proj.lora_A.weight" + ) + #expect(key == "model.layers.3.mlp.gate_proj.lora_a") + } + + @Test + func mlxKeyThrowsForUnrecognisedSuffix() { + #expect(throws: LoRAAdapterConverter.Error.self) { + _ = try LoRAAdapterConverter.mlxKey( + forPEFTKey: "base_model.model.model.embed_tokens.weight" + ) + } + } + + // MARK: Layer-index extraction + + @Test + func layerIndexFromTypicalKey() { + #expect(LoRAAdapterConverter.layerIndex(in: "model.layers.0.self_attn.q_proj.lora_a") == 0) + #expect(LoRAAdapterConverter.layerIndex(in: "model.layers.31.mlp.up_proj.lora_b") == 31) + } + + @Test + func layerIndexNilForNonLayerKey() { + #expect(LoRAAdapterConverter.layerIndex(in: "model.embed_tokens.lora_a") == nil) + } + + // MARK: Config translation + + @Test + func mlxConfigurationDerivesScaleFromAlphaOverRank() { + let peft = LocalAdapter.PEFTConfig( + baseModelNameOrPath: "Qwen3-8B-4bit", + r: 8, + loraAlpha: 16, + targetModules: ["q_proj", "v_proj"], + peftType: "LORA" + ) + let cfg = LoRAAdapterConverter.mlxConfiguration(from: peft, numLayers: 24) + #expect(cfg.numLayers == 24) + #expect(cfg.fineTuneType == "lora") + #expect(cfg.loraParameters.rank == 8) + #expect(cfg.loraParameters.scale == 2.0) // 16 / 8 + #expect(cfg.loraParameters.keys == ["q_proj", "v_proj"]) + } + + @Test + func mlxConfigurationDefaultsRankAndAlphaWhenMissing() { + let peft = LocalAdapter.PEFTConfig( + baseModelNameOrPath: nil, + r: nil, + loraAlpha: nil, + targetModules: nil, + peftType: nil + ) + let cfg = LoRAAdapterConverter.mlxConfiguration(from: peft, numLayers: 1) + #expect(cfg.loraParameters.rank == 8) // default + #expect(cfg.loraParameters.scale == 1.0) // alpha default = rank → 1.0 + #expect(cfg.loraParameters.keys == nil) + } +}