[FSDP] Fix use_orig_params=True, CPU offload, no_sync()
#100180
Conversation
[ghstack-poisoned]
awgu
requested review from
mrshenli,
zhaojuanmao,
rohan-varma,
H-Huang,
kwen2501,
wanchaol,
fegin,
kiukchung and
d4l3k
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/100180
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit b4d402e: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
release notes: distributed (fsdp)
ghstack-source-id: 8102bea96ec3551b1041a299b005e6e47f3cbbba Pull Request resolved: #100180
| @@ -117,12 +117,6 @@ def _test_grad_acc( | |||
| point to prefetch the next layer's full parameters during the | |||
| backward pass, if at all. | |||
| """ | |||
| # Gradient accumulation outside `no_sync()` is not currently compatible | |||
There was a problem hiding this comment.
Previously, we were skipping all CPU offloading tests since every tested config included use_no_sync == False 😢
That is why I thought use_orig_params=True worked with CPUOffload(True), but we were actually skipping the test.
This should fix #98494. [ghstack-poisoned]
ghstack-source-id: 8102bea96ec3551b1041a299b005e6e47f3cbbba Pull Request resolved: #100180
This should fix #98494. We follow a similar approach as in past PRs for mismatched dtype or size from running in `no_sync()`. [ghstack-poisoned]
ghstack-source-id: 4d6ae641b0016ae04de6b5ee4fc81777eeef0496 Pull Request resolved: #100180
|
|
||
| NOTE: Gradient accumulation without using the ``no_sync()`` context | ||
| manager is not currently compatible with CPU offloading, so those tests | ||
| just return directly. | ||
| manager is not currently compatible with CPU offloading. |
There was a problem hiding this comment.
curious, how do we error here?
There was a problem hiding this comment.
We do not error. It is silently incorrect if I understand correctly.
| ) | ||
|
|
||
| @skip_if_lt_x_gpu(2) | ||
| @parametrize("use_orig_params", [False, True]) |
There was a problem hiding this comment.
should we put this as a subtest?
There was a problem hiding this comment.
I think we can avoid subtesting major config differences, where we actually expect based on the implementation there may be a difference. At least that is what I have been doing so far.
| # NOTE: This is a hack using `.data` to side step the check | ||
| # that parameter/gradient sizes/dtypes/devices match. From | ||
| # calling `reshard()`, `param` has the sharded size, the full | ||
| # precision dtype, and is on CPU. Thus, one or more of the |
There was a problem hiding this comment.
is on CPU only if CPU offloading?
| # calling `reshard()`, `param` has the sharded size, the full | ||
| # precision dtype, and is on CPU. Thus, one or more of the | ||
| # following cases can hold when in `no_sync()`: | ||
| # 1. `view` can have the unsharded size. |
There was a problem hiding this comment.
view is the grad here right, can we clarify that?
There was a problem hiding this comment.
Let me do it in a follow-up to avoid re-triggering CI.
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
…100180) This should fix pytorch#98494. We follow a similar approach as in past PRs for mismatched dtype or size from running in `no_sync()`. Pull Request resolved: pytorch#100180 Approved by: https://github.com/rohan-varma
Stack from ghstack (oldest at bottom):
use_orig_params=True, CPU offload,no_sync()#100180This should fix #98494. We follow a similar approach as in past PRs for mismatched dtype or size from running in
no_sync().