Help Troubleshooting

[JimWallace]

Hi,

I'm implementing t-SNE using MLX as a learning exercise. I have setup a simple gradient descent loop using the following code:

func gradientDescent(ldEmbeddings: inout MLXArray, hdAffinities: MLXArray, learningRate: Float, nIters: Int, patience: Int) {
               
        //let computeGradient = grad(klDivergence)
        let computeValueAndGradient = valueAndGrad(loss)
        var bestLoss: Float = .greatestFiniteMagnitude
        var earlyStopCounter = 0
        
        for _ in 0..<nIters {
            // Compute the gradient
            let (value, gradient) = computeValueAndGradient([ldEmbeddings,hdAffinities])
            
            let v = value[0].asArray(Float.self)[0]
            if v < bestLoss {
                bestLoss = v
                earlyStopCounter = 0
            } else {
                earlyStopCounter += 1
            }
                                
            // Update ldEmbeddings using the gradient
            ldEmbeddings -= learningRate * gradient[0]
            
            if earlyStopCounter >= patience {
                break
            }
            
        }
    }

It works great! Until I feed it ~893 documents, at which point I'm getting "EXC_BAD_ACCESS (code=2, address=0x16ff1bff8)" from valueAndGradient (Transforms+Internal, line 14) (plus about 6,000 calls down the stack into C land). The fact that the Swift call ends with a ! makes me wonder if this is something that we should be handling that snuck through the cracks, or maybe it's just me.

I'm new to this, trying to understand how to best debug from here -- can anyone provide me a nudge in the right direction? (Not sure how to create a minimal example without sharing all the data too ... but I've tried with multiple data sets and the error pops up at the same number of documents for different data sets... works from 3 all the way up to about 893, then crashes).

Thanks!

[davidkoski]

I am not sure the shape of value, but I wonder if .item(Float.self) would be more convenient here? item() requires that it have a single value but it is a little less cumbersome than converting to an array and then indexing that.

let v = value[0].asArray(Float.self)[0]

Anyway, to your question: a segmentation fault with a stack depth of 6000 sounds like a stack overflow trying to either evaluate or deallocate a very large graph. There is code in mlx::core (the C++ side) to avoid these situations but we recently ran into one in the 0.18.0 release and cut 0.18.1 with the fix.

@awni or @angeloskath do you see anything that stands out here? We might need more context to reproduce this and track it down.

[JimWallace]

Thanks! .item() sounds like the right approach there -- I'll make sure to update my code!

Unfortunately for the crash, it's not getting to that line, so I think this would be unrelated?

I'm building with 0.18.1, at least according to Xcode.

[JimWallace]

I should have added, I’m computing affinities for all points on each cycle, so this is of shape [893,883].

I should be making this SGD to scale, but haven’t been able to figure out how those gradient functions work (what’s a module? And why do I need one to use the SGD optimizer?)

[JimWallace]

I also just tried writing a batched version, with shape e.g., [256,256], and I'm seeing the same behaviour.

I'm creating the ldAffiinities matrix inside the loss function, so this seems to be stemming from my passing a set of ~900x900 matrixes in to the loss function?

[davidkoski]

The ! in the call you mentioned:

    let vector_pair = mlx_closure_value_and_grad_apply(valueAndGrad, input_vector)!

is to handle the fact that the C API can't indicate if a return value is optional (and thus everything is optional). We use the ! to assert than in fact it isn't optional.

[JimWallace]

In case it helps, the final error is happening at line 92 of compressed_pair.h:

_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 reference __get() _NOEXCEPT { return *this; }

[awni]

Until I feed it ~893 documents

It's not implausible that you are hitting a stack overflow in value_and_grad. There is a recursion in there. But I also am not understanding why the size of your graph scales with the size of the training set. Those should be independent.

Perhaps you could provide more information so we can help out / reproduce?

[davidkoski]

The value is being evaluated but I wonder if the gradient is not?

[JimWallace]

I’m trying to distil this down to a smaller example, but in case it helps the test case is here:

https://git.uwaterloo.ca/jrwallace/swiftnlp/-/blob/main/Tests/SwiftNLPTests/6.%20Topic%20Modelling/BasicTopicModelTests.swift?ref_type=heads

It’s a big project, sorry!

The graph (affinity matrix) is scaling with the size of the input because it includes pairwise affinities.

[JimWallace]

Here's a way of replicating the error. It seems that it has something to do with my hdAffinities calculation, if I replace it with a sequential matrix the error goes away:

import XCTest
import Foundation
import System
import MLX
import MLXNN


final class mweTests: XCTestCase {
    
