nnx.Dropout in combination with pmap returns an error

[maxxxzdn]

System information

OS Platform and Distribution: Linux Ubuntu 22.04
flax, jax, jaxlib versions:
flax: 0.8.5
jax: 0.4.30
jaxlib: 0.4.30
Python version: 3.9.19

Problem you have encountered:

nnx.Dropout in combination with jax.pmap/nnx.pmap returns the following error:

TraceContextError: Cannot call RngStream from a different trace level

Steps to reproduce:

from flax import nnx

class Model(nnx.Module):
  def __init__(self, din, dmid, dout, rngs: nnx.Rngs):
    self.linear = nnx.Linear(din, dmid, rngs=rngs)
    self.bn = nnx.LayerNorm(dmid, rngs=rngs)
    self.dropout = nnx.Dropout(0.2, rngs=rngs)
    self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)

  def __call__(self, x):
    x = nnx.relu(self.dropout(self.bn(self.linear(x))))
    return self.linear_out(x)

model = Model(2, 64, 3, rngs=nnx.Rngs(0))
out = nnx.pmap(model)(jnp.ones((1, 16, 2)))

I saw a similar discussion for scan, is there a remedy for pmap as well?

[cgarciae]

Hey! I'm guessing you want to replicate the weights but have different RNGs, to do this you can use the nnx.split_rngs decorator to split the RNGs before entering pmap and, and use StateAxes to specify the parallelization axes for substates of your Module, in this case map RngState to 0 and the rest (...) to None:

state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})

@nnx.split_rngs(splits=1)
@nnx.pmap(in_axes=(state_axes, 0))
def forward(model, x):
  return model(x)

out = forward(model, jnp.ones((1, 16, 2)))

For more info, check out the Filters guide.

[maxxxzdn]

thank you, keep up the amazing work!