2.2 backport PR request list

[zpcore]

The issue is to track 2.2 release backport.

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

• Original PR link
• Reason to backport
• 2.2 backport PR link

[JackCaoG]

Master PR: #5956
Status: Merged

Backport PR: #6055
Status: Merged

Reason:
AWS need this to pr to land to enable the coalesce reduce_scatter

[JackCaoG]

Master PR: #5958
Status: Merged

Backport PR: #6037
Status: Merged

Reason:
Enhance the PT_XLA_DEBUG tool to truncate the python frame count.

[JackCaoG]

Master PR: #6039
status: Merged

Backport PR: #6095
Status: Merged

Reason:
Update the trouble shooting doc with pt_xla_debug.

[jonb377]

Master PR: #5803
Status: Merged

Backport PR: #6050
Status: Merged

Reason:
Persistent compilation caching - env hash

[jonb377]

Master PR: #5804
Status: Merged

Backport PR: #6065
Status: Merged

Reason:
Persistent compilation caching - enable via environment

[jonb377]

Master PR: #6046
Status: Merged

Backport PR: #6074
Status: Merged

Reason:
Persistent compilation caching - python API and docs

[jeffhataws]

Master PR: #5931
Status: Merged

Backport PR: #6054
Status: Merged

Reason:
AWS: Add a mesh_spmd for application to pass in topology, needed for pipeline parallel.

[jeffhataws]

Master PR: #5950
Status: Merged

Backport PR: 1271964 (got in before the branch cut)
Status: Merged

Reason:
AWS: enable all-gather coalescing

[JackCaoG]

Master PR: #5922
status: merged

Backport PR: #6052
Status: merged

Reason:
fix some dynamo model failure

[will-cromar]

Master PR: #6060
status: merged

Backport PR: #6066
status: merged

Reason: Fix transfers from eager GPU to XLA device

[ManfeiBai]

Master PR: #6053
status: merged

Backport PR: #6079
status: merged

Reason: Hot-fix random duplicated trace registered TPU CI failure

[jeffhataws]

Master PR: #6059
status: Merged

Backport PR: #6108
status: Merged

Reason: Add a separate out-of-place all-gather. Add missing test for all-gather coalesce out, and fixed a bug for list size check in output!=None case.

[jeffhataws]

Master PR: #6058
status: Merged

Backport PR: #6109
Status: Merged

Reason: Add a separate out-of-place reduce-scatter and accompanying test.

[jonb377]

Master PR: #6075
Status: Merged

Backport PR: #6091
Status: Merged

Reason: Eliminate flakiness in profiler tests

[JackCaoG]

Master PR: #6071 (comment)
Status: Merged

Backport PR: #6094
Status: Merged

Reason: Fix transfers from eager GPU to XLA device for scalar tensor

[JackCaoG]

Master PR: #6123
Status: Merged

Backport PR: #6129
Status: Merged

Reason: Bug fix for diagonal_scatter op.

[jonb377]

Master PR: #6148
Status: Merged

Backport PR: #6167
Status: Merged

Reason: Bug fix for persistent cache

[jonb377]

Master PR: #6116
Status: Merged

Backport PR: #6120
Status: Merged

Reason: Remove flaky test from CI

[jeffhataws]

Master PR: #5936
Status: Merged

Backport PR: #6237
Status: Merged

Reason: Fix a bug how zero-1 optimizer infer the local ranks with PP and DP.

[alanwaketan]

Master PR: #6097
Status: Merged

Backport PR: #6157
Status: Merged

Reason: Needed for v5e GA.

[alanwaketan]

Master PR: #6101
Status: Merged

Backport PR: #6158
Status: Merged

Reason: Needed for FSDPv2.

[qihqi]

Master PR: #6111
Status: Merged

Backport PR: #6164
Status: Merged

Reason: upstream changed torch.export;

[qihqi]

Master Pr: #6103
Status: Merged

Backport PR: #6172
Status: Merged

Reason: bug fix: currently computes totally wrong numbers for pow of int tensor and float exponent.

[wonjoolee95]

Master PR: #6110
Status: Merged

Backport PR: #6173
Status: Merged

Reason: Addresses behavior gap -- makes torch_xla's isnan op behavior consistent with torch's isnan

[ManfeiBai]

Master PR: #6166
Status: Merged

Backport PR: #6329
Status: Merged

Reason: promote int to float for tanh operation (consistent with Pytorch)

[jonb377]

Master PR: #6168
Status: Merged

Backport PR: #6204
Status: Merged

Reason: Add default XLA flags for perf on v5

[qihqi]

Master PR: #6056
Status: Merged

Backport PR: #6185
Status: Merged

Reason: this one removes excesive log message.

[alanwaketan]

Master PR: #6122
Status: Merged

Backport PR: #6187
Status: Merged

Reason: This one introduces fsdpv2.

[zpcore]

Master PR: #6221
Status: Merged

Backport PR: #6228
Status: Merged

Reason: Fix the xl-ml test expecttest issue.

[ManfeiBai]

Master PR: #6224
Status: Merged

Backport PR: #6238
Status: Merged

Reason: Promote int to float for acos operation (consistent with Pytorch

[jeffhataws]

Master PR: #6247, #6268
Status: Merged and Merged

Backport PR: #6278
Status: Merged

Reason: Bug fix for segfault when running test_zero1 on GPU (#6260)

(would you mind help to verify the modification of this comment? @jeffhataws, thanks!)

[ManfeiBai]

Mater PR: #6300
Status: Merged

Backport PR: #6304
Status: Merged

Reason: Correct order of constant vs args

[ManfeiBai]

Master PR: #6263
Status: Merged

Backport PR: #6329
Status: Merged

Reason: Fix #6252 for local test

[yeounoh]

Master PR: #6326
Status: merged

Backport PR: #6330
Status: merged

Reason: Fix #6319 (regression)

[vanbasten23]

Master PR: #6321
Status: merged

Backport PR: #6334
Status: merged

Reason: Fix regression (see #6320 for detail)

[jonb377]

Master PR: #6356
Status: merged

Backport PR: #6378
Status: merged

Reason: Update documentation for CheckpointManager

[alanwaketan]

Master PR: #6386
Status: merged

Backport PR: #6408
Status: merged

Reason: Update documentation for FSDPv2.