# PyTorch/XLA 2.4 dev update

**URL:** <https://dev-discuss.pytorch.org/t/pytorch-xla-2-4-dev-update/2356>\
**Category:** compiler\
**Created:** [August 6, 2024, 6:05pm UTC](https://dev-discuss.pytorch.org/t/pytorch-xla-2-4-dev-update/2356 "2024-08-06T18:05:12Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**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:** [August 6, 2024, 6:05pm UTC](https://dev-discuss.pytorch.org/t/pytorch-xla-2-4-dev-update/2356/1 "2024-08-06T18:05:12Z")

</div>

# PyTorch/XLA 2.4 Dev Update

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.4.0) for detailed updates. I am going to highlight some of the new features and share how I think about them.

## Eager Mode

In 2.4 release we introduce the eager mode to the PyTorch/XLA, you can enable it via

```auto
torch_xla.experimental.eager_mode(True)

```

This might sound weird to you if you are not familiar with torch\_xla. The default mode of the torch\_xla is lazy tracing

```auto
a = torch.tensor(100, device=torch_xla.device())
# execution does not happen, we created a node called `COS` with an input
b = torch.cos(a)
c = torch.sin(B)

# execution happens here, we construct a graph to calculate the value of both b and c
torch_xla.sync()

```

The motivation of this approach was to accumulate as many pytorch ops as possible and execute them in one graph to optimize the performance. However over the year we found this mode to be too confusing so we decided to add the eager mode so it is closer to the native pytorch experience. The performance of the eager mode can be found in [here](https://pytorch.org/xla/release/r2.4/eager_mode.html#benchmark) but the TLDR is that it is very model dependent and ranges from 1% - 45% of the compiled performance.

Going forward we want users to mix eager mode with a compile API similar to the native pytorch. For inference it will looks like

```auto
torch_xla.experimental.eager_mode(True)

compiled_model = torch.compile(model, backend="openxla")
res = compiled_model(input)

```

For training we currently recommend

```auto
torch_xla.experimental.eager_mode(True)

def step_fn(model, data, target, loss_fn, optimizer):
    optimizer.zero_grad()
    logits = model(data)
    loss = loss_fn(logits, target)
    loss.backward()
    optimizer.step()
    return loss

step_fn = torch_xla.experimental.compile(step_fn)

```

And hope we can also eventually converge on using the `torch.compile` directly. To learn more about the eager mode + compile, please take a look at [this doc](https://pytorch.org/xla/release/r2.4/eager_mode.html#benchmark).

## Pallas Kernel

In this release we really try to harden our pallas kernel support. [Pallas](https://jax.readthedocs.io/en/latest/pallas/index.html) is a custom kernel language that supports both TPU and GPU. Pallas was born in JAX to support the need of implementing fusions without the XLA compiler and better control the memory access pattern. It got more popular on the TPU side and now there are a couple very popular custom kernels being open sourced in [here](https://github.com/google/jax/tree/main/jax/experimental/pallas/ops/tpu).

PyTorch/XLA has made a lot of progress on porting the existing Pallas kernel like `flash_attention`, `paged_attention`, `gmm `(megablocks for moe), `tgmm`, into the PyTorch/XLA repo. Users can also call their own Pallas kernel as we have talked about in the last dev update. In order to use these kernels on training, we also did additional work to register the corresponding backward kernel, adding the support for GSPMD (GSPMD sharding propagation is disabled in Palls kernel since they are just black boxes to the compiler). You can take a look at [this example](https://github.com/pytorch/xla/blob/r2.4/examples/flash_attention/train_decoder_only_flash_attention.py) for how to use our flash attention wrapper.

PyTorch/XLA team sees a rapid need for the custom kernel and our goal is to make this experience easier. For more information please take a look at [this doc](https://github.com/pytorch/xla/blob/r2.4/docs/pallas.md).

## Triton kernel

In this release we also added the support for triton kernel on XLA:GPU. You can use it via

```auto
payload = xla_triton.triton_call(
    x, y, output, size, kernel=add_kernel, grid=grid, BLOCK_SIZE=block_size)

output = torch_xla._XLAC._xla_gpu_custom_call([x, y], payload,
                                                [output.shape], [torch.int64])

```

The triton kernel will be treated as a black box and represented as a custom call in the HLO graph generated. For more detail you can take a look at [this doc](https://github.com/pytorch/xla/blob/r2.4/docs/triton.md)

## Usability & Debuggability Improvements

In the past few releases PyTorch/XLA has come a long way to optimize the performance but the team is aware that usability still has many issues. On top of the eager mode change mentioned above, we also made a couple changes to clean up our api and add more debugging tools. For example

1. `torch_xla.sync()` to replace the `torch_xla.core.xla_model.mark_step()`
2. `torch_xla.device()` to replace the `torch_xla.core.xla_model.xla_device()`
3. add the support for `torch_xla.core.xla_model.get_memory_info` to query the TPU HBM info
4. add compiled program HBM usage analysis when running with `PT_XLA_DEBUG=1`
5. add `PT_XLA_DEBUG_LEVEL=1` to only output debug messages on new compilation.
6. refresh our [doc page](https://pytorch.org/xla/release/r2.4/index.html)

We plan to keep cleaning up our API and make it more aligned with the upstream while improving the debugging experience.