Support basic TorchBind in aot_compile and aoti_compile_and_package #148506
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/148506
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 0728785 with merge base a6e71db ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
0ed43c8 to
12d5288
Compare
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
12d5288 to
9cc78c8
Compare
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
9cc78c8 to
4f370d4
Compare
yushangdi
added a commit
to yushangdi/pytorch
that referenced
this pull request
Mar 5, 2025
…kage (pytorch#148506) Summary: Pull Request resolved: pytorch#148506 **Codegen** - Skip some codegen parts for torchbind (such as arg decleration) because they are loaded in proxy executor, so we do not need to declare torchbind args in cpp code - Added a helper method to get the schema of CallTorchBind HOP. The returned schema is only the schema of `obj.method()`. **Serialization** Add support for torchbind object in serialization - For CallTorchBind HOP, we need to handle it specially because of it's schema. The output serialized args is in the format of `(obj, method, *args, **kwargs)`. - it.TorchBindObject inputs are serialized to `as_custom_obj` Argument. **Packaging** Add torchbind objects file and `custom_objs_config.json` file to generated files output of `aot_compile`. The json file is stored in the `data/aotinductor/<model_name>` folder in pt2 archive. The torchbind objects are stored in data/constants/ folder in pt2 archive. The format of torchbind objects are `f"{CUSTOM_OBJ_FILENAME_PREFIX}{custom_obj_idx}"`. e.g. `custom_obj_0`. CustomClassHolder objects implement their own pickle methods. Note that this `custom_objs_config.json` file is different from the `model_constants_config.json` file produced in package_sigmoid(). The keys in `custom_objs_config` directly correspond to the arg name in extern nodes json. The key in `model_constants_config.json` produced by `package_sigmoid` is the attribute name in the user mode code. This is required for both internal and OSS torchbind support. For OSS torchbind support, we also need to package torchbind_constants into the .pt2 output. **Work Left** We still need to add torchbind support in ProxyExecutor for inductor.aoti_load_package to work. See other diffs in the stack. Test Plan: ``` buck run fbcode//mode/dev-nosan //caffe2/test/inductor:torchbind -- -r schema buck run fbcode//mode/dev-nosan //caffe2/test/inductor:torchbind -- -r aot_compile ``` Differential Revision: D69490718
4f370d4 to
4a2a7d0
Compare
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
4a2a7d0 to
440fbbe
Compare
yushangdi
added a commit
to yushangdi/pytorch
that referenced
this pull request
Mar 5, 2025
…kage (pytorch#148506) Summary: **Codegen** - Skip some codegen parts for torchbind (such as arg decleration) because they are loaded in proxy executor, so we do not need to declare torchbind args in cpp code - Added a helper method to get the schema of CallTorchBind HOP. The returned schema is only the schema of `obj.method()`. **Serialization** Add support for torchbind object in serialization - For CallTorchBind HOP, we need to handle it specially because of it's schema. The output serialized args is in the format of `(obj, method, *args, **kwargs)`. - it.TorchBindObject inputs are serialized to `as_custom_obj` Argument. **Packaging** Add torchbind objects file and `custom_objs_config.json` file to generated files output of `aot_compile`. The json file is stored in the `data/aotinductor/<model_name>` folder in pt2 archive. The torchbind objects are stored in data/constants/ folder in pt2 archive. The format of torchbind objects are `f"{CUSTOM_OBJ_FILENAME_PREFIX}{custom_obj_idx}"`. e.g. `custom_obj_0`. CustomClassHolder objects implement their own pickle methods. Note that this `custom_objs_config.json` file is different from the `model_constants_config.json` file produced in package_sigmoid(). The keys in `custom_objs_config` directly correspond to the arg name in extern nodes json. The key in `model_constants_config.json` produced by `package_sigmoid` is the attribute name in the user mode code. This is required for both internal and OSS torchbind support. For OSS torchbind support, we also need to package torchbind_constants into the .pt2 output. **Work Left** We still need to add torchbind support in ProxyExecutor for inductor.aoti_load_package to work. See other diffs in the stack. Test Plan: ``` buck run fbcode//mode/dev-nosan //caffe2/test/inductor:torchbind -- -r schema buck run fbcode//mode/dev-nosan //caffe2/test/inductor:torchbind -- -r aot_compile ``` Differential Revision: D69490718
440fbbe to
827a796
Compare
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
827a796 to
8c93058
Compare
yushangdi
added a commit
to yushangdi/pytorch
that referenced
this pull request
Mar 8, 2025
…kage (pytorch#148506) Summary: Pull Request resolved: pytorch#148506 **Codegen** - Skip some codegen parts for torchbind (such as arg decleration) because they are loaded in proxy executor, so we do not need to declare torchbind args in cpp code - Added a helper method to get the schema of CallTorchBind HOP. **Serialization** Add support for torchbind object in serialization - For CallTorchBind HOP, we need to handle it specially because of it's schema. - it.TorchBindObject inputs are serialized to `as_custom_obj` Argument. **Packaging** Add torchbind objects file and `custom_objs_config.json` file to generated files output of `aot_compile`. The json file is stored in the `data/aotinductor/<model_name>` folder in pt2 archive. The torchbind objects are stored in data/constants/ folder in pt2 archive. The format of torchbind objects are `f"{CUSTOM_OBJ_FILENAME_PREFIX}{custom_obj_idx}"`. e.g. `custom_obj_0`. CustomClassHolder objects implement their own pickle methods. Note that this `custom_objs_config.json` file is different from the `model_constants_config.json` file produced in package_sigmoid(). The keys in `custom_objs_config` directly correspond to the arg name in extern nodes json. The key in `model_constants_config.json` produced by `package_sigmoid` is the attribute name in the user mode code. This is required for both internal and OSS torchbind support. For OSS torchbind support, we also need to package torchbind_constants into the .pt2 output. **Work Left** We still need to add torchbind support in ProxyExecutor for inductor.aoti_load_package to work. See other diffs in the stack. Test Plan: ``` buck run fbcode//mode/dev-nosan //caffe2/test/inductor:torchbind -- -r schema buck run fbcode//mode/dev-nosan //caffe2/test/inductor:torchbind -- -r aot_compile buck run fbcode//mode/dev-nosan //caffe2/test:test_export -- -r test_custom_class buck run fbcode//mode/dev-nosan //caffe2/test/inductor:test_inductor -- -r aoti_eager_dtype_device_layout ``` Differential Revision: D69490718
f112492 to
d1e942e
Compare
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
d1e942e to
86c9ced
Compare
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
86c9ced to
b36655c
Compare
…kage (pytorch#148506) Summary: Pull Request resolved: pytorch#148506 **Codegen** - Skip some codegen parts for torchbind (such as arg decleration) because they are loaded in proxy executor, so we do not need to declare torchbind args in cpp code - Added a helper method to get the schema of CallTorchBind HOP. **Serialization** Add support for torchbind object in serialization - For CallTorchBind HOP, we need to handle it specially because of it's schema. - it.TorchBindObject inputs are serialized to `as_custom_obj` Argument. **Packaging** Add torchbind objects file and `custom_objs_config.json` file to generated files output of `aot_compile`. The json file is stored in the `data/aotinductor/<model_name>` folder in pt2 archive. The torchbind objects are stored in data/constants/ folder in pt2 archive. The format of torchbind objects are `f"{CUSTOM_OBJ_FILENAME_PREFIX}{custom_obj_idx}"`. e.g. `custom_obj_0`. CustomClassHolder objects implement their own pickle methods. Note that this `custom_objs_config.json` file is different from the `model_constants_config.json` file produced in package_sigmoid(). The keys in `custom_objs_config` directly correspond to the arg name in extern nodes json. The key in `model_constants_config.json` produced by `package_sigmoid` is the attribute name in the user mode code. This is required for both internal and OSS torchbind support. For OSS torchbind support, we also need to package torchbind_constants into the .pt2 output. **Work Left** We still need to add torchbind support in ProxyExecutor for inductor.aoti_load_package to work. See other diffs in the stack. Test Plan: ``` buck run fbcode//mode/dev-nosan //caffe2/test/inductor:torchbind -- -r schema buck run fbcode//mode/dev-nosan //caffe2/test/inductor:torchbind -- -r aot_compile buck run fbcode//mode/dev-nosan //caffe2/test:test_export -- -r test_custom_class buck run fbcode//mode/dev-nosan //caffe2/test/inductor:test_inductor -- -r aoti_eager_dtype_device_layout ``` Differential Revision: D69490718
b36655c to
796a3dd
Compare
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
1 similar comment
|
This pull request was exported from Phabricator. Differential Revision: D69490718 |
796a3dd to
0728785
Compare
|
@pytorchbot merge (Initiating merge automatically since Phabricator Diff has merged) |
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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
Codegen
obj.method().Serialization
Add support for torchbind object in serialization
(obj, method, *args, **kwargs).as_custom_objArgument.Packaging
Add torchbind objects file and
custom_objs_config.jsonfile to generated files output ofaot_compile.The json file is stored in the
data/aotinductor/<model_name>folder in pt2 archive.The torchbind objects are stored in data/constants/ folder in pt2 archive.
The format of torchbind objects are
f"{CUSTOM_OBJ_FILENAME_PREFIX}{custom_obj_idx}". e.g.custom_obj_0.CustomClassHolder objects implement their own pickle methods.
Note that this
custom_objs_config.jsonfile is different from themodel_constants_config.jsonfile produced in package_sigmoid(). The keys incustom_objs_configdirectly correspond to the arg name in extern nodes json.The key in
model_constants_config.jsonproduced bypackage_sigmoidis the attribute name in the user mode code.This is required for both internal and OSS torchbind support.
For OSS torchbind support, we also need to package torchbind_constants into the .pt2 output.
Work Left
We still need to add torchbind support in ProxyExecutor for inductor.aoti_load_package to work. See other diffs in the stack.
Test Plan:
Differential Revision: D69490718
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov