pspec error when run Maxdiffusion with new JAX version

[hx89]

We started to have pspec error when run Maxdiffusion with new JAX version like 0.6:

File "/opt/maxdiffusion/src/maxdiffusion/train_flux.py", line 36, in train
trainer.start_training()
File "/opt/maxdiffusion/src/maxdiffusion/trainers/flux_trainer.py", line 138, in start_training
p_train_step = self.compile_train_step(pipeline, params, train_states, state_shardings, data_shardings)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/maxdiffusion/src/maxdiffusion/trainers/flux_trainer.py", line 309, in compile_train_step
p_train_step = p_train_step.lower(train_states[FLUX_STATE_KEY], dummy_batch, train_rngs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: pspec PartitionSpec(('data', 'fsdp', 'tensor'), None) contains a manual axes ('data', 'fsdp', 'tensor') of mesh which is not allowed. If you are using a with_sharding_constraint under a shard_map, only use the mesh axis in PartitionSpec which are not manual.

It's likely due to API change in new JAX and Maxdiffusion may need to change accordingly.

[entrpn]

Fixed in #206.

Please try it again.

[hx89]

Thanks @entrpn for the fix!

I tested the Flux training on Maxdiffusion main and didn't see the original pspec error during lowering, but I got NCCL OOM error in the training loop, not sure if it could be related to the fix or it's a separate issue?

[hx89]

@entrpn Instead of using docker_build_dependency_image.sh that gave collective OOM error, I tried to build a new container using my dockerfile: https://github.com/hx89/maxdiffusion/blob/gpu_dockerfile/gpu_docker/Dockerfile
and it still had the pspec error:

File "/opt/maxdiffusion/src/maxdiffusion/models/attention_flax.py", line 347, in _cudnn_flash_attention
out = wrap_flash_attention(query, key, value)
File "/opt/maxdiffusion/src/maxdiffusion/models/attention_flax.py", line 345, in wrap_flash_attention
return jax.vmap(dpa_layer)(query, key, value, mask=None)
File "/usr/local/lib/python3.12/dist-packages/transformer_engine/jax/flax/transformer.py", line 687, in call
x = _FusedDotProductAttention(
File "/usr/local/lib/python3.12/dist-packages/transformer_engine/jax/flax/transformer.py", line 360, in call
x = fused_attn(
File "/usr/local/lib/python3.12/dist-packages/transformer_engine/jax/attention.py", line 1094, in fused_attn
return _legacy_fused_attn(
File "/usr/local/lib/python3.12/dist-packages/transformer_engine/jax/attention.py", line 761, in _legacy_fused_attn
output = _fused_attn(
File "/usr/local/lib/python3.12/dist-packages/transformer_engine/jax/attention.py", line 916, in _fused_attn_fwd_rule
output, softmax_aux, rng_state = tex.fused_attn_fwd(
File "/usr/local/lib/python3.12/dist-packages/transformer_engine/jax/cpp_extensions/attention.py", line 2659, in fused_attn_fwd
rng_state = with_sharding_constraint(rng_state, PartitionSpec(get_all_mesh_axes(), None))
File "/usr/local/lib/python3.12/dist-packages/transformer_engine/jax/sharding.py", line 158, in with_sharding_constraint
return jax.lax.with_sharding_constraint(x, pspec)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: pspec PartitionSpec(('data', 'fsdp', 'tensor'), None) contains a manual axes ('data', 'fsdp', 'ten
sor') of mesh which is not allowed. If you are using a with_sharding_constraint under a shard_map, only use the mesh axis in PartitionSpec which are not manual.

Could you please help double check if I missed anything?

[hx89]

The issue is from TE and is fixed in NVIDIA/TransformerEngine#2069, so close this issue