[PT-D] Relaxed `contract` to allow `Sequence[nn.Module]` #127773

awgu · 2024-06-03T17:06:55Z

Stack from ghstack (oldest at bottom):

This PR relaxes @contract to allow the 1st argument to be Sequence[nn.Module] instead of strictly nn.Module. This is required for the next PR, which allows fully_shard to take in List[nn.Module].

cc @XilunWu @H-Huang @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @penguinwu @tianyu-l @yf225 @chauhang

[ghstack-poisoned]

pytorch-bot · 2024-06-03T17:06:58Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/127773

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit d1cd203 with merge base 6f275ae ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

ghstack-source-id: 5bd8e3b579f284c3153a4f76bdf1eea698f0422c Pull Request resolved: pytorch#127773

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

This PR relaxes `contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

This PR allows `fully_shard`'s first argument to be `List[nn.Module]` instead of strictly `nn.Module`. This allows more flexible grouping of modules/parameters for communication, which can lead to memory savings and/or more efficient communication. **Approach** At a high level, we can think of a model as a tree of modules. Previously, we could only select specific module nodes in this tree as representing one FSDP parameter group. With this PR, we can select a group of module nodes, effectively becoming a single super node. To implement the runtime schedule, we define new forward hooks that run based on the following semantics: - If a module is the first to run the pre-hook, actually run the given pre-hook. Otherwise, the pre-hook is no-op. - If a module is the last to run the post-hook, actually run the given post-hook. Otherwise, the post-hook is a no-op. - First and last are determined by scoreboarding against a set of the modules. - This set must get cleared at the end of backward in the case that >=1 module in the list is never used, in which case we still want the forward hooks to run in the next forward after this backward. Beyond these new forward hooks, everything else is some simple generalization from `Module` to `List[Module]` or `Tuple[Module, ...]`. **Examples** This PR enables wrapping Llama models more efficiently by grouping the final norm and output linear together: pytorch/torchtitan#382. If at least one of the modules in the list does not run forward before backward, then there will be a warning message like: ``` 1 of the 2 modules passed to fully_shard did not run forward before backward, which is error-prone since FSDP post-forward/pre-backward logic will not run for these modules. We recommend passing only modules that run forward together. Modules that did not run forward: [FSDPLinear(in_features=1, out_features=1, bias=True)] ``` Pull Request resolved: #127786 Approved by: https://github.com/yf225, https://github.com/weifengpy ghstack dependencies: #127773

atalman · 2024-07-16T23:46:24Z

@pytorchmergebot revert -c ghfirst -m "failing internally"

pytorchmergebot · 2024-07-16T23:48:01Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2024-07-16T23:48:12Z

@awgu your PR has been successfully reverted.

…7773)" This reverts commit b276957. Reverted #127773 on behalf of https://github.com/atalman due to failing internally ([comment](#127773 (comment)))

…7773)" This reverts commit b276957. Reverted #127773 on behalf of https://github.com/atalman due to failing internally ([comment](#127773 (comment))) [ghstack-poisoned]

This PR relaxes `@contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. [ghstack-poisoned]

This PR relaxes `contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. ghstack-source-id: 3ef70104cc68f542ce1d91cd5c3b1ec37c2dd4cb Pull Request resolved: #130947

…ce[nn.Module]` (#127773)" This PR relaxes `contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. cc XilunWu H-Huang kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]

…127773)" This PR relaxes `contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. cc XilunWu H-Huang kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]

