PR #28949 [TPU] Support PyTorch/XLA FSDP via SPMD

[TPU] Support PyTorch/XLA FSDP via SPMD #28949

amyeroberts merged 10 commits into huggingface:main from alanwaketan:fsdpv2

ArthurZucker approved these changes on 2024-02-12

ArthurZucker requested a review from

muellerzr 1 year ago

muellerzr commented on 2024-02-12

alanwaketan force pushed 1 year ago

alanwaketan commented on 2024-02-14

muellerzr approved these changes on 2024-02-14

Initial commit

acb6e9f6

Add guards for the global mesh

a33c3cc9

Address more comments

a98ee8de

Move the dataloader into integrations/tpu.py

f13f8c56

Fix linters

ad1ee78a

Make karg more explicitly

36d091c8

Remove the move device logic

76ae042f

Fix the CI

b4386482

Fix linters

6a009a0a

Re-enable checkpointing

b4bdf0ae

alanwaketan force pushed to b4bdf0ae 1 year ago

amyeroberts merged 5f06053d into main 1 year ago

Reviewers

ArthurZucker

muellerzr

Assignees

No one assigned

Labels

None yet

Milestone

No milestone