[spmd_types] SpmdLayout for NamedPlacement by pianpwk · Pull Request #3501 · pytorch/torchtitan

pianpwk · 2026-06-03T20:28:15Z

Adding this to allow both spmd types specs, and DTensor-style placements. We're adding something like D107576760 in spmd_types, but in the meantime, because non-llama3 models are still in DTensor placement form, I think we have 2 options:

keep this as a hybrid spmd/DTensor container (e.g. moe_sharding.py calls dense_activation_placement(tp=Replicate()), and migrate to SpmdLayout once all models are off full_dtensor.
since we're adding a spmd->DTensor translation shim anyways, migrate all models from Replicate/Shard/Partial to spmd.R/S/P immediately, use SpmdLayout, and go deal with I/R distinctions later. This should be ok because non-llama3 models won't support spmd backend yet.

Stack from ghstack (oldest at bottom):

[ghstack-poisoned]

Adds the spmd_types integration to titan part-by-part, starting from replacing the `torch.distributed.tensor.parallel.loss_parallel()` DTensor-based context manager with a manually written `_LossParallelCrossEntropy` custom autograd function. Written to support uneven sharding (on last rank), IGNORE_INDEX, verified bitwise against loss_parallel() in test cases. The PR adds spmd_types requirement to titan, as it includes typecheck_forward & testing. Not yet wired into the training path - the later stack will wire it for the spmd backend path, but likely leaves DTensor & plain-tensor on the loss_parallel() path for now. Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0) (oldest at bottom): * #3551 * #3550 * #3278 * #3472 * #3468 * #3471 * #3253 * #3501 * #3466 * __->__ #3467

[ghstack-poisoned]

tianyu-l · 2026-06-07T00:28:17Z

Suggested change

spmd_layout = sharding_config.state_shardings.get(name)

tianyu-l · 2026-06-07T00:34:46Z

+    axis_types: dict[MeshAxisName, spmd.PerMeshAxisSpmdType]
+    partition_spec: spmd.PartitionSpec | tuple[Any, ...] | None = None
+
+    def __post_init__(self) -> None:


This post_init type check feels unnecessary.

updated a bit, this V + PartitionSpec check is helpful for standardized representation

tianyu-l · 2026-06-07T00:35:43Z

+        """Return per-axis types with PartitionSpec sharding represented as S(i).
+
+        This manually handles ``MeshAxisName`` because
+        ``partition_spec_to_shard_types`` expects resolved ``MeshAxis`` values


what is partition_spec_to_shard_types

tianyu-l · 2026-06-07T00:36:19Z

+        ``partition_spec_to_shard_types`` expects resolved ``MeshAxis`` values
+        and treats ``MeshAxisName`` as unresolved string axes.
+        """
+        result = dict(self.axis_types)


what's this line doing

this inherits the non-shard types, e.g. {DP: R, CP: V} + PartitionSpec(None, CP) -> {DP: R, CP: S(1)}

This is still ill-defined. I could hardly imagine anything general could be done with this util.

It doesn't even error out when shard order exists. If you really want to unblock DTensor, at least error out when there is shard order and add TODO to delete this function after migration is done.

I'll edit this later, but the reason I settled on this is:

It so happens that the src/dst sharding configs in torchtitan only involve single axis redistribute calls. The redistribute API requires S(*) to be specified, e.g. S(1) -> R, P -> S(1). It doesn't take V + PartitionSpec form. There isn't a redistribute API that takes this yet, so D107576758 is up for review.

To implement the existing models under SPMD backend, it's clean enough to:

retrieve src & dst shard_types, and loop over axis -> type key values

if there is a mismatch, call single-axis redistribute.

We should raise on > 1 axis mismatch, I can add this in the infra PR. Detection for unsupported src/dst is more involved (e.g. redistributing non-innermost sharding dim in shard order)

As for shard order, the only place it shows up is dense SP placement, and the only redistributions involving it are RS from TP matmul Partial, and sequence AG before lm_head, attention. So it seems enough to retrieve from the SpmdLayout, say src: {DP: S(0), CP: S(1), TP: P} -> dst: {DP: S(0), CP: S(1), TP: S(1)}.

The longer term case would be to add D107576758 in spmd types and call with general SpmdLayout, or if full_dtensor backend is removed, ShardingConfig src/dst can be replaced with the actual collective that's performed, as the placement-based redistribute might be overkill for what we actually need. With SPMD types there might also be a preference for explicit collectives, rather than DTensor-style redistribute.

sounds very reasonable

it's good to know the {DP: S(0), CP: S(1), TP: P} -> dst: {DP: S(0), CP: S(1), TP: S(1)} case is handled

we don't have to support the general cases, please just error out.

moved out of types.py, error checking belongs in the spmd infra PR

tianyu-l · 2026-06-07T00:37:18Z

+        return tuple(self.axis_types)
+
+    def shard_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]:
+        """Return per-axis types with PartitionSpec sharding represented as S(i).


Is this transitional code? I feel in the long term this could be ill-defined, e.g. when you express shard order / same tensor dim sharded twice.

In general I think it's fine to have transitional code, but let's have precise comments on the assumptions.

Updated the docstring, it's not clear if it's transitional yet, but useful to work with per-axis redistribute API

tianyu-l · 2026-06-07T00:40:06Z

I think it's fine to use MeshAxisName since @fegin prefers.
If so, my preference is minimize it to

Suggested change

class MeshAxis(StrEnum):

tianyu-l · 2026-06-07T00:40:31Z

-# will be applied against; ``resolve_placements`` errors otherwise.
-NamedPlacement = dict[MeshAxisName, Placement]
+@dataclass(frozen=True, slots=True)
+class NamedPlacement:


sounds good to me

[ghstack-poisoned]

tianyu-l · 2026-06-08T00:15:32Z

+    """Convert an SPMD layout to DTensor placements keyed by mesh axis name."""
+    result: dict[MeshAxisName, Placement] = {}
+    for axis_name, axis_type in layout.shard_types().items():
+        if axis_type == spmd.R or axis_type == spmd.I or axis_type == spmd.V:


spmd.V -> Replicate?

tianyu-l · 2026-06-08T00:15:53Z

+            raise ValueError(
+                f"Unsupported SPMD type for axis {axis_name.value!r}: {axis_type!r}."
+            )


is this else possible?

tianyu-l · 2026-06-08T00:20:00Z

+        ``partition_spec_to_shard_types`` expects resolved ``MeshAxis`` values
+        and treats ``MeshAxisName`` as unresolved string axes.
+        """
+        result = dict(self.axis_types)


This is still ill-defined. I could hardly imagine anything general could be done with this util.

It doesn't even error out when shard order exists. If you really want to unblock DTensor, at least error out when there is shard order and add TODO to delete this function after migration is done.

[ghstack-poisoned]

fegin

Please check torchtitan/components/lora.py. It doesn't seems to be migrated.

fegin · 2026-06-08T20:13:37Z

+    TODO(pianpwk): Replace this with ``spmd_types.SpmdLayout`` once that API is
+    available in TorchTitan's minimum ``spmd_types`` version.


Do you mean that you are going to upstream this to spmd_types? If so, does this mean that MeshAxisName will also in spmd_types?

MeshAxisName unlikely, I don't think spmd_types will hold a enum of valid mesh names? But it does share a str key form, for unnormalized mesh axis type.

But yes, see D107576760

fegin · 2026-06-08T20:21:57Z

-    TODO: Ideally placements would be defined on a computation mesh that
-    has a single DP axis (no DP_REPLICATE vs DP_SHARD distinction). That
-    requires a mesh switch between FSDP storage and computation — likely
-    resolved by FlexShard. Revisit once FlexShard lands.


Should we keep this TODO or any of your diffs fix this?

ah, moved around a bit but this PR should do it now

tianyu-l

some last comments.

tianyu-l · 2026-06-08T22:12:57Z


-__all__ = ["ParallelDims"]
+
+class StrEnum(str, Enum):


Maybe unify
https://github.com/search?q=repo%3Apytorch%2Ftorchtitan%20StrEnum&type=code

Maybe StrEnum would justify a protocols/types.py file, lol

tried to unify but a lot of circular imports...

e.g. parallel_dims.py
-> torchtitan.protocols.types
-> torchtitan/protocols/init.py
-> torchtitan/protocols/model.py
-> torchtitan/protocols/module.py
-> torchtitan/distributed/parallel_dims.py

We should just move to Python 3.12 to avoid defining this. lol

tianyu-l · 2026-06-08T22:20:48Z

+        This is not meant as a minimal description of the SPMD layout; shard order
+        cannot be expressed carefully. This is a helper for calling spmd.redistribute,
+        which takes per-axis types (e.g. redistribute(S(1) -> R)).


Do you mean you intentionally want to allow for the ambiguous case, where CP: S(1), TP: S(1) is a legit output? Please explicit mention that "shard order info will be lost".

tianyu-l · 2026-06-08T22:21:20Z

+    def axes(self) -> tuple[MeshAxisName, ...]:
+        return tuple(self.axis_types)
+
+    def shard_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]:


Suggested change

def shard_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]:

def per_axis_spmd_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]:

pianpwk · 2026-06-08T22:25:19Z

some last comments.

you say this, but wait to see what my next update introduces

[ghstack-poisoned]

pianpwk · 2026-06-08T23:17:00Z

