zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-10-23 22:18:06 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

use lower precision base weights

Browse Source
This commit is contained in:
Awni Hannun
2023-12-15 10:29:42 -08:00
parent d108c558fc
commit 84f02ef58b
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
lora/lora.py
View File

@@ -184,6 +184,7 @@ def load(args):
def loss(model, inputs, targets, lengths):
# Run model on inputs
logits, _ = model(inputs)
logits = logits.astype(mx.float32)
# Mask padding tokens
length_mask = mx.arange(inputs.shape[1])[None, :] < lengths[:, None]
@@ -326,7 +327,7 @@ def generate(model, prompt, tokenizer, args):
print(s, flush=True)
def load_model(folder: str, dtype=mx.float32):
def load_model(folder: str, dtype=mx.float16):
model_path = Path(folder)
tokenizer = Tokenizer(str(model_path / "tokenizer.model"))
with open(model_path / "params.json", "r") as f: