Add layer-streaming to server, convert, and verification

- --streaming / --max-memory-gb / --streaming-window on mlx_lm.server
- --split-for-streaming on mlx_lm.convert
- Forward and sampling mx.array tests for streaming path
- scripts/verify_streaming_consumer.py for end-to-end MLX tensor checks

Author: SashimiSaketoro

mlx_lm/convert.py

@@ -97,6 +97,7 @@ def convert(
         Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str]
     ] = None,
     trust_remote_code: bool = False,
+    split_for_streaming: bool = False,
 ):
     # Check the save path is empty
     if isinstance(mlx_path, str):
@@ -172,6 +173,12 @@ def set_dtype(k, v):
         config,
     )
 
+    if split_for_streaming:
+        from mlx_lm.streaming.split_model import ensure_streaming_layout
+
+        print("[INFO] Splitting weights for layer-streaming inference")
+        ensure_streaming_layout(mlx_path, verbose=True)
+
     if upload_repo is not None:
         upload_to_hub(mlx_path, upload_repo)
 
@@ -251,6 +258,12 @@ def configure_parser() -> argparse.ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--split-for-streaming",
+        help="Split saved weights into per-layer files for streaming inference.",
+        action="store_true",
+        default=False,
+    )
     return parser
 
 

mlx_lm/server.py

@@ -345,6 +345,23 @@ def _load(self, model_path, adapter_path=None, draft_model_path=None):
                 tokenizer_config=self._tokenizer_config,
                 trust_remote_code=self.cli_args.trust_remote_code,
             )
+        elif getattr(self.cli_args, "streaming", False):
+            if adapter_path is not None or draft_model_path is not None:
+                raise ValueError(
+                    "Layer streaming does not support adapters or draft models"
+                )
+            from mlx_lm.streaming import StreamingConfig, load_streaming
+
+            streaming_config = StreamingConfig(
+                max_memory_gb=self.cli_args.max_memory_gb,
+                window_size=getattr(self.cli_args, "streaming_window", None),
+                verbose=self.cli_args.log_level == "DEBUG",
+            )
+            model, tokenizer, _ = load_streaming(
+                model_path,
+                streaming_config=streaming_config,
+                tokenizer_config=self._tokenizer_config,
+            )
         else:
             model, tokenizer = load(
                 model_path,
@@ -1947,6 +1964,23 @@ def main():
         action="store_true",
         help="Use pipelining instead of tensor parallelism",
     )
+    parser.add_argument(
+        "--streaming",
+        action="store_true",
+        help="Enable layer-streaming inference for models larger than RAM.",
+    )
+    parser.add_argument(
+        "--max-memory-gb",
+        type=float,
+        default=20.0,
+        help="Memory budget (GB) for layer window when --streaming is set.",
+    )
+    parser.add_argument(
+        "--streaming-window",
+        type=int,
+        default=None,
+        help="Fixed layer window size (auto-computed if omitted).",
+    )
     args = parser.parse_args()
     if mx.metal.is_available():
         wired_limit = mx.device_info()["max_recommended_working_set_size"]

