sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

move bf16 cast hack to Tensor.llvm_bf16_cast (#3788)

Browse Source
This commit is contained in:
chenyu
2024-03-17 18:51:22 -04:00
committed by GitHub
parent 311cf2b7d3
commit 639bd5dbfc
5 changed files with 13 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/coder.py
+1 -1
View File
@@ -34,7 +34,7 @@ if __name__ == "__main__":
part2 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00002-of-00002.bin?download=true"))
# fix bf16, TODO: check if device supports bf16
def fix_bf16(weights): return {k:v.llvm().cast(dtypes.float16).to(Device.DEFAULT) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
def fix_bf16(weights): return {k:v.llvm_bf16_cast(dtypes.half).to(v.device) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
with Timing("weights -> model: "):
nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, model, 32, 8)), strict=False)