[spmd_types] manual loss parallel CE by pianpwk · Pull Request #3467 · pytorch/torchtitan

pianpwk · 2026-05-31T23:23:01Z

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 (oldest at bottom):

[ghstack-poisoned]

tianyu-l · 2026-06-02T00:54:49Z

+    """
+
+    @staticmethod
+    def typecheck_forward(


Is typecheck_forward an spmd types native solution for autograd function?

This sounds a partial solution.

In torchtitan we shouldn't do partial mesh check, instead we should do full mesh check.

There are more than one way to handle things, including local_spmd field in Module. Is there way to consolidate?

Is it worth it to introduce a Node(Configurable) which

Module inherits

can still have in / out placements

Is typecheck_forward an spmd types native solution for autograd function?

yes, see spmd_types docs, the typechecker will look for this method

# Set of autograd.Function subclasses registered with a typecheck_forward # staticmethod. When .apply() is intercepted by __torch_function__, # typecheck_forward is called INSTEAD of .apply(). It should assert_type() # on inputs, call .apply() to execute, and assert_type() on the output.

In torchtitan we shouldn't do partial mesh check, instead we should do full mesh check.

hmm, I feel this should be handled by earlier module boundaries where we declare activation placements, or the later lm_head sharding? (I should ensure that's explicit and asserted). This fn in theory shouldn't touch the DP/CP axes, so re-checking if DP/CP is on seems a bit messy.

There are more than one way to handle things, including local_spmd field in Module. Is there way to consolidate?

They should serve different purposes here - local_spmd regions are still transparent to typechecking, we just switch to local typecheck. With custom autograd typecheck the type propagation is opaque, so this is a black-box where typecheck_forward is trusted by the typechecker (no typecheck).

Node(Configurable)

this seems... useful for chunked loss, and maybe other contexts outside spmd types too. But I'm not sure about the scope of this - in MoE & VL there are some forward sections & functions that require local/no-typecheck - would the plan be to refactor all of them to Node

OK that sounds convincing!

in MoE & VL there are some forward sections & functions that require local/no-typecheck - would the plan be to refactor all of them to Node

that seems a different topic, I think we should kill the non-DTensor code as much as possible, outside the modules that we apply local_spmd on.
@shuhuayu can also help on this.

tianyu-l · 2026-06-02T02:12:07Z

+        )
+        output_type = dict(spmd.get_local_type(logits))
+        output_type[tp_group] = spmd.I
+        spmd.assert_type(result, output_type)


My worry is that such autograd functions would not enforce typing on all mesh axes

tianyu-l

The only concern is that this code not being actually integrated, but I'll stamp.

tianyu-l · 2026-06-02T22:43:56Z

    )


+def _vocab_shard_size_and_offset(


the file is already big, maybe make this inline if it's only used for loss parallel

tianyu-l · 2026-06-02T22:45:07Z

+    """
+
+    @staticmethod
+    def typecheck_forward(


OK that sounds convincing!

in MoE & VL there are some forward sections & functions that require local/no-typecheck - would the plan be to refactor all of them to Node

that seems a different topic, I think we should kill the non-DTensor code as much as possible, outside the modules that we apply local_spmd on.
@shuhuayu can also help on this.

fegin

LGTM, some change requests but not critical.

fegin · 2026-06-05T05:35:06Z

 safetensors
 einops
 pillow
+spmd_types==0.2.0


The README.md still uses pip install -r requirements.txt. So we should also add this to the general requirements.txt

I think it's just a symlink to this one?

fegin · 2026-06-05T05:44:44Z

+        ctx,
+        logits: torch.Tensor,
+        labels: torch.Tensor,
+        tp_group: dist.ProcessGroup,


Can we take a DeviceMesh? Most places in trainer still take DeviceMesh. So it is better to be consistent. Or is there any blocker?

I think once spmd infra lands we can just call current_mesh() and stop threading this. I wasn't planning to move full DTensor to this as part of the stack yet

fegin · 2026-06-05T05:51:47Z

+    ``IGNORE_INDEX`` labels.  Forward uses three TP all-reduces (max, sumexp,
+    gather); backward is fused (NLL + log-softmax) with zero collectives.


nit: use three TP all-reduce to reduce the results of three oerators.

fegin · 2026-06-05T05:54:00Z

+            safe_labels >= vocab_start + local_vocab_size
+        )
+        local_labels = safe_labels - vocab_start
+        local_labels = local_labels.clone()


