2.3 backport PR request list

[lsy323]

Similar to the backport request thread in 2.2 release #6036

The issue is to track 2.3 release backport.

For any PRs you want to backport to 2.3, please reply with following:

Original PR link
Reason to backport
2.3 backport PR link

[yeounoh]

Original PR link:

• Add IsVirtualDeivce() and refactor #6726
• [SPMD] Use UNKNOWN sharding type for implicit replication #6723
• [SPMD] auto-sharding PoC #6719
• Include auto-sharding configs as part of hash #6753

Reason to backport: new experimental auto-sharding feature

2.3 backport PR link:

• [r2.3] Backport #6726  #6733
• [r2.3] backport #6723 #6734
• [r2.3] backport: auto-sharding PoC (#6719) #6755
• [r2.3] Include auto-sharding configs as part of hash (#6753) #6763

[zpcore]

Original PR link:

• fix csv output error #6662

Reason to backport: Minor bug fixing for benchmark.

2.3 backport PR link:

• fix csv output error (#6662) #6738

[lsy323]

Original PR link:

• fix torchvision installation in tpu ci test setup script #6704
• install pinned torchvision for Github Action TPU CI #6730
• Bump JAX version in setup.py #6740
  Reason to backport: The PRs fix TPU CI at HEAD, this makes sure the CI infra on r2.3 branch is healthy

2.3 backport PR link:

• Backport #6704 to r2.3 #6736
• Backport #6730 #6739
• Backport Bump JAX version in setup.py (#6740) #6741

[lsy323]

Original PR link:

• Pin update March 2024 #6677
  Reason to backport: Update XLA pin, auto-sharding feature depends on the new XLA pin

2.3 backport PR link:

• Backport #6677 to r2.3 #6737

[alanwaketan]

Original PR link:

• [Pallas] Introduce make_kernel_from_pallas  #6713

Reason to backport: this is a new important feature for 2.3 which enables to use custom kernel like flash attention.

2.3 backport PR link:

• [Pallas] Introduce make_kernel_from_pallas (#6713) #6742

[wonjoolee95]

Original PR link:

• Properly lower add and mul #6731
• Properly lower aten::sub and aten::rsub #6668

Reason to backport: This is a series of important fixes to lowerings to reduce unnecessary decomposition.

2.3 backport PR link

• TBD

[lsy323]

The Above PRs goes into the 1st rc tag https://github.com/pytorch/xla/tree/v2.3.0-rc2

[bhavya01]

Original PR link:

• Properly lower div.scalar and div.tensor #6669

Reason to backport: Part of a series of important fixes to lowerings to reduce unnecessary decomposition.

2.3 backport PR link

• Properly lower div.scalar and div.tensor(#6669) #6808

[yeounoh]

Original PR link:

• Add IsVirtualDeivce() and refactor #6726
• [SPMD] Use UNKNOWN sharding type for implicit replication #6723
• [SPMD] auto-sharding PoC #6719
• Include auto-sharding configs as part of hash #6753

Reason to backport: new experimental auto-sharding feature

2.3 backport PR link:

• Backport #6726 to r2.3 #6733
• [r2.3] backport #6723 #6734
• [r2.3] backport: auto-sharding PoC (#6719) #6755
• [r2.3] Include auto-sharding configs as part of hash (#6753) #6763

The whole auto-sharding support will be available from the next RC tag. cc @lsy323

[JackCaoG]

Original PR:

• remove openxla_eval from the doc #6757

Reason to backport:
route user to use openxla dyanmo backend instead of openxla_eval

2.3 backport PR link:

• backport remove openxla_eval from the doc #6767

[jonb377]

Original PR:

• Address deprecation in torch.distributed.checkpoint #6773

Reason to backport:

• Address deprecation warnings when using distributed checkpointing

2.3 backport PR link:

• Address deprecation in torch.distributed.checkpoint #6786

[JackCaoG]

Original PR

• Fix preserve_rng_state for activation checkpointing #4690

Reason to backport:
AWS folks need this change.

backport PR

• Backport Fix preserve_rng_state for activation checkpointing (#4690) #6788

[alanwaketan]

Original PR

• [Pallas] Introduce jax_import_guard #6747

Reason to backport:
To support Pallas.

backport PR

• [Backport] Introduce jax_import_guard #6794

[JackCaoG]

Original PR

• ZeRO1: Add bucketting logic to control the size of tensors for all-gather/reduce-scatter #6025

Reason to backport:
help aws to bring this feature about zero1 optimization

Backport PR

• Backport ZeRO1: Add bucketting logic to control the size of tensors for all-gather/reduce-scatter (#6025) #6806

[ManfeiBai]

Original PR
- #6807

Reason to backport: to complex test case for feature fori_loop
will not backport due to only contain test case

[alanwaketan]

Original PR

• [Pallas] Add a bfloat16 flash attention test case #6748

Reason to backport:
To support Pallas.

backport PR

• [Pallas] Add a bfloat16 flash attention test case #6810

[yeounoh]

Original PR: #6797
Reason to backport: auto-sharding fix in XLA
Backport PR: #6811

[alanwaketan]

Original PR

• [Pallas] Make a FlashAttention Wrapper #6785

Reason to backport:
To support Pallas.

backport PR

• [Backport] Make a FlashAttention Wrapper #6827

[alanwaketan]

Original PR

• [Pallas] Allow setting FlashAttention's causal mask #6792

Reason to backport:
To support Pallas.

backport PR

• [Backport] Allow setting FlashAttention's causal mask #6837

[alanwaketan]

Original PR

• [Pallas] Support XLA_USE_BF16 #6817

Reason to backport:
To support Pallas.

backport PR

• [Backport] Support XLA_USE_BF16 #6841

[wonjoolee95]

Original PR

• Remove reshape lowering #6801

Reason to backport:
To fix a bug to remove unrelated warning message.

Backport PR

• [Backport] Remove reshape lowering (#6801) #6847

[yeounoh]

Original PR

• [SPMD] use UNKNOWN as default sharded data sharding type only for auto-sharding #6846

Reason to backport:

• Fix SPMD inference regression

backport PR

• [r2,3][SPMD] use UNKNOWN as default sharded data sharding type selectively #6852

[alanwaketan]

Original PR

• [Pallas] Set the major-minor layout for inputs and outputs #6826

Reason to backport:
To support Pallas.

backport PR

• [Backport] Set the major-minor layout for inputs and outputs #6876

[lsy323]

Original PR

• Fix div lowering and core aten test script enhancement #6873

Reason to backport:
Bug fix

Backport PR

• [r2.3] Fix div lowering and core aten test script enhancement (#6873) #6877

[JackCaoG]

Original PR

• Add flash attention as custom op so dynamo can trace it #6875

Reason to backport:
Make support for flash attention Pallas Kernel complete. If user can only use flash attention with LTC it won't be very useful for the inference.

Backport PR

• [Backport] Add flash attention as custom op so dynamo can trace it (#6875) #6884

[JackCaoG]

Original PR

• add torch_distributed tests to our TPU CI #6891

Reason to backport:
Without this change, we will always throw a warning message about dev_kind during mp training which is a bad ux.

Backport pr

• Fix warning message when running mp workload #6892

[ManfeiBai]

Original PR:

• [Fori_loop|While_loop] Create fori_loop.md #6942

Reason to backport: add doc for while_loop feature

Backport PR:

• [cherry-pick] [Fori_loop|While_loop] Create fori_loop.md (#6942) #6945

[lsy323]

Original PR:

• Build torchvision from the pytorch pin in ansible #6756
• Reland Install torchvision in release docker #6795 #6860
  Reason to backport: Fix torchvision version issue in release docker

Backport PR:

• Reland Install torchvision in release docker #6795 (#6860) #6955
• Build torchvision from the pytorch pin in ansible (#6756) #6960