[spmd_types] fix spmd->DTensor translation on partial mesh - #3913
Conversation
[ghstack-poisoned]
meta-cla
Bot
added
the
CLA Signed
pianpwk
changed the title
ghstack-source-id: 25f87f1 Pull-Request: pytorch/torchtitan#3913
There was a problem hiding this comment.
translating {DP: V, TP: *} into DTensor (default backend) was failing, even though DP wasn't part of the mesh.
Sounds related to https://github.com/pytorch/torchtitan/pull/3895/changes#r3574804150, so putting a hold. The previous invariance seems
- there's not local activations passing between global module boundary
This PR breaks it in two ways
- there could be local activations passing between global module boundary
- there would be config-based redistribute on local activations across global module boundary
Both sounds unsafe, especially (2) which sounds very hacky. Please think about safer solutions. Worst case I'm OK with (1), but redistribute on the local activations should be banned at boundary.
There was a problem hiding this comment.
there would be config-based redistribute on local activations across global module boundary
tbh, I don't think spmd_types has restrictions on redistribution relations b/w inner/outer axes, e.g. you can do {DP: V, TP: I} -> {DP: I, TP: I}, spmd_types treats this as an allgather on dim 0. Since V in titan might be any kind of varying, I've banned V redistributions in this PR for now?
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
ghstack-source-id: a65138d Pull Request resolved: pytorch/torchtitan#3913
[ghstack-poisoned]
ghstack-source-id: 3be2d89 Pull Request resolved: pytorch/torchtitan#3913
ghstack-source-id: 3be2d89 Pull Request resolved: pytorch/torchtitan#3913
ghstack-source-id: 3be2d89 Pull Request resolved: pytorch/torchtitan#3913
| if changed_axes and ( | ||
| src_types[changed_axes[0]] is spmd.V | ||
| or dst_types[changed_axes[0]] is spmd.V | ||
| ): | ||
| axis = changed_axes[0] | ||
| raise ValueError( | ||
| f"{name}: SpmdLayout-based redistribution changes mesh axis " | ||
| f"{axis.value!r} with spmd.V as the source or destination type. " | ||
| "Config-based redistribution requires non-V types; write an " | ||
| "explicit collective when the value semantics are unclear." | ||
| ) |
There was a problem hiding this comment.
question 1: why only guarding on index [0]?
question 2: in addition to spmd.V, how about spmd.P?
There was a problem hiding this comment.
question 1: why only guarding on index [0]?
This is the "src/dst redistribution can only redistribute one axis" thing, so it's just indexing the one redistribution.
question 2: in addition to spmd.V, how about spmd.P?
P src should be fine, I assume you mean P dst? fwiw we do support convert(R/I/S -> P), with the zero-ing out, or zero-padding: https://github.com/meta-pytorch/spmd_types/blob/60705234a3ebb7a2ddd423cb3259f80985d270b2/spmd_types/_local.py#L912-L999
we actually do support V->I/R (assumed as dim 0 allgather), but I'm banning as V could be interpreted as generally different values/shapes
tianyu-l
left a comment
tianyu-l
left a comment
There was a problem hiding this comment.
After offline discussion, this makes sense, but also a bit hard to maintain. But we'll probably move away from this code soon, so having this sounds good for now.
[ghstack-poisoned]
after #3913 broke this Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * __->__ #4128
2 participants
translating
{DP: V, TP: *}into DTensor (default backend) was failing, even though DP wasn't part of the mesh. Fix by only translating on present axesStack from ghstack (oldest at bottom):