Is this clone redundant? safe_labels - vocab_start creates a new tensor. So clone() seems to be redundant.

fegin · 2026-06-05T05:54:05Z

+        local_labels[out_of_range] = 0
+
+        local_result = torch.gather(log_probs, -1, local_labels.unsqueeze(-1))
+        local_result = local_result.clone()


fegin · 2026-06-05T05:54:24Z

+            safe_labels >= ctx.vocab_start + ctx.local_vocab_size
+        )
+        local_labels = safe_labels - ctx.vocab_start
+        local_labels = local_labels.clone()


same redundant clone()?

[ghstack-poisoned]

…tor still TODO Documents the 34-commit pull from upstream/main (27aa490..641b5f6, merge commit fb1c5a3, sitting in worktree ezpz-46th-47th-sync). Two structural refactors hit ezpz: - PR pytorch#3458 (RoPE refactor): REPLAYED in commit 02dd1e7, verified by importing + building all agpt + moe flavors under the torch 2.13 venv. Numerics smoke against baselines still pending. - PR pytorch#3269 (mixed-optimizer refactor): NOT YET STARTED. Affects OptimizersContainer.Config + 8 custom container subclasses in optimizer/containers.py + ~13 callsites in agpt/moe/competition config registries. Also documents that spmd_types==0.2.1 was installed into .venv (new upstream runtime dep from PR pytorch#3466/pytorch#3467/pytorch#3560) and notes the open follow-ups (replay PR pytorch#3269, smoke against baselines, audit RL). Doesn't update upstream-sync.md's "next-sync" reminder banner — this sync isn't complete yet (ezpz branch is not on the worktree HEAD).

pianpwk added 2 commits May 31, 2026 16:22

Update (base update)

5384c3e

[ghstack-poisoned]

Update

1f90807

[ghstack-poisoned]

pytorch-bot Bot added the ciflow/8gpu label May 31, 2026

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

pianpwk added 2 commits June 1, 2026 09:56

Update (base update)

f38ce7d

[ghstack-poisoned]

Update

51634bb

[ghstack-poisoned]

pianpwk changed the title ~~Add standalone loss parallel cross entropy~~ [spmd_types] manual loss parallel CE Jun 1, 2026

pianpwk marked this pull request as ready for review June 1, 2026 17:01

pianpwk requested review from fegin, tianyu-l, wconstab and wwwjn as code owners June 1, 2026 17:01

pianpwk requested a review from ezyang June 1, 2026 17:01

This was referenced Jun 1, 2026

[spmd_types] llama3 enablement #3278

Closed

[spmd_types] spmd infra #3253

Merged

tianyu-l reviewed Jun 2, 2026

View reviewed changes

tianyu-l approved these changes Jun 2, 2026

View reviewed changes

This was referenced Jun 3, 2026

[spmd_types] SpmdLayout for NamedPlacement #3501

Merged

Add spmd_backend config flag #3540

Closed

fegin approved these changes Jun 5, 2026

View reviewed changes

Update

0769acf

[ghstack-poisoned]

This was referenced Jun 5, 2026

Add SPMD sharding configs for MoE models #3550

Closed

Enable Llama4 SPMD type integration #3551

Closed

pianpwk changed the base branch from gh/pianpwk/18/base to main June 5, 2026 16:55

pianpwk merged commit fec0c17 into main Jun 5, 2026
9 of 13 checks passed

pytorch deleted a comment from pytorch-bot Bot Jun 5, 2026

		``IGNORE_INDEX`` labels. Forward uses three TP all-reduces (max, sumexp,
		gather); backward is fused (NLL + log-softmax) with zero collectives.

Conversation

pianpwk commented May 31, 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!

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!

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!

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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

pianpwk commented May 31, 2026 •

edited

Loading