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

[FSDPv2] Move the module to xla device #6525

Merged
merged 2 commits into from
Feb 13, 2024
Merged

[FSDPv2] Move the module to xla device #6525

merged 2 commits into from
Feb 13, 2024

Commits on Feb 13, 2024

  1. [FSDPv2] Move the module to xla device 

    Summary:
This change allows moving module to the xla device during wrapping such that
the caller doesn't need to move the module to the xla device.

Test Plan:
python test/spmd/test_fsdp_v2.py -v -k test_fsdp_v2_cpu_model
    @alanwaketan
    alanwaketan committed Feb 13, 2024
    Configuration menu
    Copy the full SHA
    9b506fc View commit details
    Browse the repository at this point in the history

  2. fix linters

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