This PR allows `fully_shard`'s first argument to be `List[nn.Module]` instead of strictly `nn.Module`. This allows more flexible grouping of modules/parameters for communication, which can lead to memory savings and/or more efficient communication. **Approach** At a high level, we can think of a model as a tree of modules. Previously, we could only select specific module nodes in this tree as representing one FSDP parameter group. With this PR, we can select a group of module nodes, effectively becoming a single super node. To implement the runtime schedule, we define new forward hooks that run based on the following semantics: - If a module is the first to run the pre-hook, actually run the given pre-hook. Otherwise, the pre-hook is no-op. - If a module is the last to run the post-hook, actually run the given post-hook. Otherwise, the post-hook is a no-op. - First and last are determined by scoreboarding against a set of the modules. - This set must get cleared at the end of backward in the case that >=1 module in the list is never used, in which case we still want the forward hooks to run in the next forward after this backward. Beyond these new forward hooks, everything else is some simple generalization from `Module` to `List[Module]` or `Tuple[Module, ...]`. **Examples** This PR enables wrapping Llama models more efficiently by grouping the final norm and output linear together: pytorch/torchtitan#382. If at least one of the modules in the list does not run forward before backward, then there will be a warning message like: ``` 1 of the 2 modules passed to fully_shard did not run forward before backward, which is error-prone since FSDP post-forward/pre-backward logic will not run for these modules. We recommend passing only modules that run forward together. Modules that did not run forward: [FSDPLinear(in_features=1, out_features=1, bias=True)] ``` Pull Request resolved: #127786 Approved by: https://github.com/yf225, https://github.com/weifengpy ghstack dependencies: #127773 ghstack-source-id: c70870546bed6108b163062b88415cbba3d37925

awgu · 2024-07-17T17:21:06Z

Closing in favor of #130947

…7773) (#130947) This PR relaxes `@contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. --- **Changes for reland:** - The previous PR assumed that any `func` decorated with `@contract` would return the same input `module` as output (which is true for PT-D composable APIs). - However, TorchRec `shard` returns a different module as output (though that module _does_ satisfy the `@contract` FQN check). - This PR removes the assumption and instead only enforces the FQN check following the input module order. In other words, if calling `func([x1, ..., xN])` for `N` modules `x1, ..., xN` that returns `[y1, ..., yM]` for `M` modules, we require that `N = M` and that FQNs are preserved coordinate-wise: `xi` and `yi` have same FQNs for all `i = 1, ..., N`. Differential Revision: [D59863438](https://our.internmc.facebook.com/intern/diff/D59863438) Pull Request resolved: #130947 Approved by: https://github.com/weifengpy, https://github.com/atalman

…7773)" This reverts commit b276957. Reverted #127773 on behalf of https://github.com/atalman due to failing internally ([comment](#127773 (comment)))

…orch#127773) (pytorch#130947) This PR relaxes `@contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. --- **Changes for reland:** - The previous PR assumed that any `func` decorated with `@contract` would return the same input `module` as output (which is true for PT-D composable APIs). - However, TorchRec `shard` returns a different module as output (though that module _does_ satisfy the `@contract` FQN check). - This PR removes the assumption and instead only enforces the FQN check following the input module order. In other words, if calling `func([x1, ..., xN])` for `N` modules `x1, ..., xN` that returns `[y1, ..., yM]` for `M` modules, we require that `N = M` and that FQNs are preserved coordinate-wise: `xi` and `yi` have same FQNs for all `i = 1, ..., N`. Differential Revision: [D59863438](https://our.internmc.facebook.com/intern/diff/D59863438) Pull Request resolved: pytorch#130947 Approved by: https://github.com/weifengpy, https://github.com/atalman

ghstack-source-id: 9841a65ca843f496f71801a6f49a65061ca82630 Pull Request resolved: pytorch/pytorch#127773

…7773)" This reverts commit b276957. Reverted pytorch/pytorch#127773 on behalf of https://github.com/atalman due to failing internally ([comment](pytorch/pytorch#127773 (comment))) ghstack-source-id: d3b876189599fa0c6c9d43a4d6ebc732ed2eab7a Pull Request resolved: pytorch/pytorch#130946

…7773)" This reverts commit b276957. Reverted #127773 on behalf of https://github.com/atalman due to failing internally ([comment](#127773 (comment))) (cherry picked from commit d7a8e8f)

