[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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[JingyaHuang]

@whn09 Thanks for contributing the PR. Left some nits, and could you add a section combining CP+TP in the doc as well?

[JingyaHuang]

Could you put these lines under TensorParallelConfig.setup?

[JingyaHuang]

@sayakpaul what's your take on the default order of the mesh? If Neuron runtime requests Ulysses mesh to be contiguous maybe we should set default value according to the device?

[JingyaHuang]

Tested on 4xA10G, TP × CP combined (ulysses) PASSED.

[sayakpaul]

Thanks for your PR! Can we see some outputs with varying TP + CP configurations and numbers on memory consumption and latency?

[github-actions]

Hi @whn09, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[JingyaHuang]

I'll share numbers for Neuron, TPU and CUDA once #14544 is merged, since peak memory is more meaningful to measure after that.

[whn09]

@sayakpaul here are numbers for MiniMax-H3 (33.1B DiT) on one trn2.48xlarge (64 logical NeuronCores, 24 GiB HBM each), now that #14544 is in.

1344×768, 124 frames, 30 steps, bf16, from_pretrained(..., parallel_config=...) (sharded load). TP is capped at 8 because H3 has 56 heads, so CP (Ulysses) is the only way past 8 cores.

cores	config	weights / core	s / step	speed-up	peak HBM / core (DiT)
8	TP=8	7.84 GiB	9.33 (mean; 8.49 → 10.8)*	1×	18.7 GiB
16	TP=4 × Ulysses=4	15.53 GiB	5.31	1.76×	21.7 GiB
32	TP=8 × Ulysses=4	7.84 GiB	3.81	2.45×	13.7 GiB
32	TP=4 × Ulysses=8	15.53 GiB	3.97	2.35×	20.5 GiB
64	TP=8 × Ulysses=8	7.84 GiB	3.42	2.73×	12.9 GiB

• Peak HBM is sampled from the Neuron driver's per-core counters over the denoising stage (so it includes the runtime's scratchpad and DMA rings, not only torch's allocator).
• TP is what controls memory (TP=4 → TP=8 halves the weights and saves ~7–8 GiB per core); CP mostly buys latency, plus a little activation memory (13.7 → 12.9 GiB going from Ulysses=4 to 8).
• At the same 32 cores, TP=8 × Ulysses=4 beats TP=4 × Ulysses=8: 4% faster and 6.9 GiB less per core. So on this model, max out TP first and add CP on top. (TP=4 × Ulysses=16 at 64 cores is not reachable: 14 heads per TP rank padded to 16 leaves fewer than one real head per rank.)
• *At TP=8 alone, the per-step time drifts up deterministically over the run (reproduced twice). The CP configs stay flat. I'm still looking into it.

Outputs (frames 0 / 60 / 120; rows = 8 / 16 / 32 / 64 cores): same seed and prompt, same content across configs.

This was run on #14609's branch (7e4a7eff, which carries this PR's commit), with three changes on top:

1. ulysses-fastest mesh order (see my comment on [feat] Tensor parallelism for MiniMax-H3 #14609),
2. applying the CP hooks on the sharded-load path; without this, from_pretrained(parallel_config=TP+CP) silently runs TP only,
3. relaxing assert H_PAD < NEW_H_LOCAL to <= in attention_dispatch.py, which is needed for 64 cores (7 heads per TP rank padded to 8 for Ulysses=8).

Env: torch 2.14.0, torch-neuronx 2.14.3.0.2008+2dda91f8d.

[whn09]

Merged main into this branch to resolve the conflicts (74f0f50). They were all in the tests: main moved the Neuron TP tests to the mp.spawn-based TensorParallelNeuronTesterMixin and dropped _neuron_tp_worker.py, so this PR now uses that mixin for plain TP and drops its own torchrun TP test and make_neuron_tp_spec. ContextAndTensorParallelTesterMixin, HybridParallelTesterMixin and the hybrid / sharded-load Neuron workers are unchanged. On a trn2.48xlarge (Neuron SDK 2.33: torch 2.14.0, torch-neuronx 2.14.3.0.1, neuronx-cc 2.28) these pass, each in its own pytest process:

• TestFluxTransformerTensorParallelNeuron
• TestQwenImageTransformerTensorParallelNeuron
• TestFluxTransformerHybridParallelNeuron (TP=2 × Ulysses=4)
• TestFluxTransformerShardedLoadNeuron (TP=2 × Ulysses=4)

@sayakpaul updated MiniMax-H3 numbers on the same SDK 2.33, at the full step count: 1344×768, 124 frames, 50 steps, bf16, 64 NeuronCores as TP=8 × Ulysses=8, sharded load via from_pretrained(..., parallel_config=...).

denoise, steady state	3.265 s / step (flat over all 49 transformer calls, 3.15–3.35)
denoise, 49 calls	177.3 s (first call 15.9 s with a warm compile cache)
weights / core	7.84 GiB
peak HBM / core	12.96 GiB reserved by torch after denoising; 22.46 GiB for the whole run, reached in the fp32 video VAE decode
video VAE decode (tiles spread over the 64 cores)	12.2 s
audio VAE decode (CPU)	2.0 s

Against the earlier run (30 steps, torch-neuronx 2.14.3 nightly), the steady-state step is 4.5% faster: 3.42 → 3.265 s. Output looks right, same content as before (fox in a snowy forest).

Tree: #14609 @ bc35ec2d merged with this branch. The only change on top is still the H_PAD < NEW_H_LOCAL → <= relaxation in attention_dispatch.py that 64 cores need (7 heads per TP rank padded to 8 for Ulysses=8). The mesh order and sharded-load CP hooks from my previous comment are both in now.

中文

已把 main 合进本分支解决冲突（74f0f501c），冲突都在测试里：main 把 Neuron TP 测试改成基于 mp.spawn 的 TensorParallelNeuronTesterMixin，删了 _neuron_tp_worker.py。因此纯 TP 改用该 mixin，本 PR 自己的 torchrun TP 测试和 make_neuron_tp_spec 删掉；ContextAndTensorParallelTesterMixin、HybridParallelTesterMixin 以及 hybrid / sharded-load 两个 Neuron worker 不变。在 trn2.48xlarge（Neuron SDK 2.33：torch 2.14.0、torch-neuronx 2.14.3.0.1、neuronx-cc 2.28）上，上面四个测试各自单独起 pytest 进程，全部通过。

MiniMax-H3 在同一 SDK 2.33、完整 50 步下的新数据：1344×768、124 帧、bf16、64 核 TP=8 × Ulysses=8、sharded load。

• 稳态 3.265 s/step（49 次 transformer 调用全程平稳，3.15–3.35）。
• 49 次 denoise 共 177.3 s，首次 15.9 s（编译缓存已热）。
• 每核权重 7.84 GiB；denoise 后 torch reserved 12.96 GiB，整跑峰值 22.46 GiB 出现在 fp32 视频 VAE 解码阶段。
• 视频 VAE 解码（tile 分到 64 核）12.2 s，音频 VAE 解码（CPU）2.0 s。
• 比上次（30 步、torch-neuronx 2.14.3 nightly）稳态快 4.5%（3.42 → 3.265 s），画面正常，内容与之前一致。

代码：#14609 @ bc35ec2d 合并本分支，只多一处改动：attention_dispatch.py 里把 H_PAD < NEW_H_LOCAL 放宽成 <=，64 核需要它（每个 TP rank 7 个头，Ulysses=8 时补到 8）。上次评论里的 mesh 维序与 sharded-load CP hook 两处都已在代码里。

[whn09]

@JingyaHuang added the CP + TP section you asked for to distributed_inference.md (0c7b4aa): Combining context and tensor parallelism, between the tensor parallelism section and "Choosing a strategy". It covers:

• the ParallelConfig usage, with a FLUX.2 TP=2 × Ulysses=2 example through sharded from_pretrained;
• the degree checks;
• why the default mesh is ("tp", "ring", "ulysses") and when to pass your own mesh=;
• MiniMax-H3 numbers on trn2 showing that TP is the memory lever and CP the latency one.

It also replaces the "not yet validated, treat it as experimental" bullet in "Choosing a strategy" with a link to the new section.

The H3 rows depend on H3's _tp_plan from #14609, so they only reproduce on main once that lands. Happy to drop the table, or move it into #14609, if you'd rather keep this PR model-agnostic.

中文

已在 distributed_inference.md 补上 CP + TP 章节（0c7b4aa1）："Combining context and tensor parallelism"，放在 tensor parallelism 与 "Choosing a strategy" 之间。内容包括：

• 用法：一个 ParallelConfig 同时传两种配置，示例是 FLUX.2 TP=2 × Ulysses=2，走 sharded from_pretrained；
• 并行度校验规则；
• 默认 mesh 为什么是 ("tp", "ring", "ulysses")，以及什么时候该自己传 mesh=；
• MiniMax-H3 在 trn2 上的数据：TP 管显存，CP 管延迟。

"Choosing a strategy" 里原来那句"尚未验证、视为实验性"也换成了指向新章节的链接。

H3 那几行依赖 #14609 里 H3 的 _tp_plan，要等它合进 main 才能复现；如果希望本 PR 保持与具体模型无关，我可以删掉这个表，或者挪到 #14609。