zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-01 04:14:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Faster sampling with mx.compile (#937)

Browse Source 
* faster sampling with compile

* fix test
This commit is contained in:
Awni Hannun
2024-08-15 11:29:09 -07:00
committed by GitHub
parent 95840f32e2
commit 9b83004631
3 changed files with 32 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
llms/mlx_lm/sample_utils.py
View File

@@ -1,6 +1,11 @@
# Copyright © 2023-2024 Apple Inc.
from functools import partial
import mlx.core as mx
@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.array:
"""
Apply top-p (nucleus) sampling to logits.
@@ -13,7 +18,7 @@ def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.arr
token selected based on the top-p criterion.
"""
# referenced implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L449-L460
probs = mx.softmax(logits / temperature, axis=-1)
probs = mx.softmax(logits * (1 / temperature), axis=-1)
# sort probs in ascending order
sorted_indices = mx.argsort(probs, axis=-1)
@@ -25,10 +30,15 @@ def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.arr
top_probs = mx.where(
cumulative_probs > 1 - top_p,
sorted_probs,
mx.zeros_like(sorted_probs),
0,
)
sorted_token = mx.random.categorical(mx.log(top_probs))
token = sorted_indices.squeeze(0)[sorted_token]
return token
@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
def categorical_sampling(logits, temp):
return mx.random.categorical(logits * (1 / temp))

4
llms/mlx_lm/utils.py
View File

@@ -20,7 +20,7 @@ from transformers import PreTrainedTokenizer
# Local imports
from .models.base import KVCache
from .sample_utils import top_p_sampling
from .sample_utils import categorical_sampling, top_p_sampling
from .tokenizer_utils import TokenizerWrapper, load_tokenizer
from .tuner.utils import apply_lora_layers
from .tuner.utils import dequantize as dequantize_model
@@ -169,7 +169,7 @@ def generate_step(
if top_p > 0 and top_p < 1.0:
token = top_p_sampling(logits, top_p, temp)
else:
token = mx.random.categorical(logits * (1 / temp))
token = categorical_sampling(logits, temp)
return token, logprobs