…7773) (#130947) This PR relaxes `@contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. --- **Changes for reland:** - The previous PR assumed that any `func` decorated with `@contract` would return the same input `module` as output (which is true for PT-D composable APIs). - However, TorchRec `shard` returns a different module as output (though that module _does_ satisfy the `@contract` FQN check). - This PR removes the assumption and instead only enforces the FQN check following the input module order. In other words, if calling `func([x1, ..., xN])` for `N` modules `x1, ..., xN` that returns `[y1, ..., yM]` for `M` modules, we require that `N = M` and that FQNs are preserved coordinate-wise: `xi` and `yi` have same FQNs for all `i = 1, ..., N`. Differential Revision: [D59863438](https://our.internmc.facebook.com/intern/diff/D59863438) Pull Request resolved: #130947 Approved by: https://github.com/weifengpy, https://github.com/atalman (cherry picked from commit ff7e021)

) This PR relaxes `@contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. Pull Request resolved: pytorch#127773 Approved by: https://github.com/weifengpy

This PR allows `fully_shard`'s first argument to be `List[nn.Module]` instead of strictly `nn.Module`. This allows more flexible grouping of modules/parameters for communication, which can lead to memory savings and/or more efficient communication. **Approach** At a high level, we can think of a model as a tree of modules. Previously, we could only select specific module nodes in this tree as representing one FSDP parameter group. With this PR, we can select a group of module nodes, effectively becoming a single super node. To implement the runtime schedule, we define new forward hooks that run based on the following semantics: - If a module is the first to run the pre-hook, actually run the given pre-hook. Otherwise, the pre-hook is no-op. - If a module is the last to run the post-hook, actually run the given post-hook. Otherwise, the post-hook is a no-op. - First and last are determined by scoreboarding against a set of the modules. - This set must get cleared at the end of backward in the case that >=1 module in the list is never used, in which case we still want the forward hooks to run in the next forward after this backward. Beyond these new forward hooks, everything else is some simple generalization from `Module` to `List[Module]` or `Tuple[Module, ...]`. **Examples** This PR enables wrapping Llama models more efficiently by grouping the final norm and output linear together: pytorch/torchtitan#382. If at least one of the modules in the list does not run forward before backward, then there will be a warning message like: ``` 1 of the 2 modules passed to fully_shard did not run forward before backward, which is error-prone since FSDP post-forward/pre-backward logic will not run for these modules. We recommend passing only modules that run forward together. Modules that did not run forward: [FSDPLinear(in_features=1, out_features=1, bias=True)] ``` Pull Request resolved: pytorch#127786 Approved by: https://github.com/yf225, https://github.com/weifengpy ghstack dependencies: pytorch#127773

…orch#127773)" This reverts commit b276957. Reverted pytorch#127773 on behalf of https://github.com/atalman due to failing internally ([comment](pytorch#127773 (comment)))

…orch#127773) (pytorch#130947) This PR relaxes `@contract` to allow the 1st argument to be `Sequence[nn.Module]` instead of strictly `nn.Module`. This is required for the next PR, which allows `fully_shard` to take in `List[nn.Module]`. --- **Changes for reland:** - The previous PR assumed that any `func` decorated with `@contract` would return the same input `module` as output (which is true for PT-D composable APIs). - However, TorchRec `shard` returns a different module as output (though that module _does_ satisfy the `@contract` FQN check). - This PR removes the assumption and instead only enforces the FQN check following the input module order. In other words, if calling `func([x1, ..., xN])` for `N` modules `x1, ..., xN` that returns `[y1, ..., yM]` for `M` modules, we require that `N = M` and that FQNs are preserved coordinate-wise: `xi` and `yi` have same FQNs for all `i = 1, ..., N`. Differential Revision: [D59863438](https://our.internmc.facebook.com/intern/diff/D59863438) Pull Request resolved: pytorch#130947 Approved by: https://github.com/weifengpy, https://github.com/atalman

[PT-D] Relaxed contract to allow Sequence[nn.Module]

6494413

[ghstack-poisoned]

pytorch-bot bot added the oncall: distributed Add this issue/PR to distributed oncall triage queue label Jun 3, 2024

This was referenced Jun 3, 2024

[BE] Enabled lintrunner on torch/distributed/utils.py #127771

Closed

[BE] Fixed CPU autocast warning #127774

Closed

Update on "[PT-D] Relaxed contract to allow Sequence[nn.Module]"

e064321

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

Update on "[PT-D] Relaxed contract to allow Sequence[nn.Module]"

38ec4ed

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

awgu mentioned this pull request Jun 3, 2024

[FSDP2] Fixed variable shadowing of module #127776

Closed

This was referenced Jun 3, 2024

[FSDP2] Allowed List[nn.Module] as arg #127786

Closed

[FSDP2] Showed forward prefetching augmented to fully_shard #127828

Closed

awgu added 2 commits June 3, 2024 17:25

Update on "[PT-D] Relaxed contract to allow Sequence[nn.Module]"

6f3d742

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

Update on "[PT-D] Relaxed contract to allow Sequence[nn.Module]"

654a89d

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

awgu added the topic: not user facing topic category label Jun 4, 2024

Update on "[PT-D] Relaxed contract to allow Sequence[nn.Module]"

4d08f3b

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

awgu added a commit to awgu/pytorch that referenced this pull request Jun 21, 2024

[PT-D] Relaxed contract to allow Sequence[nn.Module]

97b2eca

ghstack-source-id: 5bd8e3b579f284c3153a4f76bdf1eea698f0422c Pull Request resolved: pytorch#127773

awgu added 3 commits June 21, 2024 14:45

Update on "[PT-D] Relaxed contract to allow Sequence[nn.Module]"

fd6f9bb

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

Update on "[PT-D] Relaxed contract to allow Sequence[nn.Module]"

02b69a5

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

Update on "[PT-D] Relaxed contract to allow Sequence[nn.Module]"

de90873

cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

awgu marked this pull request as ready for review July 11, 2024 15:22

awgu requested review from fegin, wconstab, wanchaol and weifengpy July 11, 2024 15:22

weifengpy approved these changes Jul 11, 2024

View reviewed changes

awgu added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 15, 2024

pytorchmergebot closed this in b276957 Jul 15, 2024

pytorchmergebot added the Reverted label Jul 16, 2024

pytorchmergebot reopened this Jul 16, 2024

awgu mentioned this pull request Jul 17, 2024

Revert "[PT-D] Relaxed contract to allow Sequence[nn.Module] (#127773)" #130946

Closed

awgu closed this Jul 17, 2024

henrylhtsang mentioned this pull request Jul 17, 2024

[aoti] Unskip some aot inductor tests #130973

Closed

francograndegmailcom pushed a commit to francograndegmailcom/pytorch-pytorch that referenced this pull request Jul 23, 2024

[PT-D] Relaxed contract to allow Sequence[nn.Module]

d28d531

ghstack-source-id: 9841a65ca843f496f71801a6f49a65061ca82630 Pull Request resolved: pytorch/pytorch#127773

henrylhtsang mentioned this pull request Jul 31, 2024

[BE][typing] fix types in common pruning #132309

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[PT-D] Relaxed `contract` to allow `Sequence[nn.Module]` #127773

[PT-D] Relaxed `contract` to allow `Sequence[nn.Module]` #127773

awgu commented Jun 3, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jun 3, 2024 •

edited

Loading

atalman commented Jul 16, 2024

pytorchmergebot commented Jul 16, 2024

pytorchmergebot commented Jul 16, 2024

awgu commented Jul 17, 2024

[PT-D] Relaxed contract to allow Sequence[nn.Module] #127773

[PT-D] Relaxed contract to allow Sequence[nn.Module] #127773

Conversation

awgu commented Jun 3, 2024 • edited by pytorch-bot bot Loading

pytorch-bot bot commented Jun 3, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/127773

✅ No Failures

atalman commented Jul 16, 2024

pytorchmergebot commented Jul 16, 2024

pytorchmergebot commented Jul 16, 2024

awgu commented Jul 17, 2024

[PT-D] Relaxed `contract` to allow `Sequence[nn.Module]` #127773

[PT-D] Relaxed `contract` to allow `Sequence[nn.Module]` #127773

awgu commented Jun 3, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jun 3, 2024 •

edited

Loading