[FSDP2] Added test to show rank 0 broadcast meta-device flow #124651

awgu · 2024-04-22T20:55:41Z

Stack from ghstack (oldest at bottom):

This PR includes two things:

Changes to support load_state_dict(assign=True)
- These changes are not ideal, but until we have DTensor padding the local tensor and general swap_tensors adoption, we may need to make do.
Example of how to convert a full state dict on rank 0 to sharded state dict on all ranks via broadcast
- ~~To-do: check for recordStream from the funcol broadcast; if being called, remediate either via async_op=False c10d broadcast or use TORCH_NCCL_AVOID_RECORD_STREAMS=1~~ switched to using c10d async_op=False broadcast
- To-do: check for broadcast latency since not using any coalescing

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

[ghstack-poisoned]

pytorch-bot · 2024-04-22T20:55:44Z

🔗 Helpful Links

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

📄 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 1f5d795 with merge base c82fcb7 ():
💚 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 rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

awgu · 2024-04-22T22:03:32Z

torch/distributed/_composable/fsdp/_fsdp_param.py

+        # TODO: Remove this padding logic once DTensor pads the local tensor:
+        # https://github.com/pytorch/pytorch/issues/113045
+        self._post_load_hook_handle = (
+            module_info.module.register_load_state_dict_post_hook(
+                lambda *args, **kwargs: self.reset_sharded_param()
+            )
+        )


This is needed to support load_state_dict(assign=True) today.

In the future, we want (1) DTensor to pad its local tensor and (2) for users to use the swap_tensors path for load_state_dict() with FSDP2. Together these should allow us to remove this hook.

This PR includes two things: 1. Changes to support `load_state_dict(assign=True)` - These changes are not ideal, but until we have `DTensor` padding the local tensor and general `swap_tensors` adoption, we may need to make do. 2. Example of how to convert a full state dict on rank 0 to sharded state dict on all ranks via broadcast - To-do: check for `recordStream` from the funcol broadcast; if being called, remediate either via `async_op=False` c10d broadcast or use `TORCH_NCCL_AVOID_RECORD_STREAMS=1` - To-do: check for broadcast latency since not using any coalescing cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

test/distributed/_composable/fsdp/test_fully_shard_init.py

This PR includes two things: 1. Changes to support `load_state_dict(assign=True)` - These changes are not ideal, but until we have `DTensor` padding the local tensor and general `swap_tensors` adoption, we may need to make do. 2. Example of how to convert a full state dict on rank 0 to sharded state dict on all ranks via broadcast - ~~To-do: check for `recordStream` from the funcol broadcast; if being called, remediate either via `async_op=False` c10d broadcast or use `TORCH_NCCL_AVOID_RECORD_STREAMS=1`~~ switched to using c10d `async_op=False` broadcast - To-do: check for broadcast latency since not using any coalescing cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

ghstack-source-id: 02d6620 Pull Request resolved: #124651

awgu · 2024-04-23T16:01:52Z

test/distributed/_composable/fsdp/test_fully_shard_state_dict.py

@@ -148,9 +148,14 @@ def _test_state_dict_save_load(self, model: nn.Module):
                param.to_local(),
                torch.ones_like(param.to_local()) * new_fill_value,
            )
