diff --git a/mlx_lm/evaluate.py b/mlx_lm/evaluate.py index 13cca54..8f103bd 100644 --- a/mlx_lm/evaluate.py +++ b/mlx_lm/evaluate.py @@ -112,7 +112,7 @@ class MLXLM(LM): ) mx.eval(score, ig) - mx.metal.clear_cache() + mx.clear_cache() is_greedy.append(ig) scores.append(score) diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index e90a742..0681636 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -211,7 +211,7 @@ def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None): "MB. This can be slow. See the documentation for possible work-arounds: " "https://github.com/ml-explore/mlx-lm/tree/main#large-models" ) - old_limit = mx.metal.set_wired_limit(max_rec_size) + old_limit = mx.set_wired_limit(max_rec_size) try: yield None finally: @@ -220,7 +220,7 @@ def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None): mx.synchronize(s) else: mx.synchronize() - mx.metal.set_wired_limit(old_limit) + mx.set_wired_limit(old_limit) @dataclass @@ -362,7 +362,7 @@ def generate_step( prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens) prompt_processed_tokens += prefill_step_size y = y[prefill_step_size:] - mx.metal.clear_cache() + mx.clear_cache() y, logprobs = _step(y) @@ -379,7 +379,7 @@ def generate_step( break yield y.item(), logprobs if n % 256 == 0: - mx.metal.clear_cache() + mx.clear_cache() y, logprobs = next_y, next_logprobs n += 1 @@ -490,7 +490,7 @@ def speculative_generate_step( quantize_cache_fn(cache) mx.eval([c.state for c in cache]) y = y[prefill_step_size:] - mx.metal.clear_cache() + mx.clear_cache() return y def _rewind_cache(num_draft, num_accept): @@ -634,7 +634,7 @@ def stream_generate( prompt_tps=prompt_tps, generation_tokens=n + 1, generation_tps=(n + 1) / (time.perf_counter() - tic), - peak_memory=mx.metal.get_peak_memory() / 1e9, + peak_memory=mx.get_peak_memory() / 1e9, finish_reason=None, ) @@ -648,7 +648,7 @@ def stream_generate( prompt_tps=prompt_tps, generation_tokens=n + 1, generation_tps=(n + 1) / (time.perf_counter() - tic), - peak_memory=mx.metal.get_peak_memory() / 1e9, + peak_memory=mx.get_peak_memory() / 1e9, finish_reason="stop" if token in tokenizer.eos_token_ids else "length", ) diff --git a/mlx_lm/server.py b/mlx_lm/server.py index c638891..5a009af 100644 --- a/mlx_lm/server.py +++ b/mlx_lm/server.py @@ -770,13 +770,6 @@ def main(): choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Set the logging level (default: INFO)", ) - parser.add_argument( - "--cache-limit-gb", - type=int, - default=None, - help="Set the MLX cache limit in GB", - required=False, - ) parser.add_argument( "--chat-template", type=str, @@ -795,11 +788,6 @@ def main(): level=getattr(logging, args.log_level.upper(), None), format="%(asctime)s - %(levelname)s - %(message)s", ) - - if args.cache_limit_gb is not None: - logging.debug(f"Setting cache limit to {args.cache_limit_gb} GB") - mx.metal.set_cache_limit(args.cache_limit_gb * 1024 * 1024 * 1024) - run(args.host, args.port, ModelProvider(args)) diff --git a/mlx_lm/tuner/trainer.py b/mlx_lm/tuner/trainer.py index 64e26af..b0d41d6 100644 --- a/mlx_lm/tuner/trainer.py +++ b/mlx_lm/tuner/trainer.py @@ -295,7 +295,7 @@ def train( it_sec = args.steps_per_report / train_time tokens_sec = float(n_tokens) / train_time trained_tokens += n_tokens - peak_mem = mx.metal.get_peak_memory() / 1e9 + peak_mem = mx.get_peak_memory() / 1e9 if rank == 0: print( f"Iter {it}: Train loss {train_loss:.3f}, " diff --git a/requirements.txt b/requirements.txt index d176161..2e57ac7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -mlx>=0.23.3 +mlx>=0.24.1 numpy transformers[sentencepiece]>=4.39.3 protobuf