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