-            self.assertEqual(


The assertion was too strict before, and our implementation was actually incorrect since we did not re-pad the local tensor.

This PR includes two things: 1. Changes to support `load_state_dict(assign=True)` - These changes are not ideal, but until we have `DTensor` padding the local tensor and general `swap_tensors` adoption, we may need to make do. 2. Example of how to convert a full state dict on rank 0 to sharded state dict on all ranks via broadcast - ~~To-do: check for `recordStream` from the funcol broadcast; if being called, remediate either via `async_op=False` c10d broadcast or use `TORCH_NCCL_AVOID_RECORD_STREAMS=1`~~ switched to using c10d `async_op=False` broadcast - To-do: check for broadcast latency since not using any coalescing cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

wanchaol

lgtm!

weifengpy · 2024-04-24T08:02:54Z

test/distributed/_composable/fsdp/test_fully_shard_init.py

+                sharded_tensor = distribute_tensor(
+                    full_param, mesh, sharded_meta_param.placements
+                )
+                sharded_sd[param_name] = nn.Parameter(sharded_tensor)


curious if buffers are sometimes included in state dict? for llama2, the buffer is RotaryPositionalEmbeddings.theta but it's not included in HF state dict

This PR adds a unit test to show how we can convert FSDP2 GPU sharded state dicts to a CPU full state dict on rank 0. Pull Request resolved: #124741 Approved by: https://github.com/wanchaol, https://github.com/wz337 ghstack dependencies: #124651

This PR makes sure to construct the `DeviceMesh`'s `mesh` tensor on CPU device in `init_device_mesh()`. This means that we can call `init_device_mesh()` under meta-device context and still construct the correct `mesh` tensor. Pull Request resolved: #124767 Approved by: https://github.com/wz337 ghstack dependencies: #124651, #124741

Pull Request resolved: #124768 Approved by: https://github.com/wz337 ghstack dependencies: #124651, #124741, #124767

ghstack-source-id: a7b16b5 Pull Request resolved: pytorch#124651

This PR adds a `DeviceMesh.from_group()` static method to convert an existing process group to a device mesh. Motivation: We need `DeviceMesh.from_group()` to allow FSDP2 to interoperate with distributed libraries that do not use `DeviceMesh` for all parallelisms. Pull Request resolved: #124787 Approved by: https://github.com/wanchaol ghstack dependencies: #124651, #124741, #124767, #124768, #124780

This PR adds a `DeviceMesh.from_group()` static method to convert an existing process group to a device mesh. Motivation: We need `DeviceMesh.from_group()` to allow FSDP2 to interoperate with distributed libraries that do not use `DeviceMesh` for all parallelisms. Pull Request resolved: pytorch#124787 Approved by: https://github.com/wanchaol ghstack dependencies: pytorch#124651, pytorch#124741, pytorch#124767, pytorch#124768, pytorch#124780

This PR renames the `FSDP` class to `FSDPModule`. This is a BC breaking change. The rationale is that `FSDPModule` is more descriptive since `fully_shard` is a module-level API (applied to a `module` arg), so the `FSDP` class will always correspond to a module. Also, users commonly import `FullyShardedDataParallel` as `FSDP`, so this can help avoid some name conflict in some cases. Pull Request resolved: #124955 Approved by: https://github.com/wanchaol, https://github.com/wconstab ghstack dependencies: #124651, #124741, #124767, #124768, #124780, #124787

…#124651) This PR includes two things: 1. Changes to support `load_state_dict(assign=True)` - These changes are not ideal, but until we have `DTensor` padding the local tensor and general `swap_tensors` adoption, we may need to make do. 2. Example of how to convert a full state dict on rank 0 to sharded state dict on all ranks via broadcast - ~~To-do: check for `recordStream` from the funcol broadcast; if being called, remediate either via `async_op=False` c10d broadcast or use `TORCH_NCCL_AVOID_RECORD_STREAMS=1`~~ switched to using c10d `async_op=False` broadcast - To-do: check for broadcast latency since not using any coalescing Pull Request resolved: pytorch#124651 Approved by: https://github.com/wanchaol

…24741) This PR adds a unit test to show how we can convert FSDP2 GPU sharded state dicts to a CPU full state dict on rank 0. Pull Request resolved: pytorch#124741 Approved by: https://github.com/wanchaol, https://github.com/wz337 ghstack dependencies: pytorch#124651

This PR makes sure to construct the `DeviceMesh`'s `mesh` tensor on CPU device in `init_device_mesh()`. This means that we can call `init_device_mesh()` under meta-device context and still construct the correct `mesh` tensor. Pull Request resolved: #124767 Approved by: https://github.com/wz337 ghstack dependencies: #124651, #124741

Pull Request resolved: #124768 Approved by: https://github.com/wz337 ghstack dependencies: #124651, #124741, #124767

This PR adds a `DeviceMesh.from_group()` static method to convert an existing process group to a device mesh. Motivation: We need `DeviceMesh.from_group()` to allow FSDP2 to interoperate with distributed libraries that do not use `DeviceMesh` for all parallelisms. Pull Request resolved: #124787 Approved by: https://github.com/wanchaol ghstack dependencies: #124651, #124741, #124767, #124768, #124780

This PR renames the `FSDP` class to `FSDPModule`. This is a BC breaking change. The rationale is that `FSDPModule` is more descriptive since `fully_shard` is a module-level API (applied to a `module` arg), so the `FSDP` class will always correspond to a module. Also, users commonly import `FullyShardedDataParallel` as `FSDP`, so this can help avoid some name conflict in some cases. Pull Request resolved: pytorch#124955 Approved by: https://github.com/wanchaol, https://github.com/wconstab ghstack dependencies: pytorch#124651, pytorch#124741, pytorch#124767, pytorch#124768, pytorch#124780, pytorch#124787

[FSDP2] Added test to show rank 0 broadcast meta-device flow

e49a2b5

[ghstack-poisoned]

pytorch-bot bot added ci-td-distributed oncall: distributed Add this issue/PR to distributed oncall triage queue topic: not user facing topic category labels Apr 22, 2024

awgu commented Apr 22, 2024

View reviewed changes

test/distributed/_composable/fsdp/test_fully_shard_init.py Show resolved Hide resolved

awgu commented Apr 22, 2024

View reviewed changes

test/distributed/_composable/fsdp/test_fully_shard_init.py Show resolved Hide resolved

awgu marked this pull request as ready for review April 22, 2024 22:33

awgu pushed a commit that referenced this pull request Apr 23, 2024

[FSDP2] Added test to show rank 0 broadcast meta-device flow

98823a0

ghstack-source-id: 02d6620 Pull Request resolved: #124651

awgu requested a review from weifengpy April 23, 2024 01:00

awgu added the release notes: distributed (fsdp2) release notes category label Apr 23, 2024

awgu requested review from fegin and wz337 April 23, 2024 14:34

awgu commented Apr 23, 2024

View reviewed changes

awgu mentioned this pull request Apr 23, 2024

[FSDP2] Added test to show rank 0 CPU full state dict flow #124741

Closed

This was referenced Apr 23, 2024

[DeviceMesh] Initialized mesh tensor with CPU context #124767

Closed

[DeviceMesh] Removed unneeded .to(cpu) #124768

Closed

wanchaol approved these changes Apr 23, 2024

View reviewed changes

pytorchmergebot added the Merged label Apr 24, 2024

pytorchmergebot closed this in e8f9f37 Apr 24, 2024

pytorchmergebot removed the merging label Apr 24, 2024

weifengpy mentioned this pull request Apr 24, 2024

enable LoRA + FSDP2 pytorch/torchtune#855

Merged

6 tasks

weifengpy reviewed Apr 24, 2024

View reviewed changes

pytorchmergebot pushed a commit that referenced this pull request Apr 24, 2024

[DeviceMesh] Removed unneeded .to(cpu) (#124768)

48312a7

Pull Request resolved: #124768 Approved by: https://github.com/wz337 ghstack dependencies: #124651, #124741, #124767

awgu pushed a commit to awgu/pytorch that referenced this pull request Apr 24, 2024

[FSDP2] Added test to show rank 0 broadcast meta-device flow

8156a76

ghstack-source-id: a7b16b5 Pull Request resolved: pytorch#124651

This was referenced Apr 25, 2024

[RFC][FSDP2] Renamed FSDP to FSDPModule #124936

Closed

[RFC][FSDP2] Renamed FSDP to FSDPModule #124955

Closed

This was referenced Apr 29, 2024

[FSDP2] Fixed fp32 param dtype/bf16 reduce dtype test #125190

Closed

[FSDP2] Accumulated in reduce_dtype if not syncing grads #125191

Closed

pytorch-bot bot pushed a commit that referenced this pull request May 3, 2024

[DeviceMesh] Removed unneeded .to(cpu) (#124768)

7943f1c

Pull Request resolved: #124768 Approved by: https://github.com/wz337 ghstack dependencies: #124651, #124741, #124767

github-actions bot deleted the gh/awgu/566/head branch June 2, 2024 02:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[FSDP2] Added test to show rank 0 broadcast meta-device flow #124651

[FSDP2] Added test to show rank 0 broadcast meta-device flow #124651

Uh oh!

awgu commented Apr 22, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Apr 22, 2024 •

edited

Loading

Uh oh!

awgu Apr 22, 2024

Uh oh!

Uh oh!

Uh oh!

awgu Apr 23, 2024

Uh oh!

wanchaol left a comment

Uh oh!

weifengpy Apr 24, 2024

Uh oh!

Uh oh!

[FSDP2] Added test to show rank 0 broadcast meta-device flow #124651

[FSDP2] Added test to show rank 0 broadcast meta-device flow #124651

Uh oh!

Conversation

awgu commented Apr 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Apr 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

awgu Apr 22, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

awgu Apr 23, 2024

Choose a reason for hiding this comment

Uh oh!

wanchaol left a comment

Choose a reason for hiding this comment

Uh oh!

weifengpy Apr 24, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

awgu commented Apr 22, 2024 •

edited

Loading

pytorch-bot bot commented Apr 22, 2024 •

edited

Loading