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.

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

[FSDP2] allow meta tensors during loading state dict and cpu offloading #126267

Closed
wants to merge 2 commits into from

Commits on May 15, 2024

  1. [FSDP2] allow meta tensors during loading state dict and cpu offloading

    Summary:
    
    Test Plan:
    
    Reviewers:
    
    Subscribers:
    
    Tasks:
    
    Tags:
    
    [ghstack-poisoned]
    weifengpy committed May 15, 2024
    Configuration menu
    Copy the full SHA
    12d6770 View commit details
    Browse the repository at this point in the history
  2. Update on "[FSDP2] allow meta tensors during loading state dict and c…

    …pu offloading"
    
    
    unit test: ``pytest test/distributed/_composable/fsdp/test_fully_shard_state_dict.py``
    
    with meta init and cpu offloading, we have meta tensors after`model.load_state_dict(assign=True, strict=False)`. This PR avoided calling `.cpu` on meta tensors otherwise it's a runtime error
    
    cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k
    
    [ghstack-poisoned]
    weifengpy committed May 15, 2024
    Configuration menu
    Copy the full SHA
    5eb2e7d View commit details
    Browse the repository at this point in the history