Allow tensor parallelism and context parallelism in one ParallelConfig

[whn09]

What this does

ParallelConfig.__post_init__ currently rejects a config that carries both a context_parallel_config and a tensor_parallel_config, so a model can only use one form of parallelism at a time. That caps some models well below the hardware they run on: MiniMax-H3 has 56 attention heads, so tp_degree cannot exceed 8, and on a 64-accelerator host tensor parallelism alone leaves 56 of them idle — even though the model already ships a _cp_plan.

Both configs already accept a caller-supplied mesh, documented as being for "combining CP with other parallelism strategies that share the same mesh", so the intent was there. What was missing was building that shared mesh and splitting it between the two configs.

• __post_init__ now only rejects a config that carries neither.
• ParallelConfig.setup hands each config its own dimensions: CP keeps the whole mesh, since its own setup already selects "ring" and "ulysses" out of it, while TP gets the "tp" dimension alone because it reads its degree from mesh.size().
• enable_parallelism builds one ("tp", "ring", "ulysses") mesh when both are requested.
• The Neuron tensor-parallel backend used the global rank as its shard index. That equals the rank's coordinate on the TP mesh only while that mesh spans the whole world, so sharing a mesh made it index past the end of the weight (an empty all-gather input). The generic backend beside it already reads the coordinate.

Usage is what you'd expect:

transformer.enable_parallelism(
    config=ParallelConfig(
        tensor_parallel_config=TensorParallelConfig(tp_degree=8),
        context_parallel_config=ContextParallelConfig(ulysses_degree=4),
    )
)

On the mesh dimension order

The default puts the context-parallel dimensions fastest-varying, so a CP group is a contiguous run of ranks and a TP group takes one rank out of each run. That is a default, not a universal optimum, and the comment in the code says so.

It is what the Neuron runtime requires: its all-to-all only accepts contiguous replica groups ("replica group start ranks must be multiples of world size"), while all-reduce and all-gather also accept strided ones. Only one axis of a 2-D mesh can be contiguous, and Ulysses is the axis that needs it — I first wrote this the other way round and it failed outright at 16 ranks. On multi-node CUDA the opposite order is usually preferable, since TP is the most bandwidth-hungry collective and wants to stay inside one NVLink domain. Callers who need the other layout pass mesh= on either config, which is the escape hatch that already exists; happy to flip the default or expose it as an argument if you'd rather.

Measurements

MiniMax-H3, 1344x768x124f, 30 steps, on a trn2.48xlarge. Steady-state seconds per denoising step (the first two model evaluations are discarded — one compiles, one is still warming the kernel cache):

cores	config	s/step	vs. TP=8	scaling efficiency
8	TP=8	9.285	1.00x	—
16	TP=4 x ulysses=4	5.66	1.64x	82%
32	TP=8 x ulysses=4	3.941	2.36x	59%

The paired denoise phase goes 329.7 s -> 137.8 s. Output is visually equivalent to the TP-only run at the same seed (same composition, same scene, same sharpness); Ulysses changes the attention reduction order, so it is a different sample of equal quality rather than a bit-comparable one.

Tests

• HybridParallelTesterMixin in tests/models/testing_utils/parallelism.py, next to the existing TP and CP mixins, wired into the Flux transformer tests. It needs tp_degree * ulysses_degree accelerators (4 at the degrees used) and skips below that, so a 2-device runner is unaffected.
• A Neuron torchrun worker following the _neuron_tp_worker.py convention already in the tree, since Neuron needs the "neuron" distributed backend and cannot use the NCCL spawn path. It runs at tp_degree=2 x ulysses_degree=4 and passes on a trn2.48xlarge: max_abs_diff=4.2e-05 against the single-device reference. The pre-existing _neuron_tp_worker test still passes with the shard-index change.

I have not been able to run the CUDA mixin — I only have Trainium hardware here — so that half needs a 4-GPU runner. Draft for that reason.

Related

Follow-up, not in this PR: from_pretrained(..., parallel_config=...) in #14609 shards weights as it reads them and so cannot call enable_parallelism afterwards, which means it applies TP but not the CP hooks. Extracting the CP half of enable_parallelism into a small _apply_context_parallel and calling it from that path is a ~20-line change that stacks on top of this one; I can open it against #14609 or hand it to @JingyaHuang.