PyTorch/XLA 1.13 release

Cloud TPUs now support the PyTorch 1.13 release, via PyTorch/XLA integration. The release has daily automated testing for the supported models: Torchvision ResNet, FairSeq Transformer and RoBERTa, HuggingFace GLUE and LM, and Facebook Research DLRM.

On top of the underlying improvements and bug fixes in PyTorch's 1.13 release, this release adds several features and PyTorch/XLA specified bug fixes.

New Features

• GPU enhancement
  • Add upsample_nearest/bilinear implementation for CPU and GPU (#3990)
  • Set three_fry as the default RNG for GPU (#3951)
• FSDP enhancement
  • allow FSDP wrapping and sharding over modules on CPU devices (#3992)
  • Support param sharding dim and pinning memory (#3830)
• Lower torch::einsum using xla::einsum which provide significant speedup (#3843)
• Support large models with >3200 graph input on TPU + PJRT (#3920)

Experimental Features

• PJRT experimental support on Cloud TPU v4
  • Check the instruction and example code in here
• DDP experimental support on Cloud TPU and GPU
  • Check the instruction, analysis and example code in here

Ongoing development

• Ongoing Dynamic Shape implementation (POC completed)
• Ongoing SPMD implementation (POC completed)
• Ongoing LTC migration

Bug fixes and improvements

• Make XLA_HLO_DEBUG populate the scope metadata (#3985)

PyTorch/XLA 1.12 release

Cloud TPUs now support the PyTorch 1.12 release, via PyTorch/XLA integration. The release has daily automated testing for the supported models: Torchvision ResNet, FairSeq Transformer and RoBERTa, HuggingFace GLUE and LM, and Facebook Research DLRM.

On top of the underlying improvements and bug fixes in PyTorch's 1.12 release, this release adds several features and PyTorch/XLA specified bug fixes.

New feature

• FSDP
  • Check the instruction and example code in here
  • FSDP support for PyTorch/XLA (#3431)
  • Bfloat 16 and float 16 support in FSDP (#3617)
• PyTorch/XLA gradident checkpoint api (#3524)
• Optimization_barrier which enables gradient checkpointing (#3482)
• Ongoing LTC migration
• Device lock position optimization to speed up tracing (#3457)
• Experimental support for PJRT TPU client (#3550)
• Send/Recv CC op support (#3494)
• Performance profiling tool enhancement (#3498)
• TPU-V4 pod official support (#3440)
• Roll lowering (#3505)
• Celu, celu_, selu, selu_ lowering (#3547)

Bug fixes and improvements

• Fixed a view bug which will create unnecessary IR graph (#3411)

PyTorch/XLA 1.11 release

Cloud TPUs now support the PyTorch 1.11 release, via PyTorch/XLA integration. The release has daily automated testing for the supported models: Torchvision ResNet, FairSeq Transformer and RoBERTa, HuggingFace GLUE and LM, and Facebook Research DLRM.

On top of the underlying improvements and bug fixes in PyTorch's 1.11 release, this release adds several features and PyTorch/XLA specified bug fixes.

New feature

• Enable asynchronous RNG seed sending by environment variable XLA_TRANSFER_SEED_ASYNC
• Add a native torch.distributed backend
• Introduce a Eager debug mode by environment variable XLA_USE_EAGER_DEBUG_MODE
• Add synchronous free Adam and AdamW optimizers for PyTorch/XLA:GPU AMP
• Add synchronous free SGD optimizers for PyTorch/XLA:GPU AMP
• linspace lowering
• mish lowering
• prelu lowering
• slogdet lowering
• stable sort lowering
• index_add with alpha scaling lowering

Bug fixes && improvements

• Improve torch.var performance and numerical stability on TPU
• Improve torch.pow performance
• Fix the incorrect output dtype when divide a f32 by a f64
• Fix the incorrect result of nll_loss when reduction = "mean" and whole target is equal to ignore_index

PyTorch/XLA 1.10 release

Cloud TPUs now support the PyTorch 1.10 release, via PyTorch/XLA integration. The release has daily automated testing for the supported models: Torchvision ResNet, FairSeq Transformer and RoBERTa, HuggingFace GLUE and LM, and Facebook Research DLRM.

On top of the underlying improvements and bug fixes in PyTorch's 1.10 release, this release adds several PyTorch/XLA-specific bug fixes:

• Add support for reduce_scatter
• Introduce the AMP Zero gradients optimization for XLA:GPU
• Introduce the environment variable XLA_DOWN_CAST_BF16 and XLA_DOWNCAST_FP16 to downcast input tensors
• adaptive_max_pool2d lowering
• nan_to_num lowering
• sgn lowering
• logical_not/logical_xor/logical_or/logical_and lowering
• amax lowering
• amin lowering
• std_mean lowering
• var_mean lowering
• lerp lowering
• isnan lowering

PyTorch/XLA 1.8 release

Summary

Cloud TPUs now support the PyTorch 1.8 release, via PyTorch/XLA integration. The release has daily automated testing for the supported models: Torchvision ResNet, FairSeq Transformer and RoBERTa, HuggingFace GLUE and LM, and Facebook Research DLRM.

This release focused on making PyTorch XLA easier to use and debug. See below for a list of new features.

New Features

• Enhanced usability:
  • Profiler tools to help you pinpoint the areas where you can improve the memory usage or speed of your TPU models. The tools are ready to use; check out our main README for some upcoming tutorials.
  • Simpler error messages (#2771)
  • Less log spam using TPU Pods (#2662)
  • Able to view images in Tensorboard (#2679)
• TriangularSolve (#2498) (example)
• New ops supported by PyTorch/XLA:
  • random_ (#2617)
  • adaptive_avg_pool3d (#2616)
  • UpsampleNearest2D (#2597)

Bug Fixes

• Crashing while using dynamic shapes (#2602)
• all_to_all crashing on TPU pods (#2601)
• SiLU fix (#2721)

PyTorch/XLA 1.6 Release (GA)

Highlights

Cloud TPUs now support the PyTorch 1.6 release, via PyTorch/XLA integration. With this release we mark our general availability (GA) with the models such as ResNet, FairSeq Transformer and RoBERTa, and HuggingFace GLUE task models that have been rigorously tested and optimized.

In addition, with our PyTorch/XLA 1.6 release, you no longer need to run the env-setup.py script on Colab/Kaggle as those are now compatible with native torch wheels. See here for an example of the new Colab/Kaggle install step. You can still continue to use that script if you would like to run with our latest unstable releases.

New Features

• XLA RNG state checkpointing/loading (#2096)
• Device Memory XRT API (#2295)
• [Kaggle/Colab] Small host VM memory environment utility (#2025)
• [Advanced User] XLA Builder Support (#2125)
• New ops supported on PyTorch/XLA
  • Hardsigmoid (#1940)
  • true_divide (#1782)
  • max_unpool2d (#2188)
  • max_unpool3d (#2188)
  • Replication_pad1d (#2188)
  • Replication_pad2d (#2188)
• Dynamic shape support on XLA:CPU and XLA:GPU (experimental)

Bug Fixes

• RNG Fix (proper randomness with bernoulli and dropout) (#1932)
• Manual all-reduce in backward pass (#2325)

PyTorch/XLA 1.5 release

Cloud TPUs and Cloud TPU Pods now support PyTorch 1.5 via the PyTorch/XLA integration. This integration aims to make it possible for PyTorch users to do everything they can do on GPUs on Cloud TPUs as well while minimizing changes to the user experience. You can try out PyTorch on an 8-core Cloud TPU device for free via Google Colab, and you can use PyTorch on Cloud TPUs at a much larger scale on Google Cloud (all the way up to full Cloud TPU Pods).

Three PyTorch models have been added to our list of supported models, which are rigorously and continuously tested:

• ResNet-50
• Fairseq Transformer
• Fairseq RoBERTa

Additional notes:

• New Operators added
  • masked_scatter (#1525)
  • logdet (#1407)
  • reflection_pad2d (#1409)
  • gelu(#1278)
  • take(#1389)
  • upsample_bilinear2d (#1239)
  • upsample_nearest2d (#1233)
  • mse_loss(#1206)
• Exposed APIs to enable different types of cross-replica reduce operations using the TPU interconnect link (#1709)
• Exposed API to perform rendezvous operations among the different replica processes (#1669)
• Added support for reading/writing GCS files (#1230)
• Added support to read TFRecords (#1220)
• Miscellaneous bug fixes

PyTorch/XLA 1.9 release

Cloud TPUs now support the PyTorch 1.9 release, via PyTorch/XLA integration. The release has daily automated testing for the supported models: Torchvision ResNet, FairSeq Transformer and RoBERTa, HuggingFace GLUE and LM, and Facebook Research DLRM.

On top of the underlying improvements and bug fixes in PyTorch's 1.9 release, this release adds several PyTorch/XLA-specific bug fixes:

• Floor division fix
• Clamp fix
• Softplus fix
• as_strided_
• std abd var lowering

PyTorch/XLA 1.8.1 release

Cloud TPUs now support the PyTorch 1.8.1 release, via PyTorch/XLA integration. The release has daily automated testing for the supported models: Torchvision ResNet, FairSeq Transformer and RoBERTa, HuggingFace GLUE and LM, and Facebook Research DLRM.

On top of the underlying bug fixes in PyTorch's 1.8.1 release, this release adds a few bug fixes on the PyTorch XLA side around the XRT server and TPU Pods training.

PyTorch/XLA 1.7 release

Summary

Cloud TPUs now support the PyTorch 1.7 release, via PyTorch/XLA integration. The release has daily automated testing for the supported models: Torchvision ResNet, FairSeq Transformer and RoBERTa, HuggingFace GLUE and LM, and Facebook Research DLRM.

New Features

• TriangularSolve (#2498) (example)
• New ops supported by PyTorch/XLA:
  • round (#2457)
  • baddbmm_out (#2451)
  • var_out (#2495)
  • asinh, acosh, atanh (#2501)
• Documentation on adding more supported ops: (#2458)

Bug Fixes

• exponential_() returning 0 (#2562)
• cross_entropy on inf input (#2553)