SIGN IN SIGN UP

Create model with device='meta'

Summary:
See discussion: D54825007

Two optimizations:
1. Use `mmap=True` to load the checkpoint.

2. Create model with device="meta". Tensors created in this context do not carry data. Previously, llama7b model was created with fp32 (default), using up 25GB ram. With device="meta", tensors are assigned only when we load the state dict.

- Note: non-persistent buffers and tensors that do not have keys in the state dict will be created with device="meta" as well. These have to be manually initialized when creating the model. See D46784302.

Checkpoint loading time: 10s -> 0.011s

Peak memory usage: [37.8GB](https://lookaside.facebook.com/intern/diff/file/data/?number=1467921211&download=1) ->[25.5GB](https://lookaside.facebook.com/intern/diff/file/data/?number=1468357208&download=1)
Model creation time:[ 77s](https://lookaside.facebook.com/intern/diff/file/data/?number=1468360493&download=1) -> [11.6s](https://lookaside.facebook.com/intern/diff/file/data/?number=1468364061&download=1)

Follow on: iterate over params/buffers and initialize uninitialized tensors (instead of manually initializing, which is model-specific) T182328293

thanks iseeyuan for the tips:
https://pytorch.org/tutorials/recipes/recipes/module_load_state_dict_tips.html

bypass-github-export-checks

Reviewed By: iseeyuan

Differential Revision: D54871495

fbshipit-source-id: f6a8d01c88ce45bb5d2358cc522ba81c0e1fbd5b
L
Lucy Qiu committed
d5f898dbf163e4a5672cf58db971e8def9b3c489
Parent: d0512b6
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> on 3/13/2024, 11:40:32 PM