rixty
/
KoboldAI-Client
mirror of https://github.com/KoboldAI/KoboldAI-Client.git
1
1
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #167 from VE-FORBRYDERNE/accelerate 

Fix base fairseq dense models when using accelerate with a GPU
This commit is contained in:
henk717 2022-07-02 02:19:41 +02:00 committed by GitHub
parent 5d957e33ae aeed9bd8f7
commit 74547b31d6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
aiserver.py
View File

@ -642,6 +642,11 @@ def move_model_to_devices(model):
import breakmodel
if(utils.HAS_ACCELERATE):
import accelerate.utils
for key, value in model.state_dict().items():
target_dtype = torch.float32 if breakmodel.primary_device == "cpu" else torch.float16
if(value.dtype is not target_dtype):
accelerate.utils.set_module_tensor_to_device(model, key, target_dtype)
disk_blocks = breakmodel.disk_blocks
gpu_blocks = breakmodel.gpu_blocks
ram_blocks = len(utils.layers_module_names) - sum(gpu_blocks)