diff --git a/.ci/docker/ci_commit_pins/pytorch.txt b/.ci/docker/ci_commit_pins/pytorch.txt index e7c2ac92ddf..3d53f1cef1d 100644 --- a/.ci/docker/ci_commit_pins/pytorch.txt +++ b/.ci/docker/ci_commit_pins/pytorch.txt @@ -1 +1 @@ -e47e8794499a4a0130ff4efb8713ff93f4b40c36 \ No newline at end of file +c8a648d4dffb9f0133ff4a2ea0e660b42105d3ad diff --git a/examples/models/llama3_2_vision/preprocess/export_preprocess.py b/examples/models/llama3_2_vision/preprocess/export_preprocess.py index d82f79c2f35..550d1bfb48d 100644 --- a/examples/models/llama3_2_vision/preprocess/export_preprocess.py +++ b/examples/models/llama3_2_vision/preprocess/export_preprocess.py @@ -24,7 +24,15 @@ def main(): strict=False, ) - # Executorch + # AOTInductor. Note: export AOTI before ExecuTorch, as + # ExecuTorch will modify the ExportedProgram. + torch._inductor.aot_compile( + ep.module(), + model.get_example_inputs(), + options={"aot_inductor.output_path": "preprocess_aoti.so"}, + ) + + # Executorch. edge_program = to_edge( ep, compile_config=EdgeCompileConfig(_check_ir_validity=False) ) @@ -32,21 +40,6 @@ def main(): with open("preprocess_et.pte", "wb") as file: et_program.write_to_file(file) - # Export. - # ep = torch.export.export( - # model.get_eager_model(), - # model.get_example_inputs(), - # dynamic_shapes=model.get_dynamic_shapes(), - # strict=False, - # ) - # - # # AOTInductor - # torch._inductor.aot_compile( - # ep.module(), - # model.get_example_inputs(), - # options={"aot_inductor.output_path": "preprocess_aoti.so"}, - # ) - if __name__ == "__main__": main() diff --git a/examples/models/llama3_2_vision/preprocess/test_preprocess.py b/examples/models/llama3_2_vision/preprocess/test_preprocess.py index 83a05877495..4c0a5635e5c 100644 --- a/examples/models/llama3_2_vision/preprocess/test_preprocess.py +++ b/examples/models/llama3_2_vision/preprocess/test_preprocess.py @@ -26,6 +26,7 @@ ) from PIL import Image +from torch._inductor.package import package_aoti from torchtune.models.clip.inference._transform import CLIPImageTransform @@ -55,8 +56,10 @@ def initialize_models(resize_to_max_canvas: bool) -> Dict[str, Any]: possible_resolutions=None, ) + # Eager model. model = CLIPImageTransformModel(config) + # Exported model. exported_model = torch.export.export( model.get_eager_model(), model.get_example_inputs(), @@ -64,22 +67,35 @@ def initialize_models(resize_to_max_canvas: bool) -> Dict[str, Any]: strict=False, ) - # aoti_path = torch._inductor.aot_compile( - # exported_model.module(), - # model.get_example_inputs(), - # ) + # AOTInductor model. + so = torch._export.aot_compile( + exported_model.module(), + args=model.get_example_inputs(), + options={"aot_inductor.package": True}, + dynamic_shapes=model.get_dynamic_shapes(), + ) + aoti_path = "preprocess.pt2" + package_aoti(aoti_path, so) edge_program = to_edge( exported_model, compile_config=EdgeCompileConfig(_check_ir_validity=False) ) executorch_model = edge_program.to_executorch() + # Re-export as ExecuTorch edits the ExportedProgram. + exported_model = torch.export.export( + model.get_eager_model(), + model.get_example_inputs(), + dynamic_shapes=model.get_dynamic_shapes(), + strict=False, + ) + return { "config": config, "reference_model": reference_model, "model": model, "exported_model": exported_model, - # "aoti_path": aoti_path, + "aoti_path": aoti_path, "executorch_model": executorch_model, } @@ -265,11 +281,13 @@ def run_preprocess( ), f"Executorch model: expected {reference_ar} but got {et_ar.tolist()}" # Run aoti model and check it matches reference model. - # aoti_path = models["aoti_path"] - # aoti_model = torch._export.aot_load(aoti_path, "cpu") - # aoti_image, aoti_ar = aoti_model(image_tensor, inscribed_size, best_resolution) - # self.assertTrue(torch.allclose(reference_image, aoti_image)) - # self.assertEqual(reference_ar, aoti_ar.tolist()) + aoti_path = models["aoti_path"] + aoti_model = torch._inductor.aoti_load_package(aoti_path) + aoti_image, aoti_ar = aoti_model(image_tensor, inscribed_size, best_resolution) + assert_expected(aoti_image, reference_image, rtol=0, atol=1e-4) + assert ( + reference_ar == aoti_ar.tolist() + ), f"AOTI model: expected {reference_ar} but got {aoti_ar.tolist()}" # This test setup mirrors the one in torchtune: # https://github.com/pytorch/torchtune/blob/main/tests/torchtune/models/clip/test_clip_image_transform.py diff --git a/install_requirements.py b/install_requirements.py index f94c403c43a..40d2a3820c1 100644 --- a/install_requirements.py +++ b/install_requirements.py @@ -112,7 +112,7 @@ def python_is_compatible(): # NOTE: If a newly-fetched version of the executorch repo changes the value of # NIGHTLY_VERSION, you should re-run this script to install the necessary # package versions. -NIGHTLY_VERSION = "dev20241030" +NIGHTLY_VERSION = "dev20241101" # The pip repository that hosts nightly torch packages. TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cpu"