[dynamo] FSDP + AC + torch.compile #103953
Conversation
[ghstack-poisoned]
anijain2305
requested review from
H-Huang,
awgu,
fduwjj,
fegin,
kwen2501,
mrshenli,
rohan-varma,
wanchaol and
zhaojuanmao
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/103953
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 7c3c297: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 [ghstack-poisoned]
anijain2305
changed the title
| @skip_if_lt_x_gpu(1) | ||
| @unittest.skipIf(not has_triton(), "Inductor+gpu needs triton and recent GPU arch") | ||
| def test_fsdp_activation_checkpointing(self): | ||
| from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( |
There was a problem hiding this comment.
do we want to also test the native torch.utils one too?
There was a problem hiding this comment.
Thats already tested in test/dynamo/test_activation_checkpointing.py
They are also tested heavily on HF models - more data here - #102935
| with _dynamo_dist_per_rank_init(self.rank, self.world_size): | ||
| model, inputs = get_toy_model_for_activation_checkpointing(f"cuda:{self.rank}") | ||
| is_inner = lambda module: isinstance(module, ToyInnerModel) # noqa: E731 | ||
| wrap_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=is_inner) |
There was a problem hiding this comment.
are there any valid cases where we don't use the same policy for fsdp/ac that we need to test? cc @awgu
|
|
||
| gm.__name__ = next_name | ||
| src = NNModuleSource(GetItemSource(self.source, next_name)) | ||
| if self.source.guard_source().is_fsdp_module(): |
There was a problem hiding this comment.
what does 'add_subgraph' do generally? are there more than one subgraph per higher order operator?
There was a problem hiding this comment.
It is just a util that inserts the subgraph into the output graph module of Dynamo. Yes, there is no limit on how many subgraphs we can have. torch.cond, for example, has 2 - one for true branch and one for false.
In the case of activation checkpointing, there is only 1.
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 [ghstack-poisoned]
There was a problem hiding this comment.
lgtm, the CI failure looks correlated https://hud.pytorch.org/pr/pytorch/pytorch/103953#14451950178
| model = FSDP( | ||
| copy.deepcopy(model), | ||
| auto_wrap_policy=wrap_policy, | ||
| use_orig_params=True |
There was a problem hiding this comment.
shall we paramterize the tests and test both use_orig_params == True/False similar to https://github.com/pytorch/pytorch/blob/main/test/distributed/fsdp/test_fsdp_checkpoint.py#L137?
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 [ghstack-poisoned]
| k: variables.ConstantVariable(v) for k, v in self.value.keywords.items() | ||
| } | ||
| partial_kwargs.update(kwargs) | ||
| if requires_higher_order_op(self.value.func): |
There was a problem hiding this comment.
This seems like a good opportunity to use the walrus operator.
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 [ghstack-poisoned]
|
@pytorchbot merge |
Merge failedReason: This PR needs a If not, please add the To add a label, you can comment to pytorchbot, for example For more information, see Details for Dev Infra teamRaised by workflow job |
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 [ghstack-poisoned]
|
@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 |
|
Do you mind testing this against #103711 ? I know its an onerous ask, but I want to make sure we also keep in mind the relative future direction of PT2 + FSDP for things like this. |
Oh missed this message. Will certainly do in next couple of weeks. Trying to wrap up some work, and will pick this up. |
6 participants
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @chenyang78