Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Re-support kwargs at run time #929

Merged
merged 5 commits into from
Jan 19, 2024
Merged

Re-support kwargs at run time #929

merged 5 commits into from
Jan 19, 2024

Conversation

kwen2501
Copy link
Contributor

Description

Implements #928

Users want the first pipeline stage to accept kwargs if the original program does.
This is controlled by the _codegen field of the graph as @angelayi suggests, so we make a copy from the traced program to submod0.

Feature/Issue validation/testing

Added kwargs in test_fwd.py.
Also changed a few HF examples to directly kwargs.

pippy/IR.py Outdated
# To prevent conflicts with field names, the method and attribute names
# start with an underscore
submod0.graph._codegen.pytree_info = (
submod0.graph._codegen.pytree_info._replace(out_spec=None)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do you want to set it to be None? Then the output structure will not match the eager module -- but not sure if you care about this.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because the original program's output format should be applied to the last stage, not stage 0 :)

@kwen2501 kwen2501 merged commit 5025063 into main Jan 19, 2024
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants