[simple_fsdp] Turn on bucketing by default #2103
Conversation
[ghstack-poisoned]
meta-cla
bot
added
the
CLA Signed
[ghstack-poisoned]
tianyu-l
left a comment
tianyu-l
left a comment
There was a problem hiding this comment.
what's the behavior before vs. after? I thought the code you are modifying already does bucketing.
Before this bucketing was not enabled. Config collective bucketing was not applied if the schedule_overlap is called manually. |
| gm: torch.fx.GraphModule, example_inputs: Any | ||
| ) -> torch.fx.GraphModule: | ||
| schedule_overlap_bucketing(gm) | ||
| schedule_overlap_bucketing(gm, collective_bucketing=True) |
There was a problem hiding this comment.
collective_bucketing and insert_overlap_deps configs are turned on in this PR: #1965. Could you confirm which is the correct way to enable this pass?
There was a problem hiding this comment.
And probably remove the unused configs
There was a problem hiding this comment.
Sorry, it's a bit confusing because we had some internal usage that didnt want the pass to depend on inductor. configs. today, those inductor configs are only used in the inductor psot grad application.
Potentially we can have a :
schedule_overlap_bucketing
and
schedule_overlap_bucketing_from_configs
where the latter reads in inductor configs. I'm not sure. open to ideas here.
There was a problem hiding this comment.
oh i see, then probably we can use this PR's config to enable aten-level aot_eager_autobucketing_reordering_pass, and the inductor config to enable inductor post grad passes in inductor_autobucketing_reordering_pass. 🤔
There was a problem hiding this comment.
Sorry, didn't fully get it. Does it mean we can remove some code for the aot_eager / inductor option in this PR? Do we have to use multiple toggles for one thing? e.g. I see the following for aot_eager
dist_opts.collective_bucketing = True
But I didn't see any special inductor configs for bucketing.
There was a problem hiding this comment.
Yes, I mean @IvanKobzarev need to update the code such that dist_opts are only put to inductor scheduling pass entry before he merges the PR....
There was a problem hiding this comment.
Could we add some comment on what each steps are doing, for better readability
There was a problem hiding this comment.
I will add a function in pytorch that schedules this from inductor configs. i think that will be clearest.
There was a problem hiding this comment.
we can now just call schedule_overlap_bucketing_from_inductor_configs and use the configs.
| gm: torch.fx.GraphModule, example_inputs: Any | ||
| ) -> torch.fx.GraphModule: | ||
| schedule_overlap_bucketing(gm) | ||
| schedule_overlap_bucketing(gm, collective_bucketing=True) |
There was a problem hiding this comment.
Sorry, didn't fully get it. Does it mean we can remove some code for the aot_eager / inductor option in this PR? Do we have to use multiple toggles for one thing? e.g. I see the following for aot_eager
dist_opts.collective_bucketing = True
But I didn't see any special inductor configs for bucketing.
Stack from ghstack (oldest at bottom):