LLM/airllm
1
0
Fork 0
mirror of https://github.com/0xSojalSec/airllm.git synced 2026-03-07 22:33:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

make sure its on cuda

Browse Source
This commit is contained in:
Yu Li
2023-12-01 18:43:52 -06:00
parent 9a5d4dc730
commit d3511affe0
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
air_llm/airllm/airllm.py
View File

@@ -98,7 +98,11 @@ def uncompress_layer_state_dict(layer_state_dict):
absmax = layer_state_dict[k + ".8bit.absmax"]
code = layer_state_dict[k + ".8bit.code"]
dqv = bnb.functional.dequantize_blockwise(v.cuda(), bnb.functional.QuantState(absmax=absmax, code=code, blocksize=2048, dtype=torch.float16))
dqv = bnb.functional.dequantize_blockwise(v.cuda(),
bnb.functional.QuantState(absmax=absmax.cuda(),
code=code.cuda(),
blocksize=2048,
dtype=torch.float16))
uncompressed_layer_state_dict[k] = dqv
del layer_state_dict