Skip to content
This repository was archived by the owner on Aug 1, 2025. It is now read-only.

Try finally block for with context on graph break instruction#213

Merged
anijain2305 merged 6 commits intopytorch:mainfrom
anijain2305:with-ctx
May 6, 2022
Merged

Try finally block for with context on graph break instruction#213
anijain2305 merged 6 commits intopytorch:mainfrom
anijain2305:with-ctx

Conversation

@anijain2305
Copy link
Copy Markdown
Contributor

Fixes #207 and also hf_Reformer with AOTAutograd training.

Relying on CI to test for different python versions.

@anijain2305
Copy link
Copy Markdown
Contributor Author

@jansel this is ready for review.

return [
codegen.create_load_global("torch"),
codegen.create_load_attr("_C"),
create_instruction("LOAD_METHOD", 2, "_set_grad_enabled"),
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I don't think we will know that name index "2" will be "_set_grad_enabled".

LOAD_METHOD looks up code.co_names[2] so we need to add _set_grad_enabled to co_names if it is not there already.

We have some helpers for this already in codegen.

return ([], [])

def set_grad_insts(mode):
codegen.load_import_from("torch", "_C")
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

This has the side effect creating a bunch of instruction in codegen.output_instructions. These won't be split into the prologue/epilogue. Perhaps we should refactor this to return the instructions you want.

codegen.load_import_from("torch", "_C")
codegen.load_import_from("torch._C", "_set_grad_enabled")
return [
codegen.create_load_global("torch"),
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The user might have overloaded the global variable torch (load_import_from checks for this)

@anijain2305
Copy link
Copy Markdown
Contributor Author

Thanks for the review. I have incorporated your comments. PTAL.

@anijain2305 anijain2305 merged commit d5fc452 into pytorch:main May 6, 2022
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Bug - Missing with context on the graph break instruction

3 participants