[JAX] Migrate EP collective-stream annotation to region-based compute_on - #3304
Conversation
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
|
Pipeline #60891724 with the nightly JAX image. |
Greptile SummaryThe PR migrates JAX expert-parallel operations to the region-based
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains within the eligible follow-up review scope. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[MoE input] --> B[EP dispatch compute_on region]
B --> C[Collective-stream async wrapper]
C --> D[Per-expert MLP]
D --> E[EP combine compute_on region]
E --> F[Output constrained to input logical axes]
Reviews (3): Last reviewed commit: "Merge branch 'main' into jax_compute_on_..." | Re-trigger Greptile |
jberchtold-nvidia
left a comment
jberchtold-nvidia
left a comment
There was a problem hiding this comment.
LGTM, thanks!
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
|
pipeline #61310617 with the nightly JAX image. |
phu0ngng
changed the title
| def _on_collective_stream(func): | ||
| """Pin ``func``'s ops to XLA's collective stream so the scheduler serializes | ||
| them with native collectives. No-op on JAX that lacks the annotation.""" | ||
| if not is_collective_stream_supported(): |
There was a problem hiding this comment.
Can we update is_collective_stream_supported's definition to attempt to import compute_on in a try/catch, and return False if an import error occurs? Since this compute_on is experimental, it seems JAX may freely rename this function without warning, so any future JAX versions could cause TE to break until we update TE
If we instead tried to import compute_on as part of the is_collective_stream_supported check, we'd just skip the compute_on logic but TE would still work and import correctly. We could log a warning in this case
We could even go a step further and wrap the body of _on_collective_stream in a try/catch, so if something subtler happens like JAX keeps compute_on as the same function name but changes the params, we would try/catch then log a warning and fallback to the original function without the compute_on annotation.
Since this can fail at import-time, it affects all of TE usage even for modules unrelated to TE EP, so I think try/catch'ing at least the first import case would be helpful
What do you think?
There was a problem hiding this comment.
Can be in follow-up
| out_partition_spec=out_partition_spec, | ||
| ) | ||
| # output of MLP should be sharded the same way as the activation input | ||
| output = with_sharding_constraint_by_logical_axes(output, input_axes) |
There was a problem hiding this comment.
input_axes and out_partition_spec should be equivalent here, right? If so, can we simplify this be replacing out_partition_spec above with input_axes
There was a problem hiding this comment.
Can be in follow-up
2 participants
Description
JAX rewrote
jax.experimental.compute_on.compute_oninto a region-based transform. This migrates_on_collective_stream, which usescompute_onto the new signature.Type of change
Changes
_on_collective_streamcallscompute_on(compute_type="gpu_stream:collective", out_memory_spaces=Device)on a nullary thunk that closes over the staticEpLayerConfig/PartitionSpecargs, since region-basedcompute_onabstract-evals every argument, then invokes the annotated call.compute_onmoves_xla_stream_annotationonto the async wrapper XLA generates, sotest_z_dispatch_combine_on_collective_streamchecks each EP call is reachable from a collective-annotated wrapper.compute_onandjax.memory.Spaceexist since JAX 0.10.1, the current collective-stream floor, so no min-version bump is needed.Checklist: