[JAX] Support Flax sharding constraints

[jberchtold-nvidia]

Description

Updates our with_sharding_constraint_by_logical_axes helper function to support and prefer Flax logical axis rules when they exist in the current context. If no Flax logical axes rules exist in the current context, it will fall back to TE's hardcoded logical axes, though this functionality is no deprecated and will be removed in the future.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Update with_sharding_constraint_by_logical_axes to check if Flax logical axis rules exist in the current context, and if so call the sharding constraint with flax.linen.with_logical_constraint
• Update unfused attention in transformers.py to remove duplicate SEQLEN_AXIS usage in a single sharding constraint. This was not flagged before since our TE hardcoded axis system didn't check for this and it was always mapped to None/replicated.
• Corrected logical axis rule context setup in encoder examples

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[jberchtold-nvidia]

/te-ci L0 L1

[jberchtold-nvidia]

/te-ci L0 L1

[jberchtold-nvidia]

/te-ci L0 L1

[jberchtold-nvidia]

/te-ci L0 L1 jax

[jberchtold-nvidia]

/te-ci L0 L1

[phu0ngng]

/te-ci L1

[phu0ngng]

Hi,
I'm curious to learn what the difference is between having the jax.sharding.Mesh context before vs after the te.fp8_autocast context.

[jberchtold-nvidia]

Afaik it doesn't matter the order of the Mesh vs. fp8_autocast, as long as both are before your model. But the order of te_flax.extend_logical_axis_rules and fp8_autocast does matter according to the docstring of the former.

So the ordering needs to be:

1. fp8_autocast context
2. Create Flax logical axis rule context. To support TE's hardcoded axis system in Flax (which UnfusedDotProductAttention requires), we extend the logical access rules with TE's rule table via te_flax.extend_logical_axis_rules, which must be inside an fp8_autocast.
3. Create and initialize the model, training loop, etc.

Afaik, the Mesh can come anywhere before item 3.

I just pulled up the fp8_autocast to the top and merged the Mesh with the with block for the logical axis rule context in item 2 to reduce indentation. But if a smaller diff is preferred, I can do Mesh -> fp8_autocast -> logical axis rules