Ok, pushed a change to restructure, review-worthy but not major, and I think inline with what we've discussed. Mostly triggered by SpmdLayout.post_init establishing some invariants to make the infra commit work (no double sharding in local_types, forcing V + PartitionSpec).

To avoid breaking CI, I had to include unfold_dp_axis (DP_REPLICATE: S(0), DP_SHARD: S(0) was erroring in validation), so this PR starts using unified DP axis, and unfolds in translation / DTensor mesh construction.

Also had to define dense_sequence_parallel_placement as using V + PartitionSpec for seq-dim sharding, and call it in place of dense_activation_placement(tp=S(1)). This means the later decoder sharding PR should mostly be I/R reasoning.

[ghstack-poisoned]

pianpwk · 2026-06-08T23:26:59Z

last update was just lint

ezyang · 2026-06-09T03:05:45Z

Since Tianyu has approved this, nothing I say here is binding. But in the spirit of making things easier to review, I would have tried much harder NOT to have to do all of the

-        state_shardings={"weight": dense_param_placement(tp=Replicate())},
+        state_shardings={"weight": dense_param_placement(tp=spmd.R)},

modifications, and do those all in a purely mechanical PR later, if that's actually the way to do it. Like the types are basically isomorphic, it really doesn't cost you anything to accept either Replicate() or spmd.R, and it sure saves you a heck of a lot of fuzz in the diff.

[ghstack-poisoned]

…harding Upstream PR pytorch#3501 ([spmd_types] SpmdLayout for NamedPlacement) renamed set_dense_ffn_sharding's first argument: def set_dense_ffn_sharding(feed_forward_cfg, *, attn_x_placement: Placement, enable_sp: bool) → def set_dense_ffn_sharding(feed_forward_cfg, *, attn_x_layout: SpmdLayout, enable_sp: bool) ezpz/agpt/sharding.py + ezpz/moe/sharding.py both called the old signature directly. After the 2026-06-09 upstream merge (commit 92f764a) the first attempted training run died at trainer init: TypeError: set_dense_ffn_sharding() got an unexpected keyword argument 'attn_x_placement' at torchtitan/experiments/ezpz/agpt/sharding.py:84 Caught by bitwise_sync_check.sh on its first invocation (job 12468296) — exactly what the bitwise smoke is for. Fix matches the new llama3 / deepseek_v3 pattern: build the layout via `dense_sequence_parallel_placement()` / `dense_activation_placement (tp=spmd.R)` helpers instead of constructing a raw Placement. Also pulled the same helpers into moe/sharding.py's MLA attention sharding_config (`in_src_shardings={"x": attn_x_layout, ...}`) since the construction is shared. NOT a numerics-equivalent change yet — the prior code passed `Shard(1) if enable_sp else Replicate()` directly; the new helpers wrap those in SpmdLayout's typed envelope. Bitwise verdict pending from 12468298 — if DRIFT, the upstream helpers compute a different layout than the raw Placement and we need to investigate. Earlier upstream-sync.md claimed PR pytorch#3501 had no ezpz impact — that was wrong, will fix in the followup doc commit.

The root CLAUDE.md project rule says non-computation changes (e.g. refactoring, import renames, upstream merges) must produce identical loss before vs. after with --debug.seed=42 --debug.deterministic. We had the rule + the pattern (scripts/loss_compare.py) but no turn-key script for "run on HEAD then run on pre-merge commit then diff" — every upstream sync was on the honor system. This script automates the pre/post comparison: qsub -A datascience -q workq -l select=2 -l walltime=01:00:00 \ -l filesystems=tegu:home \ -v PRE_MERGE_COMMIT=<sha>,STEPS=20 \ torchtitan/experiments/ezpz/scripts/bitwise_sync_check.sh Phase 1: git stash any working-tree changes, run an agpt_2b N-step smoke with --debug.seed=42 --debug.deterministic, extract per-step (loss, grad_norm) tuples. Phase 2: git checkout the pre-merge commit, same smoke, same extraction. Then diff the two metrics streams; verdict is IDENTICAL or DRIFT. Output: logs/bitwise-sync-check-${PBS_JOBID}/ head.log / pre.log full run output head.metrics / pre.metrics extracted (step, loss, grad_norm) diff.txt unified diff of pre vs head metrics verdict IDENTICAL or DRIFT run.log phase boundaries + final summary On exit, trap restores the original branch + pops the stash. First invocation (job 12468296) caught a real upstream-sync regression in 2:34 wall time: ezpz/agpt/sharding.py + ezpz/moe/sharding.py called set_dense_ffn_sharding's old kwarg, post-merge upstream renamed it (PR pytorch#3501). Fixed in commit 786e568.

