MLX error: [Event::wait] Timed out at array.cpp

[dexterleng]

I ran into a timed out error and my M1 14" 16 RAM was freezing while using MLXEmbeddings to generate embeddings for 500 strings (100 words each):

MLX error: [Event::wait] Timed out at /Users/macintosh/Library/Developer/Xcode/DerivedData/ProNotes-gymsstwowowjulfsqthohdwcworn/SourcePackages/checkouts/mlx-swift/Source/Cmlx/include/mlx/c/array.cpp:290

Apologies if this is a really basic or silly question but how do I prevent this from happening? Also how do I prevent it from slowing my computer down to a crawl? Does MLX have some kind of throttling?

Code:

    func embed() async throws {
        // Helper function to generate a random word
        func randomWord() -> String {
            let letters = "abcdefghijklmnopqrstuvwxyz"
            let length = Int.random(in: 3...10) // Random word length between 3 and 10
            return String((0..<length).map { _ in letters.randomElement()! })
        }
        
        // Generate 500 strings, each with 100 random words
        let texts = (1...500).map { _ in
            (1...100).map { _ in randomWord() }.joined(separator: " ")
        }
        
        let modelContainer = try await MLXEmbedders.loadModelContainer(
            configuration: ModelConfiguration.nomic_text_v1_5)
                
        let result = await modelContainer.perform {
            (model: EmbeddingModel, tokenizer, pooling) -> [[Float]] in
            let inputs = texts.map {
                tokenizer.encode(text: $0, addSpecialTokens: true)
            }
            
            // Pad to longest
            let maxLength = inputs.reduce(into: 16) { acc, elem in
                acc = max(acc, elem.count)
            }

            let padded = stacked(
                inputs.map { elem in
                    MLXArray(
                        elem
                            + Array(
                                repeating: tokenizer.eosTokenId ?? 0,
                                count: maxLength - elem.count))
                })
            let mask = (padded .!= tokenizer.eosTokenId ?? 0)
            let tokenTypes = MLXArray.zeros(like: padded)
            let result = pooling(
                model(padded, positionIds: nil, tokenTypeIds: tokenTypes, attentionMask: mask),
                normalize: true, applyLayerNorm: true
            )
            result.eval()
            return result.map { $0.asArray(Float.self) }
        }
        
        print(result)
    }

[davidkoski]

First the question of the timeout: you would need to see what it was doing at the time. It is possible that it took more physical memory than you had and you got into swap. You can use something like Instruments to measure it or you can just instrument your code with something like this:

• https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/gpu/snapshot

You might also take a look at the size of your inputs / padded -- that looks like a lot of tokens. I am not sure what is reasonable here (@awni ?) but perhaps smaller batches would work? You could try that and see if it fixes the hang (and potential memory) problem.

Apologies if this is a really basic or silly question but how do I prevent this from happening? Also how do I prevent it from slowing my computer down to a crawl? Does MLX have some kind of throttling?

No throttling -- this is an API and it is up to the application developer to implement throttling or batching/batch-splitting as needed.

Give that a look and see if it helps!