Conversation
`ParallelConfig.__post_init__` rejects a config carrying 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.
* `__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` 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, with the context-parallel dimensions varying fastest. That
order is a default rather than a universal optimum, and the comment says
so: pass `mesh=` on either config to choose the layout yourself.
* 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; the generic backend beside it already reads the coordinate.
Measured on a trn2.48xlarge (MiniMax-H3, 1344x768x124f, 30 steps): TP=8 alone
is 9.285 s/step on 8 cores, TP=4 x ulysses=4 is 5.66 s/step on 16, and
TP=8 x ulysses=4 is 3.941 s/step on 32 -- 2.36x faster than the widest
configuration reachable before this change, with output that is visually
equivalent to the tensor-parallel-only run.
Tests: a `HybridParallelTesterMixin` next to the existing TP and CP mixins,
wired into the Flux transformer tests, plus a Neuron `torchrun` worker
following the `_neuron_tp_worker.py` convention already in the tree. The Neuron
test 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).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
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. |
| self.context_parallel_config.setup(rank, world_size, device, mesh) | ||
| if self.tensor_parallel_config is not None: | ||
| self.tensor_parallel_config.setup(rank, world_size, device, mesh) | ||
| # `TensorParallelConfig.setup` reads `tp_degree` off `mesh.size()`, so when the mesh is shared with |
There was a problem hiding this comment.
Could you put these lines under TensorParallelConfig.setup?
| or torch.distributed.device_mesh.init_device_mesh( | ||
| device_type=device_type, | ||
| mesh_shape=(tp_config.tp_degree, *cp_config.mesh_shape), | ||
| mesh_dim_names=("tp", *cp_config.mesh_dim_names), |
There was a problem hiding this comment.
@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?
|
Tested on 4xA10G, TP × CP combined (ulysses) PASSED. |
|
Thanks for your PR! Can we see some outputs with varying TP + CP configurations and numbers on memory consumption and latency? |
|
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. 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 |
What this does
ParallelConfig.__post_init__currently rejects a config that carries both acontext_parallel_configand atensor_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, sotp_degreecannot 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.setuphands each config its own dimensions: CP keeps the whole mesh, since its ownsetupalready selects"ring"and"ulysses"out of it, while TP gets the"tp"dimension alone because it reads its degree frommesh.size().enable_parallelismbuilds one("tp", "ring", "ulysses")mesh when both are requested.Usage is what you'd expect:
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):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
HybridParallelTesterMixinintests/models/testing_utils/parallelism.py, next to the existing TP and CP mixins, wired into the Flux transformer tests. It needstp_degree * ulysses_degreeaccelerators (4 at the degrees used) and skips below that, so a 2-device runner is unaffected.torchrunworker following the_neuron_tp_worker.pyconvention already in the tree, since Neuron needs the"neuron"distributed backend and cannot use the NCCL spawn path. It runs attp_degree=2 x ulysses_degree=4and passes on a trn2.48xlarge:max_abs_diff=4.2e-05against the single-device reference. The pre-existing_neuron_tp_workertest 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 callenable_parallelismafterwards, which means it applies TP but not the CP hooks. Extracting the CP half ofenable_parallelisminto a small_apply_context_paralleland 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.