From 84de27e07cab99f186bbbfd6bab84010170686fc Mon Sep 17 00:00:00 2001 From: Mason Chang Date: Fri, 2 Feb 2024 20:43:54 +0000 Subject: [PATCH 1/3] Add docs on how to integrate with dynamo --- CONTRIBUTING.md | 1 + TROUBLESHOOTING.md | 4 ++++ docs/dynamo.md | 27 +++++++++++++++++++++++++-- 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index a06905086c5..9d7e33e7ae3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -107,6 +107,7 @@ Then run `test/run_tests.sh` and `test/cpp/run_tests.sh` to verify the setup is ### Useful materials 1. [OP Lowering Guide](https://github.com/pytorch/xla/blob/master/OP_LOWERING_GUIDE.md) 2. [CODEGEN MIGRATION GUIDE](https://github.com/pytorch/xla/blob/master/CODEGEN_MIGRATION_GUIDE.md) +3. [Dynamo Integration Guide](https://github.com/pytorch/xla/blob/master/docs/dynamo.md) ### Sharp Edges diff --git a/TROUBLESHOOTING.md b/TROUBLESHOOTING.md index 79b071945a1..92ff18c31ee 100644 --- a/TROUBLESHOOTING.md +++ b/TROUBLESHOOTING.md @@ -58,6 +58,10 @@ report sent to us if you have it. You can enable the PyTorch/XLA debugging tool by setting `PT_XLA_DEBUG=1`, which provides a couple useful debugging features. +## PyTorch/XLA + Dynamo Debugging Tool + +You can enable the PyTorch/XLA + Dynamo debugging tool by setting `XLA_DYNAMO_DEBUG=1`, which provides a couple useful debugging features. + ### Perform A Auto-Metrics Analysis The debugging tool will analyze the metrics report and provide a summary. Some example output would be diff --git a/docs/dynamo.md b/docs/dynamo.md index 1a9ba529646..e9e81770534 100644 --- a/docs/dynamo.md +++ b/docs/dynamo.md @@ -1,9 +1,32 @@ ## TorchDynamo(torch.compile) integration in PyTorch XLA -[TorchDynamo](https://pytorch.org/docs/stable/dynamo/index.html) is a Python-level JIT compiler designed to make unmodified PyTorch programs faster. It provides a clean API for compiler backends to hook in and its biggest feature is to dynamically modify Python bytecode right before it is executed. In the pytorch/xla 2.0 release, PyTorch/XLA provided an experimental backend for the TorchDynamo for both inference and training. +[TorchDynamo](https://pytorch.org/docs/stable/torch.compiler.html) is a Python-level JIT compiler designed to make unmodified PyTorch programs faster. It provides a clean API for compiler backends to hook in and its biggest feature is to dynamically modify Python bytecode right before it is executed. In the pytorch/xla 2.0 release, PyTorch/XLA provided an experimental backend for the TorchDynamo for both inference and training. The way that XLA bridge works is that Dynamo will provide a TorchFX graph when it recognizes a model pattern and PyTorch/XLA will use existing Lazy Tensor technology to compile the FX graph and return the compiled function. +### Integration + +Support for PyTorch/XLA and Dynamo currently exists by adding the `backend='openxla'` argument to `torch.compile`. For example: + +``` +import torch +import torch_xla.core.xla_model as xm + +def add(a, b): + a_xla = a.to(xm.xla_device()) + b_xla = b.to(xm.xla_device()) + return a_xla + b_xla + +compiled_code = torch.compile(add, backend='openxla') +print(compiled_code(torch.randn(10), torch.randn(10))) +``` + +Currently there are two different backends, that eventually will be merged into a single 'openxla' backend: + +* `backend='openxla'` - Useful for training. +* `backend='openxla_eval'` - Useful for inference. + + ### Inference Here is a small code example of running resnet18 with `torch.compile` @@ -40,7 +63,7 @@ timm_vision_transformer | 3.52 geomean | 3.04 Note -1. User will likely see better inference perfomrance by putting the inference execution in a `torch.no_grad` context. `openxla` is a `aot-autograd` backend of `torch.compile`. `Aot-autograd` will attempt to save some states for potential backward. `torch.no_grad` will help `aot-autograd` understand that it is being executed in a inference context. +1. User will likely see better inference performance by putting the inference execution in a `torch.no_grad` context. `openxla` is a `aot-autograd` backend of `torch.compile`. `Aot-autograd` will attempt to save some state for potential backward. `torch.no_grad` will help `aot-autograd` understand that it is being executed in a inference context. 2. User can also use the `openxla_eval` backend directly without `torch.no_grad`, since `openxla_eval` is not an `aot-autograd` backend and only works for inference. ### Training From d927ce170a2f702beb92ee1ecf3f6b443998c06b Mon Sep 17 00:00:00 2001 From: Mason Chang Date: Tue, 6 Feb 2024 09:52:55 -0800 Subject: [PATCH 2/3] Update TROUBLESHOOTING.md Co-authored-by: Emilio Cota --- TROUBLESHOOTING.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/TROUBLESHOOTING.md b/TROUBLESHOOTING.md index 92ff18c31ee..8829e831fe3 100644 --- a/TROUBLESHOOTING.md +++ b/TROUBLESHOOTING.md @@ -60,7 +60,7 @@ You can enable the PyTorch/XLA debugging tool by setting `PT_XLA_DEBUG=1`, which ## PyTorch/XLA + Dynamo Debugging Tool -You can enable the PyTorch/XLA + Dynamo debugging tool by setting `XLA_DYNAMO_DEBUG=1`, which provides a couple useful debugging features. +You can enable the PyTorch/XLA + Dynamo debugging tool by setting `XLA_DYNAMO_DEBUG=1`. ### Perform A Auto-Metrics Analysis From 7c50cdef90e9dcdcbf82523a75bc4106d53a3191 Mon Sep 17 00:00:00 2001 From: Mason Chang Date: Tue, 6 Feb 2024 09:53:20 -0800 Subject: [PATCH 3/3] Update docs/dynamo.md Co-authored-by: Emilio Cota --- docs/dynamo.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dynamo.md b/docs/dynamo.md index e9e81770534..67d36ff536b 100644 --- a/docs/dynamo.md +++ b/docs/dynamo.md @@ -63,7 +63,7 @@ timm_vision_transformer | 3.52 geomean | 3.04 Note -1. User will likely see better inference performance by putting the inference execution in a `torch.no_grad` context. `openxla` is a `aot-autograd` backend of `torch.compile`. `Aot-autograd` will attempt to save some state for potential backward. `torch.no_grad` will help `aot-autograd` understand that it is being executed in a inference context. +1. User will likely see better inference performance by putting the inference execution in a `torch.no_grad` context. `openxla` is an `aot-autograd` backend of `torch.compile`; `aot-autograd` attempts to save some state for a potential backward pass. Setting `torch.no_grad` helps `aot-autograd` understand that it is being executed in an inference context. 2. User can also use the `openxla_eval` backend directly without `torch.no_grad`, since `openxla_eval` is not an `aot-autograd` backend and only works for inference. ### Training