Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 10 commits into from
Feb 14, 2024
Merged

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

merged 10 commits into from
Feb 14, 2024

Commits on Feb 14, 2024

  1. Initial commit

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    acb6e9f View commit details
    Browse the repository at this point in the history

  2. Add guards for the global mesh

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    a33c3cc View commit details
    Browse the repository at this point in the history

  3. Address more comments

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    a98ee8d View commit details
    Browse the repository at this point in the history

  4. Move the dataloader into integrations/tpu.py

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    f13f8c5 View commit details
    Browse the repository at this point in the history

  5. Fix linters

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    ad1ee78 View commit details
    Browse the repository at this point in the history

  6. Make karg more explicitly

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    36d091c View commit details
    Browse the repository at this point in the history

  7. Remove the move device logic

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    76ae042 View commit details
    Browse the repository at this point in the history

  8. Fix the CI

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    b438648 View commit details
    Browse the repository at this point in the history

  9. Fix linters

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    6a009a0 View commit details
    Browse the repository at this point in the history

  10. Re-enable checkpointing

    @alanwaketan
    alanwaketan committed Feb 14, 2024
    Configuration menu
    Copy the full SHA
    b4bdf0a View commit details
    Browse the repository at this point in the history