scripts/verify_streaming_consumer.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+# Copyright © 2024 Apple Inc.
+
+"""Consumer verification for mlx-lm layer-streaming (MLX-first tensor paths)."""
+
+import json
+import sys
+import tempfile
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+import mlx.core as mx
+
+from mlx_lm.generate import generate_step
+from mlx_lm.models.llama import Model, ModelArgs
+from mlx_lm.streaming import StreamingConfig, load_streaming
+from mlx_lm.streaming.split_model import ensure_streaming_layout, split_model_by_layers
+from mlx_lm.utils import save_model
+
+
+def _tiny_config():
+    return {
+        "model_type": "llama",
+        "vocab_size": 128,
+        "hidden_size": 64,
+        "intermediate_size": 128,
+        "num_hidden_layers": 2,
+        "num_attention_heads": 4,
+        "num_key_value_heads": 2,
+        "rms_norm_eps": 1e-5,
+        "tie_word_embeddings": True,
+    }
+
+
+def _build_model_dir(tmp_path: Path) -> Path:
+    config = _tiny_config()
+    with open(tmp_path / "config.json", "w") as f:
+        json.dump(config, f)
+    args = ModelArgs.from_dict(config)
+    model = Model(args)
+    mx.eval(model.parameters())
+    save_model(tmp_path, model)
+    ensure_streaming_layout(tmp_path, verbose=False)
+    return tmp_path
+
+
+def main() -> int:
+    with tempfile.TemporaryDirectory() as tmp:
+        model_dir = _build_model_dir(Path(tmp))
+
+        # Split I/O uses mx.load / mx.save_safetensors
+        weights = mx.load(str(model_dir / "layer_0.safetensors"))
+        assert all(isinstance(v, mx.array) for v in weights.values())
+        print("SPLIT OK", isinstance(weights["self_attn.q_proj.weight"], mx.array))
+
+        model, _, _ = load_streaming(
+            str(model_dir),
+            StreamingConfig(window_size=1, verbose=False),
+            load_tokenizer=False,
+        )
+        print("LOAD OK", type(model).__name__)
+
+        cache = model.make_cache()
+        logits = model(mx.array([[1, 2, 3]]), cache=cache)
+        mx.eval(logits)
+        assert isinstance(logits, mx.array)
+        print("FORWARD OK", isinstance(logits, mx.array), logits.shape)
+
+        prompt = mx.array([1, 2, 3])
+        for token, logprobs in generate_step(prompt, model, max_tokens=2):
+            assert isinstance(logprobs, mx.array)
+            assert isinstance(token, (int, mx.integer_types if hasattr(mx, "integer_types") else int))
+        print("SAMPLE OK", isinstance(logprobs, mx.array))
+
+        # Standalone split entry path
+        mono = tmp + "/mono.safetensors"
+        mx.save_safetensors(
+            mono,
+            {
+                "model.layers.0.self_attn.q_proj.weight": mx.random.normal((8, 8)),
+                "model.embed_tokens.weight": mx.random.normal((100, 8)),
+            },
+        )
+        out = Path(tmp) / "split2"
+        split_model_by_layers(Path(mono), out)
+        assert (out / "layer_0.safetensors").exists()
+        print("SPLIT_CLI OK", True)
+
+    print("ALL CONSUMER TESTS PASSED")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/test_streaming_generate.py

@@ -60,6 +60,26 @@ def test_load_streaming_auto_split():
         assert config["num_hidden_layers"] == 2
 
 
+def test_streaming_forward_returns_mx_array():
+    """StreamingModelWrapper forward must return mx.array logits."""
+    with tempfile.TemporaryDirectory() as tmp:
+        model_dir = _build_tiny_model_dir(Path(tmp))
+        model, _, _ = load_streaming(
+            str(model_dir),
+            StreamingConfig(window_size=1, verbose=False),
+            load_tokenizer=False,
+        )
+
+        cache = model.make_cache()
+        inputs = mx.array([[1, 2, 3]])
+        logits = model(inputs, cache=cache)
+        mx.eval(logits)
+
+        assert isinstance(logits, mx.array)
+        assert logits.ndim == 3
+        assert logits.shape[-1] == _tiny_llama_config()["vocab_size"]
+
+
 def test_generate_step_on_streaming_model():
     with tempfile.TemporaryDirectory() as tmp:
         model_dir = _build_tiny_model_dir(Path(tmp))
@@ -77,4 +97,21 @@ def test_generate_step_on_streaming_model():
             else:
                 tokens.append(int(token))
 
-        assert len(tokens) == 3
+        assert len(tokens) == 3
+
+
+def test_generate_step_sampling_returns_mx_array_logprobs():
+    """generate_step logprobs must be mx.array through the streaming path."""
+    with tempfile.TemporaryDirectory() as tmp:
+        model_dir = _build_tiny_model_dir(Path(tmp))
+        model, _, _ = load_streaming(
+            str(model_dir),
+            StreamingConfig(window_size=1, verbose=False),
+            load_tokenizer=False,
+        )
+
+        prompt = mx.array([1, 2, 3])
+        for _token, logprobs in generate_step(prompt, model, max_tokens=2):
+            assert isinstance(logprobs, mx.array)
+            assert logprobs.ndim == 1
+            assert logprobs.shape[0] == _tiny_llama_config()["vocab_size"]