PR #7917 Add torch_xla TPU support for ZeRO-1/2

Add torch_xla TPU support for ZeRO-1/2 #7917

PKUWZP wants to merge 9 commits into master from codex/tpu-xla-zero12

Add torch_xla TPU support for ZeRO-1/2

8acc916e

PKUWZP requested a review from

tjruwase 22 days ago

PKUWZP requested a review from

tohtana 22 days ago

PKUWZP requested a review from

GuanhuaWang 22 days ago

PKUWZP requested a review from

loadams 22 days ago

PKUWZP removed review request from

GuanhuaWang 22 days ago

PKUWZP requested a review from

delock 22 days ago

chatgpt-codex-connector commented on 2026-03-21

Fix TPU device selection for XLA workers

7f82c20d

Fix formatting hooks for XLA backend test

dbfd0a91

Fix XLA backend unit test import

41d0f7e9

Patch fake torch_xla before import in test

7debfbbb

Import XLA test target module explicitly

623df9ea

Relax XLA backend test initialization assertion

65cf60fd

tohtana commented on 2026-03-23

sfc-gh-truwase commented on 2026-03-24

Address review feedback for XLA/TPU support

7f6ea955

Address tohtana's review feedback for XLA/TPU support

b14e5b46

Reviewers

tohtana

sfc-gh-truwase

chatgpt-codex-connector

tjruwase

loadams

delock

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

DeepSpeed Add torch_xla TPU support for ZeRO-1/2 #7917 Open

Add torch_xla TPU support for ZeRO-1/2 #7917

DeepSpeed
Add torch_xla TPU support for ZeRO-1/2
#7917

Open