# PyTorch/XLA 2.3 dev update **URL:** \ **Category:** compiler\ **Created:** [May 29, 2024, 4:19am UTC](https://dev-discuss.pytorch.org/t/pytorch-xla-2-3-dev-update/2124 "2024-05-29T04:19:31Z")\ **Posts on this page:** 1\ **Page:** 1
**Author:** ![JackCaoG](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/jackcaog/32/819_2.png) [@JackCaoG](https://dev-discuss.pytorch.org/u/JackCaoG)\ **Post date:** [May 29, 2024, 4:19am UTC](https://dev-discuss.pytorch.org/t/pytorch-xla-2-3-dev-update/2124/1 "2024-05-29T04:19:31Z")
Hey I am here to give a late update for the PyTorch/XLA 2.3 release. Similar to my previous update, you can check our [release note](https://github.com/pytorch/xla/releases/tag/v2.3.0) for detailed updates. I am going to highlight some of the new features and share how I think about them. ## SPMD Auto-sharding We launched the experimental support for single host TPU auto-sharding in the 2.3 release. It can be enabled via torch-xla’s custom API ```auto import torch_xla.runtime as xr xr.use_spmd(auto=True) ``` Or you can enable it with the distributed module API ```auto import torch_xla.runtime as xr from torch.distributed._tensor import DeviceMesh, distribute_module from torch_xla.distributed.spmd import auto_policy device_count = xr.global_runtime_device_count() device_mesh = DeviceMesh("xla", list(range(device_count))) # Currently, model should be loaded to xla device via distribute_module. model = MyModule() # nn.module sharded_model = distribute_module(model, device_mesh, auto_policy) ``` You can find more details about the auto-sharding in [here](https://github.com/pytorch/xla/blob/r2.3/docs/spmd.md#auto-sharding). IMO auto-sharding is a great way to experience the capabilities of the GSPMD. We benchmarked the popular transformer models like GPT2, Llama2, and Gemma, auto-sharding can achieve ~ 95% of MFU compared to the carefully optimized alternatives on Cloud TPU v4-8. Auto-sharding does incur a higher compile time, hence we strongly recommend user to turn on the persistent compilation cache following [this guide](https://github.com/pytorch/xla/blob/master/API_GUIDE.md#compilation-caching). My 2 cents about the auto-sharding is that it is a cool feature that can give us a performance baseline and inspire us on how to shard the model. However long compilation time makes it a bit challenging to use if you are constantly changing the model code. For most of the transformer-like models my recommendation will be first try the FSDPv2(I am going to talk about this in the following section) and see if it works. ## FSDPv2 (FSDP via SPMD) We are still trying to figure out whether we can come up with a more fancy name but currently we are settled on FSDPv2. I already talked about this in my laste [update](https://dev-discuss.pytorch.org/t/pytorch-xla-2-2-release-dev-update/1897#fsdp-via-spmd-fsdp-v2-3). FSDPv2 has become our standard way of training large language models. We integrated FSDPv2 with [HF recently](https://huggingface.co/blog/gemma-peft#accelerate-with-fsdp-via-spmd-on-tpu) and have achieved 60%+ MFUs on llama3 and Geema on Cloud TPU v5p. I think this should be the first thing users try when they want to scale up their model. ## Pallas [Pallas](https://jax.readthedocs.io/en/latest/pallas/index.html) is a custom kernel language that supports both TPU and GPU. In the 2.3 release Pytorch/XLA integrated with Pallas on TPU and provided support using Pallas based FlashAttention kernel for the model forward. We also added the support for the `torch.compile` with the flash attention ```auto from torch_xla.experimental.custom_kernel import flash_attention output = flash_attention(q, k, v) ``` In the nightly we also added the support for the backward along with a couple other popular kernels like `paged_attention` and `gmm`. With Pallas integration PyTorch/XLA will be able to quickly experiment with the most advanced features. PyTorch/XLA will keep adding the popular Pallas kernels and examples to the repo so users can directly use them. For more detail please take a look at this [user guide](https://github.com/pytorch/xla/blob/r2.3/docs/pallas.md). ## Export From the 2.3 release PyTorch/XLA will consider `export` to be a stable feature. In this release we added support for all of the core aten ops that’s used by the export. We also enabled the dynamism in the export and fixed a bunch of bugs. Please give this feature a try if you are interested in exporting a pytorch model to StableHLO. ## Torch Dynamo In the 2.3 release PyTorch/XLA added the custom op `dynamo_mark_sharding` which can be used to perform the activation sharding in a `torch.compile` region. This is the first step to make `torch.compile` + `GSPMD` to be the recommended way of doing the model inference using PyTorch/XLA. We plan to provide more examples and performance benchmarks in the 2.4 release. ## GPU PyTorch/XLA added the official support for running SPMD on either a single node GPU or multinode GPU. You can find the instructions in [this doc.](https://github.com/pytorch/xla/blob/master/docs/spmd.md#running-spmd-on-gpu) We were able to get ~50% MFU on Llama2 using a single GPU node with 8 A100 GPU. PyTorch/XLA:GPU is also now a plugin that you can install with ```auto pip torch~=2.3.0 torch_xla~=2.3.0 https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla_cuda_plugin-2.3.0-py3-none-any.whl ``` ## While loop A frequent criticism of PyTorch/XLA is the extended compilation duration. One of the root causes of this prolonged compilation process is the unrolling of all Python loops, leading to a substantial HLO. In the 2.3 release PyTorch/XLA added the experimental support for the `torch._higher_order_ops.while_loop`. Currently it can only handle very basic torch operations, but we aim to add the nn module support in the upcoming release. The goal is to have the `while_loop` to wrap around a module(like a DecoerLayer in Llama) and produce a compact HLO. For more detail please take a look at [this doc](https://github.com/pytorch/xla/blob/r2.3/docs/fori_loop.md).