[JAX] Support axes orthogonal to EP in ep_bootstrap via mesh-derived domain grouping

[phu0ngng]

Support axes orthogonal to EP in ep_bootstrap via mesh-derived domain grouping

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

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

[greptile-apps]

Greptile Summary

This PR fixes ep_bootstrap so that JAX meshes with axes orthogonal to EP (tp, pp, cp, …) are correctly handled. The old code derived the NCCL communicator root with rank // ep_size, which only works when processes are laid out as (dp, ep); the new _ep_domain_for_rank helper moves the ep axis to the last position, reshapes the rank grid into (num_domains, ep_size), and locates each rank's domain root and intra-group column directly from the mesh — no assumed axis ordering.

• _ep_domain_for_rank (ep.py lines 74–99): reads process ranks from mesh.devices, rearranges the ep axis last, and returns (root_rank, rank_within_group, num_domains) for any mesh shape.
• ep_bootstrap validation (ep.py lines 154–164): replaces the num_ep_groups * ep_size == world_size check with get_num_devices_in_mesh(mesh) == world_size, which is valid regardless of how many orthogonal axes exist.
• TestEpDomainGrouping (test_multi_process_ep.py): single-process unit test using a 4×2 (expert×tensor) mock mesh; asserts per-rank ndom inside the loop and verifies the resulting domain partition.

Confidence Score: 5/5

Safe to merge — the new mesh-driven domain grouping is correct for all axis layouts, the old dp-only path is a strict special case, and the tightened world_size validation makes mismatches explicit.

The core change replaces a simple arithmetic assumption with a general mesh traversal that handles any number of orthogonal axes. The single-process unit test exercises a concrete 4x2 mesh and asserts ndom inside the loop. No existing call sites are broken because the new path is a superset of the old one.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/jax/ep.py	Adds _ep_domain_for_rank to compute NCCL-domain membership from mesh layout; replaces the old rank // ep_size arithmetic in ep_bootstrap so orthogonal axes like tp/pp/cp are handled correctly. Validation tightened to compare get_num_devices_in_mesh(mesh) against world_size.
tests/jax/test_multi_process_ep.py	Adds TestEpDomainGrouping — a single-process test using a mock device_to_rank mapping to exercise _ep_domain_for_rank in a 4×2 (expert×tensor) mesh, with ndom asserted inside the loop.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant P as Process (rank r)
    participant D as _ep_domain_for_rank
    participant N as libnccl
    participant G as _allgather_uid (all ranks)
    participant B as set_ep_bootstrap_params

    P->>D: mesh, ep_resource, rank
    D-->>P: root_rank, rank_within_group, num_domains

    alt "rank_within_group == 0"
        P->>N: ncclGetUniqueId()
        N-->>P: uid_bytes (128 B)
    else
        P->>P: "uid_bytes = zeros(128)"
    end

    P->>G: uid_arr (local 128 B)
    G-->>P: all_uids[world_size x 128]
    P->>P: "uid_bytes = all_uids[root_rank]"

    P->>B: uid_bytes, ep_size, rank_within_group

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant P as Process (rank r)
    participant D as _ep_domain_for_rank
    participant N as libnccl
    participant G as _allgather_uid (all ranks)
    participant B as set_ep_bootstrap_params

    P->>D: mesh, ep_resource, rank
    D-->>P: root_rank, rank_within_group, num_domains

    alt "rank_within_group == 0"
        P->>N: ncclGetUniqueId()
        N-->>P: uid_bytes (128 B)
    else
        P->>P: "uid_bytes = zeros(128)"
    end

    P->>G: uid_arr (local 128 B)
    G-->>P: all_uids[world_size x 128]
    P->>P: "uid_bytes = all_uids[root_rank]"

    P->>B: uid_bytes, ep_size, rank_within_group

Loading

_{Reviews (3): Last reviewed commit: "[JAX] Support axes orthogonal to EP in e..." | Re-trigger Greptile}

[phu0ngng]

/te-ci JAX L1