diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index a06905086c59..9d7e33e7ae3a 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 79b071945a19..8829e831fe39 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`. + ### 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 1a9ba5296468..67d36ff536b5 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 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