MLX and determinism

[billziss-gh]

Is the MLX framework deterministic if the PRNGs are seeded?

For debugging purposes I initialize my PRNG's as follows:

random.seed(value)
numpy.random.seed(value)
mlx.core.random.seed(value)

And then I get the following outputs from multiple training runs with the same seed:

best 80  tloss 0.006779   vloss 0.016929
curr 99  tloss 0.003995   vloss 0.028343

best 80  tloss 0.006801   vloss 0.016929
curr 99  tloss 0.003992   vloss 0.028244

best 80  tloss 0.006802   vloss 0.016943
curr 99  tloss 0.003979   vloss 0.028244

best 80  tloss 0.006772   vloss 0.016939
curr 99  tloss 0.003982   vloss 0.028248

best 80  tloss 0.006786   vloss 0.016937
curr 99  tloss 0.003990   vloss 0.028277

Notice that the best/last training/validation losses are quite similar, but not quite the same. For example, for the last two runs the best validation losses were 0.016939 and 0.016937 respectively. Close but not quite.

Is this a known issue? I may have missed some other source of non-determinism in my network, but my investigation appears to show that the non-determinism comes from within the MLX framework.

System: Apple M4 Max, 128 GB, Tahoe 26.0.1
Version: mlx==0.30.0

[awni]

No in general it's not deterministic. I would say it's 99% deterministic but there are some GPU kernels which can add floating point numbers in different orders and that can cause small numerical differences.

[billziss-gh]

Thank you for confirming. I will mark your response as the answer.

I am assuming there is no way to force it to be deterministic (even at the expense of performance) for the purposes of testing?

[awni]

An option is to run on the CPU (I believe it should be deterministic). But it will be quite slow so you'd need to run a much smaller model.

[billziss-gh]

Understood. Thank you.