remove metal in memory APIs (#50)

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)

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",
         )
 

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))
 
 

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}, "

requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.23.3
+mlx>=0.24.1
 numpy
 transformers[sentencepiece]>=4.39.3
 protobuf