Is it normal that different CPUs output different deterministic results for the Same Code?

[YousefElbrolosy]

Hi, I have a program that uses jax, flax and optax.

I ran this code on CPU, using three different CPUs. Two MacOs Systems (one on Sequoia (M1 Pro), other on Sonoma (M2)) and one on a linux system.

All three systems output different results for the same output, however they output that output deterministically.

Is there any solution to this? Is this a bug in jax, or is this related to XLA?

If there are any solutions to this I would be grateful.

[dfm]

It's hard to give a concrete answer without a minimal reproducible example, but it's not unexpected that different platforms could exhibit deterministically different numerics. And these can even result in large errors in absolute terms when accumulated. Without more details, it's not obvious that this is a bug, but it could be! If you can narrow it down and produce a minimal pure-JAX example of these errors, I'd be happy to help diagnose.

[YousefElbrolosy]

Hello, thanks for your reply!

Here is as narrow as I could make it.

import jax
import optax
import flax.linen as nn
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

variables = jnp.array([0.1, -3 * jnp.pi / 2])


class RNN(nn.Module):
    hidden_size: int
    output_size: int

    @nn.compact
    def __call__(self, input, hidden_state):
        gru_cell = nn.GRUCell(features=self.hidden_size)
        new_hidden_state, _ = gru_cell(hidden_state, input)
        output = nn.Dense(features=self.output_size)(new_hidden_state)
        return output, new_hidden_state

def _optimize(
    loss_fn,
    init_params,
    max_iter,
    learning_rate,
):
    optimizer = optax.adam(learning_rate)
    opt_state = optimizer.init(init_params)

    @jax.jit
    def step(params, state):
        grads = jax.grad(loss_fn)(params)
        updates, new_state = optimizer.update(grads, state, params)
        new_params = optax.apply_updates(params, updates)
        return new_params, new_state

    params = init_params
    for iter_idx in range(max_iter):
        params, opt_state = step(params, opt_state)
    return params, iter_idx + 1

def fun(gamma, delta):
    op = jnp.array([[0, -1j], [1j, 0]])
    angle = (gamma * op) + delta / 2
    return (jax.scipy.linalg.expm(1j * angle) + jax.scipy.linalg.expm(-1j * angle)) / 2

def loss(params):
    rnn = RNN(hidden_size=10, output_size=2)
    input = variables
    hidden_state = jnp.zeros((10,))
    output, _ = rnn.apply({'params': params}, input, hidden_state)
    params_out = output
    return jnp.real(jnp.trace(fun(params_out[0], params_out[1])))

if __name__ == "__main__":
    rng = jax.random.PRNGKey(0)
    rnn = RNN(hidden_size=10, output_size=2)
    input = variables
    hidden_state = jnp.zeros((10,))
    params = rnn.init(rng, input, hidden_state)['params']

    max_iter = 100
    learning_rate = 0.01
    convergence_threshold = 1e-6
    optimized_params, num_iterations = _optimize(
        loss,
        params,
        max_iter,
        learning_rate,
    )
    final_loss = loss(optimized_params)
    print("Final Loss:", final_loss)

This outputs on a macos system:

-1.9979573829398634

and on a linux system:

-1.9979573808129485

Differing in the last 8 digits, and I suppose given a much larger complicated system this difference can be quite large.

Also note that sometimes differences between two devices can be larger than between two other devices.

[jakevdp]

I suspect the difference here is stemming from linear algebra libraries. JAX on CPU uses your system LAPACK, and different systems use different LAPACK implementations.

[YousefElbrolosy]

Yes, this makes sense, and I indeed tried running this reproducible code on two mac systems and they output the same number, however the problem I was having was that they were inconsistent on two different mac systems as well - not reproducible by this code-. is it possible that different versions of the operating system may cause this?

[YousefElbrolosy]

@jakevdp I tried this again, and when I removed @jax.jit decorator both platforms led to almost the same numbers. So, the issue is not in lapack but rather jiting?

Macos: Sonoma
-1.9979573823846777

Linux
-1.997957382384678

[jakevdp]

This is unrelated to jax.jit – you're getting different results on different platforms (both with and without JIT) because the different platforms are executing using different LAPACK libraries.

[YousefElbrolosy]

Okay thank you!