[tp] add kwargs support to prepare_module_input #124114
Conversation
as titled, this PR adds kwargs support to PrepareModuleInput style, where there might be modules who have only kwargs inputs but no positional args, so we should support this [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124114
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 60cb931 with merge base afa78ad ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ci-td-distributed
oncall: distributed
| if desired_layout is not None and desired_layout != input_layout: | ||
| kwarg_val = kwarg_val.redistribute(placements=(desired_layout,)) | ||
|
|
||
| prepared_kwarg_inputs[kwarg_key] = kwarg_val.to_local() if self.use_local_output else kwarg_val |
There was a problem hiding this comment.
kwarg_val can be anything(None or whatever type here) and it fails on HF models
There was a problem hiding this comment.
should be fixed now
as titled, this PR adds kwargs support to PrepareModuleInput style, where there might be modules who have only kwargs inputs but no positional args, so we should support this cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu fduwjj wz337 tianyu-l wconstab yf225 chauhang [ghstack-poisoned]
as titled, this PR adds kwargs support to PrepareModuleInput style, where there might be modules who have only kwargs inputs but no positional args, so we should support this cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu fduwjj wz337 tianyu-l wconstab yf225 chauhang [ghstack-poisoned]
wanchaol
added
release notes: distributed (dtensor)
as titled, this PR adds kwargs support to PrepareModuleInput style, where there might be modules who have only kwargs inputs but no positional args, so we should support this cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]
| input_kwarg_layouts: Optional[Dict[str, Placement]] = None, | ||
| desired_input_kwarg_layouts: Optional[Dict[str, Placement]] = None, |
There was a problem hiding this comment.
Can kwarg layouts include non-Tensor arguments?
There was a problem hiding this comment.
no only tensor arguments should allow layouts to appear, we explicitly assert that if user passed in a layout but if it's not a tensor, then we throw
🤔 yeah I think we can interate the |
|
@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 |
as titled, this PR adds kwargs support to PrepareModuleInput style, where there might be modules who have only kwargs inputs but no positional args, so we should support this Pull Request resolved: #124114 Approved by: https://github.com/XilunWu
4 participants
Stack from ghstack (oldest at bottom):
as titled, this PR adds kwargs support to PrepareModuleInput style,
where there might be modules who have only kwargs inputs but no
positional args, so we should support this
cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k