From 19e81ee573f5682509bc9590204e616f5606fba0 Mon Sep 17 00:00:00 2001 From: Aegis-AI Date: Sun, 26 Apr 2026 22:17:17 -0700 Subject: [PATCH 1/3] fix(moe): gracefully catch safetensors I/O errors to prevent deadlock crashes --- Libraries/MLXLMCommon/ConcurrentError.swift | 27 +++++++++++ Libraries/MLXLMCommon/SwitchLayers.swift | 28 +++++++++++ .../MLXLMTests/CorruptSafetensorsTests.swift | 47 +++++++++++++++++++ 3 files changed, 102 insertions(+) create mode 100644 Libraries/MLXLMCommon/ConcurrentError.swift create mode 100644 Tests/MLXLMTests/CorruptSafetensorsTests.swift diff --git a/Libraries/MLXLMCommon/ConcurrentError.swift b/Libraries/MLXLMCommon/ConcurrentError.swift new file mode 100644 index 000000000..fc14ebb5e --- /dev/null +++ b/Libraries/MLXLMCommon/ConcurrentError.swift @@ -0,0 +1,27 @@ +import Foundation +import MLX + +package final class ThreadSafeError: @unchecked Sendable { + package let lock = NSLock() + package var error: Swift.Error? + + package init() {} + + package func catchError(_ block: () throws -> Void) { + do { + try MLX.withError { + try block() + } + } catch { + lock.withLock { + if self.error == nil { self.error = error } + } + } + } + + package func check() { + if let error = error { + fatalError("MLX SSD Streaming Error: \(error.localizedDescription). (The model safetensors file may be corrupted, truncated, or incomplete).") + } + } +} diff --git a/Libraries/MLXLMCommon/SwitchLayers.swift b/Libraries/MLXLMCommon/SwitchLayers.swift index f1affcfa9..7cb112dda 100644 --- a/Libraries/MLXLMCommon/SwitchLayers.swift +++ b/Libraries/MLXLMCommon/SwitchLayers.swift @@ -272,7 +272,9 @@ public class SwitchGLU: Module, @unchecked Sendable { if !specTargets.isEmpty { let bpe = _stackedBytesPerExpert let downBpe = _stackedDownBytesPerExpert + let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: specTargets.count * 3) { [specTargets] i in + errState.catchError { let mIdx = i / 3 let proj = i % 3 let info = specTargets[mIdx] @@ -301,7 +303,9 @@ public class SwitchGLU: Module, @unchecked Sendable { MLXFast.preadIntoOffset(self._stackedDown!, safetensorsPath: downSSD.path, tensorName: downSSD.tensorName, expertIndex: UInt32(info.expertId), dstOffset: info.slot * downBpe) } + } } + errState.check() } if idx.size == 0 { @@ -372,7 +376,9 @@ public class SwitchGLU: Module, @unchecked Sendable { if !missesNeedingPread.isEmpty { let bpe = _stackedBytesPerExpert let downBpe = _stackedDownBytesPerExpert + let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: missesNeedingPread.count * 3) { [missesNeedingPread] i in + errState.catchError { let mIdx = i / 3 let proj = i % 3 let info = missesNeedingPread[mIdx] @@ -399,7 +405,9 @@ public class SwitchGLU: Module, @unchecked Sendable { MLXFast.preadIntoOffset(self._stackedDown!, safetensorsPath: downSSD.path, tensorName: downSSD.tensorName, expertIndex: UInt32(info.expertId), dstOffset: info.slot * downBpe) } + } } + errState.check() } _previousExpertIds = ranges.map { $0.id } @@ -596,7 +604,9 @@ public class SwitchGLU: Module, @unchecked Sendable { // Full concurrent pread (baseline path) let totalReads = ranges.count * 3 + let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: totalReads) { [ranges] i in + errState.catchError { let expertIdx = i / 3 let projIdx = i % 3 let r = ranges[expertIdx] @@ -611,7 +621,9 @@ public class SwitchGLU: Module, @unchecked Sendable { MLXFast.preadInto(self._persistentDown![expertIdx], safetensorsPath: downSSD.path, tensorName: downSSD.tensorName, expertIndex: UInt32(r.id)) } + } } + errState.check() // Store routing for next token's predictions _previousExpertIds = ranges.map { $0.id } @@ -640,7 +652,9 @@ public class SwitchGLU: Module, @unchecked Sendable { if let prevIds = _previousExpertIds { let specCount = min(prevIds.count, maxBuffers) let specReads = specCount * 3 + let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: specReads) { i in + errState.catchError { let slot = i / 3 let proj = i % 3 let expertId = prevIds[slot] @@ -655,7 +669,9 @@ public class SwitchGLU: Module, @unchecked Sendable { MLXFast.preadInto(self._persistentDown![slot], safetensorsPath: downSSD.path, tensorName: downSSD.tensorName, expertIndex: UInt32(expertId)) } + } } + errState.check() } // Sync on idx (blocks until GPU finishes attention + router) @@ -719,7 +735,9 @@ public class SwitchGLU: Module, @unchecked Sendable { // Pread only misses (~30% of experts, ~6 reads at QD=6) if !missInfo.isEmpty { let totalMissReads = missInfo.count * 3 + let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: totalMissReads) { [missInfo] i in + errState.catchError { let mIdx = i / 3 let proj = i % 3 let info = missInfo[mIdx] @@ -740,7 +758,9 @@ public class SwitchGLU: Module, @unchecked Sendable { tensorName: downSSD.tensorName, expertIndex: UInt32(info.expertId)) } + } } + errState.check() } } else { // No predictions available — full pread fallback @@ -750,7 +770,9 @@ public class SwitchGLU: Module, @unchecked Sendable { usedDown.append(_persistentDown![i]) } let totalReads = ranges.count * 3 + let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: totalReads) { [ranges] i in + errState.catchError { let expertIdx = i / 3 let projIdx = i % 3 let r = ranges[expertIdx] @@ -765,7 +787,9 @@ public class SwitchGLU: Module, @unchecked Sendable { MLXFast.preadInto(self._persistentDown![expertIdx], safetensorsPath: downSSD.path, tensorName: downSSD.tensorName, expertIndex: UInt32(r.id)) } + } } + errState.check() } // Update routing for next token's predictions @@ -814,7 +838,9 @@ public class SwitchGLU: Module, @unchecked Sendable { // Concurrent pread (same as fast path) let totalReads = ranges.count * 3 + let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: totalReads) { [ranges] i in + errState.catchError { let expertIdx = i / 3 let projIdx = i % 3 let r = ranges[expertIdx] @@ -829,7 +855,9 @@ public class SwitchGLU: Module, @unchecked Sendable { MLXFast.preadInto(downBuffers[expertIdx], safetensorsPath: downSSD.path, tensorName: downSSD.tensorName, expertIndex: UInt32(r.id)) } + } } + errState.check() // Lazy compute (no eval — next layer forces it) let xGate = qGate.computeExperts(x, buffers: gateBuffers, ranges: ranges) diff --git a/Tests/MLXLMTests/CorruptSafetensorsTests.swift b/Tests/MLXLMTests/CorruptSafetensorsTests.swift new file mode 100644 index 000000000..05ee71e0e --- /dev/null +++ b/Tests/MLXLMTests/CorruptSafetensorsTests.swift @@ -0,0 +1,47 @@ +import Foundation +import MLX +import MLXLMCommon +import Testing + +@Suite +struct CorruptSafetensorsTests { + @Test + func testDeadlock() throws { + let path = "dummy_corrupt.safetensors" + let dict: [String: MLXArray] = ["weight": MLXArray.zeros([256, 1024])] + try MLX.save(arrays: dict, url: URL(fileURLWithPath: path)) + + let fd = open(path, O_RDWR) + ftruncate(fd, 100) + close(fd) + + let dst = MLXArray.zeros([256, 1024]) + dst.eval() + + final class ThreadSafeError: @unchecked Sendable { + let lock = NSLock() + var error: Error? + func catchError(_ block: () throws -> Void) { + do { + try MLX.withError { + try block() + } + } catch { + lock.withLock { + if self.error == nil { self.error = error } + } + } + } + } + + let errState = ThreadSafeError() + + DispatchQueue.concurrentPerform(iterations: 16) { i in + errState.catchError { + MLXFast.preadIntoOffset(dst, safetensorsPath: path, tensorName: "weight", expertIndex: UInt32(i), dstOffset: i * 1024 * 4) + } + } + + #expect(errState.error != nil) + } +} From a7d3f5067a63f095b30b63b654d6b1e101c58295 Mon Sep 17 00:00:00 2001 From: Aegis-AI Date: Sun, 26 Apr 2026 22:28:25 -0700 Subject: [PATCH 2/3] test: address copilot review feedback for safetensors deadlock test --- .../MLXLMTests/CorruptSafetensorsTests.swift | 40 +++++++++---------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/Tests/MLXLMTests/CorruptSafetensorsTests.swift b/Tests/MLXLMTests/CorruptSafetensorsTests.swift index 05ee71e0e..a67ef4fce 100644 --- a/Tests/MLXLMTests/CorruptSafetensorsTests.swift +++ b/Tests/MLXLMTests/CorruptSafetensorsTests.swift @@ -7,33 +7,33 @@ import Testing struct CorruptSafetensorsTests { @Test func testDeadlock() throws { - let path = "dummy_corrupt.safetensors" + let tempDir = FileManager.default.temporaryDirectory + let pathUrl = tempDir.appendingPathComponent("dummy_corrupt_\(UUID().uuidString).safetensors") + let path = pathUrl.path + let dict: [String: MLXArray] = ["weight": MLXArray.zeros([256, 1024])] - try MLX.save(arrays: dict, url: URL(fileURLWithPath: path)) + try MLX.save(arrays: dict, url: pathUrl) + + defer { + try? FileManager.default.removeItem(at: pathUrl) + } let fd = open(path, O_RDWR) - ftruncate(fd, 100) - close(fd) + #expect(fd != -1) + guard fd != -1 else { return } + + defer { + let closeResult = close(fd) + #expect(closeResult == 0) + } + + let truncateResult = ftruncate(fd, 100) + #expect(truncateResult == 0) + guard truncateResult == 0 else { return } let dst = MLXArray.zeros([256, 1024]) dst.eval() - final class ThreadSafeError: @unchecked Sendable { - let lock = NSLock() - var error: Error? - func catchError(_ block: () throws -> Void) { - do { - try MLX.withError { - try block() - } - } catch { - lock.withLock { - if self.error == nil { self.error = error } - } - } - } - } - let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: 16) { i in From 2d6b2b42f185dfaa07ea19ebdd7c1a49696ab826 Mon Sep 17 00:00:00 2001 From: Aegis-AI Date: Sun, 26 Apr 2026 23:13:09 -0700 Subject: [PATCH 3/3] test: fix swift 6 strict concurrency Sendable error in safetensors test --- Tests/MLXLMTests/CorruptSafetensorsTests.swift | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/Tests/MLXLMTests/CorruptSafetensorsTests.swift b/Tests/MLXLMTests/CorruptSafetensorsTests.swift index a67ef4fce..446501bca 100644 --- a/Tests/MLXLMTests/CorruptSafetensorsTests.swift +++ b/Tests/MLXLMTests/CorruptSafetensorsTests.swift @@ -33,12 +33,17 @@ struct CorruptSafetensorsTests { let dst = MLXArray.zeros([256, 1024]) dst.eval() + + struct SendableArray: @unchecked Sendable { + let array: MLXArray + } + let sendableDst = SendableArray(array: dst) let errState = ThreadSafeError() DispatchQueue.concurrentPerform(iterations: 16) { i in errState.catchError { - MLXFast.preadIntoOffset(dst, safetensorsPath: path, tensorName: "weight", expertIndex: UInt32(i), dstOffset: i * 1024 * 4) + MLXFast.preadIntoOffset(sendableDst.array, safetensorsPath: path, tensorName: "weight", expertIndex: UInt32(i), dstOffset: i * 1024 * 4) } }