Adding this to allow both spmd types specs, and DTensor-style placements. We're adding something like D107576760 in spmd_types, but in the meantime, because non-llama3 models are still in DTensor placement form, I think we have 2 options: 1) keep this as a hybrid spmd/DTensor container (e.g. moe_sharding.py calls `dense_activation_placement(tp=Replicate())`, and migrate to SpmdLayout once all models are off full_dtensor. 2) since we're adding a spmd->DTensor translation shim anyways, migrate all models from `Replicate/Shard/Partial` to `spmd.R/S/P` immediately, use SpmdLayout, and go deal with I/R distinctions later. This should be ok because non-llama3 models won't support spmd backend yet. Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0) (oldest at bottom): * #3278 * #3472 * #3468 * #3471 * #3253 * #3587 * #3586 * __->__ #3501

pianpwk added 2 commits June 3, 2026 13:28

Update

7bc7ae3

[ghstack-poisoned]

Update (base update)

c8e1e0d

[ghstack-poisoned]

pytorch-bot Bot added the ciflow/8gpu label Jun 3, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jun 3, 2026

pianpwk changed the title ~~Introduce NamedPlacement container~~ NamedPlacement as container Jun 3, 2026

pianpwk marked this pull request as ready for review June 3, 2026 22:11

pianpwk requested review from fegin, tianyu-l, wconstab and wwwjn as code owners June 3, 2026 22:11

pianpwk changed the title ~~NamedPlacement as container~~ [spmd_types] NamedPlacement as container Jun 3, 2026

pianpwk added 4 commits June 3, 2026 15:43

Update

ecec423

[ghstack-poisoned]

Update (base update)

3552429

[ghstack-poisoned]

Update

69257d0

[ghstack-poisoned]

Update (base update)

8a3e07e

[ghstack-poisoned]

pianpwk mentioned this pull request Jun 4, 2026

Add spmd_backend config flag #3540

Closed

pianpwk added 2 commits June 5, 2026 09:54

Update

3f3729b

[ghstack-poisoned]

Update (base update)

5f385aa

[ghstack-poisoned]

This was referenced Jun 5, 2026

Add SPMD sharding configs for MoE models #3550

Draft

Enable Llama4 SPMD type integration #3551

Draft

pianpwk added 2 commits June 5, 2026 14:42

Update

8a4597f

[ghstack-poisoned]

Update (base update)

5ebdaf9

[ghstack-poisoned]

pianpwk requested a review from tianyu-l June 6, 2026 22:05

tianyu-l reviewed Jun 7, 2026

View reviewed changes

pianpwk added 2 commits June 7, 2026 13:12

Update

f50c09c

[ghstack-poisoned]

Update (base update)

720c72e

[ghstack-poisoned]

pianpwk requested a review from tianyu-l June 7, 2026 20:18

Update

163d951

[ghstack-poisoned]

tianyu-l reviewed Jun 8, 2026

View reviewed changes

Update

3cff270

[ghstack-poisoned]

fegin reviewed Jun 8, 2026

View reviewed changes

tianyu-l reviewed Jun 8, 2026

View reviewed changes

Update

73fb36a

[ghstack-poisoned]

Update

0930cc9

[ghstack-poisoned]

tianyu-l approved these changes Jun 8, 2026

View reviewed changes

pianpwk changed the base branch from gh/pianpwk/25/base to main June 8, 2026 23:40

pianpwk changed the base branch from main to gh/pianpwk/25/base June 9, 2026 05:06

pianpwk added 2 commits June 8, 2026 22:06

Update (base update)

c55d848

[ghstack-poisoned]

Update

5799d26

[ghstack-poisoned]

pianpwk changed the base branch from gh/pianpwk/25/base to main June 9, 2026 05:28

pianpwk changed the base branch from main to gh/pianpwk/25/base June 9, 2026 06:17

This was referenced Jun 9, 2026

[spmd_types] global_valid_tokens: float | None #3586

Merged

[spmd_types] state sharding and redistribution infra #3587

Merged

pianpwk changed the base branch from gh/pianpwk/25/base to main June 9, 2026 07:53

pianpwk merged commit 465cd67 into main Jun 9, 2026
9 of 11 checks passed

		TODO(pianpwk): Replace this with ``spmd_types.SpmdLayout`` once that API is
		available in TorchTitan's minimum ``spmd_types`` version.

	def shard_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]:
	def per_axis_spmd_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]:

Uh oh!

Conversation

pianpwk commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pianpwk Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fegin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pianpwk commented Jun 8, 2026

Uh oh!

pianpwk commented Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pianpwk commented Jun 8, 2026

Uh oh!

ezyang commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

pianpwk commented Jun 3, 2026 •

edited

Loading

pianpwk Jun 8, 2026 •

edited

Loading

pianpwk commented Jun 8, 2026 •

edited

Loading

ezyang commented Jun 9, 2026 •

edited

Loading