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

Conversation

anijain2305
Copy link
Contributor

Fixes #207 and also hf_Reformer with AOTAutograd training.

Relying on CI to test for different python versions.

@anijain2305
Copy link
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
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
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
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
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