Does using fori_loop accelerate training performance?

[Hojunian]

I recently started to use JAX and implemented a simple linear regression code using a 3-layer MLP. Following "Performance Considerations", I tried various methods and found that using a functional training loop (jax.jit and nnx.split/merge) yields the fastest performance. (114 seconds for 1M iterations on an RTX 3080 GPU)
Then I tried using nnx.fori_loop and got an interesting result--It is much faster than running a Python for loop with a JAX-complied function (31 seconds for 1M iterations on an RTX 3080 GPU).

Question:

Is it common practice to use nnx.fori_loop for faster training? or might there be other factors contributing to this performance boost?

Code

Using Python for loop:

from flax import nnx
import optax
import jax
import jax.numpy as jnp
import time

class MLP(nnx.Module):
    def __init__(self, *, dims, rngs: nnx.Rngs):
        assert len(dims) >= 2
        self.layers = []
        depth = len(dims) - 1
        for i in range(depth - 1):
            self.layers.append(nnx.Linear(dims[i], dims[i + 1], rngs=rngs))
            self.layers.append(nnx.gelu)
        self.layers.append(nnx.Linear(dims[depth - 1], dims[depth], rngs=rngs))

    def __call__(self, x: jax.Array):
        for layer in self.layers:
            x = layer(x)

        return x


@nnx.jit
def train_step(model, optimizer, gt_weights, gt_bias, key):
    key, subkey = jax.random.split(key)
    x = jax.random.normal(subkey, shape=(64, 2))

    y_gt = (
        jnp.dot(x, gt_weights)
        + gt_bias
        + jax.random.normal(subkey, shape=(64, 1)) * 0.1
    )

    def loss_fn(model: MLP):
        y_pred = model(x)
        return jnp.mean((y_gt - y_pred) ** 2)

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(grads)

    return loss, key

@jax.jit
def jax_train_step(graphdef, state, gt_weights, gt_bias, key):
    model, optimizer = nnx.merge(graphdef, state)
    key, subkey = jax.random.split(key)
    x = jax.random.normal(subkey, shape=(64, 2))

    y_gt = (
        jnp.dot(x, gt_weights)
        + gt_bias
        + jax.random.normal(subkey, shape=(64, 1)) * 0.1
    )

    def loss_fn(model: MLP):
        y_pred = model(x)
        return jnp.mean((y_gt - y_pred) ** 2)

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(grads)

    state = nnx.state((model, optimizer))
    return loss, key, state

model = MLP(dims=[2, 64, 64, 1], rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3))

key = jax.random.key(0)

key, subkey = jax.random.split(key)
gt_weights = jax.random.normal(subkey, shape=(2, 1))

key, subkey = jax.random.split(key)
gt_bias = jax.random.normal(subkey, shape=(1,))

# cached_train_step = nnx.cached_partial(train_step, model, optimizer, gt_weights, gt_bias)
graphdef, state = nnx.split((model, optimizer))

start = time.time()
for i in range(100000):
    # # Method 1: using naive nnx implemetation
    # loss, key = train_step(model, optimizer, gt_weights, gt_bias, key)
    
    # # Method 2: using partial cached function, faster
    # loss, key = cached_train_step(key)
    
    # Method 3: using nnx.split and jax.jit, fastest
    loss, key, state = jax_train_step(graphdef, state, gt_weights, gt_bias, key)

print(time.time() - start)
print(loss)

Using nnx.fori_loop:

class MLP(nnx.Module):
    def __init__(self, *, dims, rngs: nnx.Rngs):
        assert len(dims) >= 2
        self.layers = []
        depth = len(dims) - 1
        for i in range(depth - 1):
            self.layers.append(nnx.Linear(dims[i], dims[i + 1], rngs=rngs))
            self.layers.append(nnx.gelu)
        self.layers.append(nnx.Linear(dims[depth - 1], dims[depth], rngs=rngs))

    def __call__(self, x: jax.Array):
        for layer in self.layers:
            x = layer(x)

        return x


def train_step(i, input):
    model, optimizer, gt_weights, gt_bias, loss, key = input
    key, subkey = jax.random.split(key)
    x = jax.random.normal(subkey, shape=(64, 2))
    y = (
        jnp.dot(x, gt_weights)
        + gt_bias
        + jax.random.normal(subkey, shape=(64, 1)) * 0.1
    )

    def loss_fn(model: MLP):
        y_pred = model(x)
        return jnp.mean((y - y_pred) ** 2)

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(grads)

    return model, optimizer, gt_weights, gt_bias, loss, key


@nnx.jit
def run(model, optimizer, gt_weights, gt_bias, loss, key):
    model, optimizer, __, __, loss, key = nnx.fori_loop(
        0, 100000, train_step, (model, optimizer, gt_weights, gt_bias, loss, key)
    )

    return model, loss


model = MLP(dims=[2, 64, 64, 1], rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3))

key = jax.random.key(0)

key, subkey = jax.random.split(key)
gt_weights = jax.random.normal(subkey, shape=(2, 1))

key, subkey = jax.random.split(key)
gt_bias = jax.random.normal(subkey, shape=(1,))

loss = 0.0

start = time.time()
*_, loss = run(model, optimizer, gt_weights, gt_bias, loss, key)
print(time.time() - start)
print(loss)