2.4 backport PR request list

[bhavya01]

The issue is to track 2.4 release backport.

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

Original PR link
Reason to backport
2.4 backport PR link

[ManfeiBai]

• original PR: [Fori_loop|While_loop] Enable while_loop/fori_loop, add test case #7157

• Reason to backport: The PR is related to a planned release feature urgent fix: urgent fix for while_loop in 2.4 release

• 2.4 backport PR link: [backport][Fori_loop|While_loop] Enable while_loop/fori_loop, add test case (#7157) #7306

[jonb377]

• Original PR: Use dest_offsets directly in LoadPlanner #7155
• Reason to backport: Fix a bug in distributed checkpointing with padded tensors
• 2.4 backport PR link: Use dest_offsets directly in LoadPlanner #7243

[bhavya01]

• original PR: Add build trigger for 2.4.0-rc1 release #7236
• Reason to backport: Include build dependencies for CUDA 12.4
• 2.4 backport PR link: Update cuda_deps.yaml to include CUDA 12.4 #7244

[bhavya01]

• original PR: Register custom call with CUDA plugin #7219
• Reason to backport: Minor addition to Triton functionality to support CUDA plugins
• 2.4 backport PR link: [backport] Register custom call with CUDA plugin #7303

[vanbasten23]

• original PR: Update the XLA pin #7254 Bump JAX version #7258
• Reason to backport: OpenXLA pin update requested because new telemetry in PyTorch/xla needs to be enabled, requested by Aman
• 2.4 backport PR link: Backport xla pin update (0612) #7261 Backport: Bump jaxversion corresponding to libtpu 0612 #7262

[will-cromar]

• Original PRs: Supporting passing ML framework name and version through the PjRt Client Create options in PyTorch/xla. #7249 and Enable PJRT plugins by default #7268
• Reason to backport: Enables new libtpu telemetry and detection of CUDA plugin package (torch_xla_cuda_plugin) by default.
• Backport link: Cherry-pick: PJRT plugin defaults (#7249 and #7268) #7270

[jeffhataws]

• Original PR: Enable bucketized all-reduce for gradients #7216
• Reason to backport: Parity with Neuron branch r2.1_aws_neuron
• Backport link: WIP

[JackCaoG]

• original PR: Avoid log spam #7278

• backport PR: Avoid log spam (#7278) #7284

Reason:
This is a regression. Without this fix, PyTorch/XLA will always output a log spam about PJRT version upon loading.

Risk:
Low.

[JackCaoG]

• origional pr: Revert the mul change #7271
• backport pr: Revert the mul change (#7271) #7285

Reason:
This is a regression fix. Without this fix a bunch of xlml test will fail

Risk:
Low since it is just a revert.

[ManfeiBai]

• original pr: Use np.prod instead of np.product #7301
• backport pr: [backport] Use np.prod instead of np.product (#7301) #7310

Reason:
It's a cherry-pick hot fix for AttributeError: module 'numpy' has no attribute 'product' in r2.4 CI

Risk:
Low since master issue has beem fixed by the original pr

[JackCaoG]

original prs

• Enable eager mode for PyTorch/XLA #7234
• add torch_xla.experimental.compile for eager mode #7246
• support eager mode with torch.compile #7256
• Egaerly execute inplace ops if in eager mode #7322
• Support eager mode for multi-process training #7327
• Handle random seed for eager mode #7328
• Enable eager spmd #7341
• fix eager mode spmd module loading with fsdpv2 #7631

backport prs

• backport Enable eager mode for PyTorch/XLA #7611
• Backport add torch_xla.experimental.compile for eager mode (#7246) #7643
• backport support eager mode with torch.compile #7256 #7649
• Cherry-pick Egaerly execute inplace ops if in eager mode #7322 #7666
• Backport Support eager mode for multi-process training #7327 #7668
• cherry-pick Handle random seed for eager mode (#7328) #7669
• Backport Enable eager spmd #7341 and fix eager mode spmd module loading with fsdpv2 #7631 #7673

Reason:
I want to enable the eager mode for 2.4 release since I want to make it default for the 2.6

Risk
low, all of the features are guarded by torch_xla.experimental.eager_mdoe, these prs should be no op is eager mode is disabled.

[bhavya01]

Original PR

• Update ansible config to separately update submodules when torch_pin is specified #7617
• Avoid OOM when build PyTorch with CUDA enabled. #7329

Backport PR

• Fetch submodules separately when torch pin is specified #7616
• Backport #7329: Avoid OOM when build PyTorch with CUDA enabled #7618

Reason:
To fix CI

Risk:
Low, since this doesn't change anything in torch_xla library.

[bhavya01]

Original PR:
#7640

Backport PR:
#7684

Reason:
To fix upstream pytorch build

Risk:
Low, this doesn't change anything in torch_xla library

[alanwaketan]

Original: #7231

Backport: #7708

Support MoE.

[JackCaoG]

Original PR

• Add docs for eager mode #7700
• Update eager.md #7710

Backport PR:

• Add docs for eager mode (#7700) #7717

Reason:
Back port docs for eager mode

Risk:
low, just a doc