    func testMWE() async throws {
        
        let numDocuments = 900
        let initialDimension = 50
        let lowerDimension = 5
        
        // Original data
        let sequence = Array(stride(from: 0.0, to: Double(numDocuments * initialDimension), by: 1.0))
        let input: MLXArray = MLXArray(converting: sequence, [numDocuments,initialDimension])
        
        // truncate for first approximation in low dimension
        let ldEmbeddings = input[0... , 0 ..< lowerDimension]
        
        
        // Step 1: Compute Pairwise Affinities in High-Dimensional Space
        func computeHighDimAffinities(data: MLXArray) -> MLXArray {
            // Step 1: Calculate pairwise Euclidean distances
            let diff = data.expandedDimensions(axes: [1]) - data.expandedDimensions(axes: [0])
            let squaredDiff = diff.pow(2)
            let distances = sqrt(squaredDiff.sum(axis: 2))

            // Step 2: Prepare the affinities matrix
            let nPoints = distances.shape[0]
            var affinities = MLXArray.zeros([nPoints, nPoints])

            // Precompute standard deviations for Gaussian kernel
            let stdDevs = (0..<nPoints).map { 2 * pow(std(distances[$0]), 2) }

            // Compute Gaussian affinities for all pairs
            for i in 0..<nPoints {
                let std = stdDevs[i]
                let rowAffinities = exp(-distances[i] / std)

                // Set the diagonal to 0 (remove self-affinity)
                rowAffinities[i] = MLXArray(0)

                // Normalize row affinities
                let rowSum = rowAffinities.sum()
                affinities[i] = rowAffinities / rowSum
            }

            // Step 3: Symmetrize affinities
            affinities = (affinities + affinities.transposed()) / 2.0

            // Step 4: Normalize entire affinities matrix to sum to 1
            let totalSum = affinities.sum()
            affinities /= totalSum
                        
            return affinities
        }
        
        // If I just set hdAffinities to values like this, no error
        //let hdSequence = Array(stride(from: 0.0, to: Double(numDocuments * numDocuments), by: 1.0))
        //let hdAffinities = MLXArray(converting: hdSequence, [numDocuments,numDocuments])
        
        // Calling this function gives me the error
        let hdAffinities = computeHighDimAffinities(data: input)
        
        
        // Step 3: Compute Pairwise Affinities in Low-Dimensional Space
        func computeLowDimAffinities(Y: MLXArray) -> MLXArray {
            let sequence = Array(stride(from: 0.0, to: Double(Y.shape[0] * Y.shape[0]), by: 1.0))
            return MLXArray(converting: sequence, [Y.shape[0], Y.shape[0]])
        }
        
        func loss(arrays: [MLXArray]) -> [MLXArray] {                        
            return [klDivLoss(inputs: arrays[1] + 1e-10, targets: computeLowDimAffinities(Y: arrays[0]), reduction: .sum)]
        }
        
        let computeValueAndGradient = valueAndGrad(loss)
        
        for _ in 0 ..< 10 {
            
            // Compute the gradient for the current batch
            let (_, gradient) = computeValueAndGradient([ldEmbeddings, hdAffinities])
            
            ldEmbeddings -= 0.1 * gradient[0]
        }
        
    }
}

So my sense is that 1) I'm doing something wrong, but also 2) it's maybe hitting a weird case in MLX that's causing the error?

[awni]

So my sense is that 1) I'm doing something wrong, but also 2) it's maybe hitting a weird case in MLX that's causing the error?

You're not doing anything wrong per se.. it's just very inefficient to do the affinity calculation that way. We probably shouldn't crash.. but it's also ok that it served as a signal in this case that it's not a good way to go about computing the affinities.

The entire affinity computation should be vectorized to avoid blowing up the graph size and it will also be much much faster.

Replace the computeHighDimAffinities with something like the following (sorry I didn't test it, but the basic ideas are all there):

    func computeHighDimAffinities(data: MLXArray) -> MLXArray {
        // Step 1: Calculate pairwise Euclidean distances
        
        // Assuming data has shape [num examples x per example dimension]
        let dataSqSums = data.square().sum(axis: 1, keepDims: true)
        let diff = dataSqSums + dataSqSums.T - 2 * data.matmul(data.T)
        let distances = sqrt(diff)

        // Step 2: Prepare the affinities matrix
        let nPoints = distances.shape[0]

        // Precompute standard deviations for Gaussian kernel
        let stdDevs = 2 * pow(std(distances, axis: 1, keepDims: true), 2)

        // Compute Gaussian affinities for all pairs
        var affinities = exp(-distances / stdDevs)
        
        // Set diagonal to 0
        let indices = MLXArray(0 < nPoints)
        affinities[indices, indices] = MLXArray(0)

        // Normalize rows
        let rowSum = affinities.sum(axis: 1, keepDims: true)
        affinities = affinities / rowSum

        // Step 3: Symmetrize affinities
        affinities = (affinities + affinities.T) / 2.0

        // Step 4: Normalize entire affinities matrix to sum to 1
        let totalSum = affinities.sum()
        affinities /= totalSum
                    
        return affinities
    }

[JimWallace]

Thank you! Yes, that has fixed the crash!

Now I can focus on making this all work on larger data sets, I really appreciate the help!