| """Reader for artifacts/mlx-nvfp4/ produced by scripts/convert.py. |
| |
| Tensor names are the artifact's own. Every layer's 128 experts arrive stacked |
| and gate/up fused, exactly as the model consumes them, so loading copies nothing: |
| |
| model.decoder.layers.<i>.experts.gate_up_proj.{weight,scales,global_scale} |
| model.decoder.layers.<i>.experts.down_proj.{weight,scales,global_scale} |
| |
| `weight` is uint32 (8 NVFP4 values per word), `scales` is one E4M3 byte per 16 |
| input values, and `global_scale` is the checkpoint's per-expert float32 scale |
| (its amax, 448*6*weight_scale_2) which `mx.gather_qmm` takes as an argument. |
| Everything else is BF16 under its checkpoint name. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from collections.abc import Iterator, Mapping |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| LAYOUT = "stacked-global-scale" |
| DEFAULT_DIR = Path(__file__).resolve().parent.parent / "artifacts" / "mlx-nvfp4" |
| if not DEFAULT_DIR.is_dir(): |
| |
| DEFAULT_DIR = Path(__file__).resolve().parent.parent |
|
|
|
|
| def load_manifest(directory: str | Path = DEFAULT_DIR) -> dict: |
| """Quantisation mode, group size, shard list and tensor -> shard map.""" |
| manifest = json.loads((Path(directory) / "manifest.json").read_text()) |
| layout = manifest["quantization"].get("layout") |
| if layout != LAYOUT: |
| raise ValueError( |
| f"{directory}: expected a '{LAYOUT}' artifact from scripts/convert.py, got {layout!r}; " |
| "the old absolute-scale artifact is ~5% off and must be regenerated" |
| ) |
| return manifest |
|
|
|
|
| def load_weights(directory: str | Path = DEFAULT_DIR) -> dict[str, mx.array]: |
| """Every tensor in memory (~17.5 GiB) as name -> mx.array, stacked experts included.""" |
| directory = Path(directory) |
| tensors: dict[str, mx.array] = {} |
| for shard in load_manifest(directory)["shards"]: |
| tensors.update(mx.load(str(directory / shard["file"]))) |
| return tensors |
|
|
|
|
| class LazyWeights(Mapping[str, mx.array]): |
| """Name -> tensor mapping that loads a 2 GiB shard on first access and keeps it. |
| |
| One shard holds about three layers of stacked experts, so a layer needs one |
| shard, not the whole artifact. |
| """ |
|
|
| def __init__(self, directory: str | Path = DEFAULT_DIR): |
| self.directory = Path(directory) |
| self.manifest = load_manifest(self.directory) |
| self._weight_map: dict[str, str] = self.manifest["weight_map"] |
| self._shards: dict[str, dict[str, mx.array]] = {} |
|
|
| def __getitem__(self, key: str) -> mx.array: |
| shard = self._weight_map[key] |
| if shard not in self._shards: |
| self._shards[shard] = mx.load(str(self.directory / shard)) |
| return self._shards[shard][key] |
|
|
| def __iter__(self) -> Iterator[str]: |
| return iter(self._weight_map) |
|
|
| def __len__(self) -> int: |
| return len(self._weight_map) |
|
|
