Ronghang Hu

use XLA patched linear in FSDP (fix #3811 and #3718) and expose options on padding in all-gather and pinning memory

2

This PR applies a patch to `nn.Linear` (`torch.nn.functional.linear`) in XLA FSDP so that the `nn.Linear`'s backward pass will use its weight parameter rather than an intermediate result. This resolves the...

[RFC] A high-level GSPMD API in PT/XLA (based on `xs.mark_sharding`)

14

## 🚀 [RFC] A high-level GSPMD API in PT/XLA (based on `xs.mark_sharding`) This RFC proposes a high-level API for GSPMD through a wrapper class and a partitioning rule function, based...

Autograd discrepancy in `nn.Linear` (`torch.nn.functional.linear`) between native PyTorch and PyTorch/XLA

7

## 🐛 Bug There seems to be a discrepancy (in addition to https://github.com/pytorch/xla/issues/3718) in how `torch.nn.Linear` (`torch.nn.functional.linear`) is implemented and dispatched between the native PyTorch and PyTorch/XLA. In particular, **the...

## 🐛 Bug For the PJRT runtime, currently `xm.all_reduce` works well on v3-8 after https://github.com/pytorch/xla/pull/3704 is merged, but `xm.reduce_scatter` and `xm.all_to_all` still still do not work well, and the `xm.all_gather`...

allowing `xm.get_ordinal()` and default device in `xm.xla_device()` in PJRT

3

## 🚀 Feature **Context**: The new PRJT runtime provides many advantages over the XRT runtime and is a really great new feature in PyTorch/XLA. On the other hand, the current...

XLA profiler issues: 1) TPU device trace does not show trace annotations, and 2) PRJT only captures 2 cores

10

## 🐛 Bug The XLA profiler has been a crucial tool for performance analysis on TPUs. In https://github.com/pytorch/xla/pull/3866, the profiler was also integrated into the PJRT runtime. However, a few...