[ExecuTorch][MLX] lowering memory/export improvements (pytorch#21795)

Muse Glimmer dflash + image (16 GB Q4_K target + 1.5 GB draft + 1.3 GB
mmproj):

| | before | after |
|---|---|---|
| export time | 4.2 min | **2.8 min** (−32%) |
| peak memory | 66.6 GB | **48.3 GB** (−27%) |

## What changed

1. **Support checks no longer repack weights.** The partitioner decided
"can MLX
lower this?" by *running* each handler, which for quantized ops means
fully
dequantizing and repacking the weight. exir asks several times per
export, so
every weight was repacked 4x. Handlers can now answer from node metadata
via an
optional `supported()` / `register_support_check`; anything without one
keeps
   the old run-and-catch, so nothing regresses.

2. **Raw blobs are freed as they're repacked.** `repack_mlx` left the
original
GGUF bytes in the ExportedProgram until a sweep that only ran after the
whole
graph was built, so raw and repacked weights were live simultaneously
(+16 GB).
   It now releases each blob on the spot, memoized per weight so tied
   embedding/lm_head repacks once.

3. **Two guards**, since a wrongly-released weight used to fail
silently:
serialization now raises naming the constant instead of emitting a
`.pte` with
a missing weight, and release is a no-op during support checks, where it
would
   corrupt the caller's own program.

4. **`ET_MLX_MEM_PROFILE=1`** adds phase-attributed peak-footprint
logging (off by
   default); it's what located the regression above.

Author: metascroy

.github/workflows/mlx.yml

@@ -194,6 +194,52 @@ jobs:
         echo "Success: built $RUNNER"
         echo "::endgroup::"
 
+  # Build-only: the Muse Glimmer GGUF checkpoints are far too large to export in
+  # CI, so this covers the compile / link / metallib wiring of the runners --
+  # which is what MLX backend changes are most likely to break.
+  test-mlx-muse-glimmer-build:
+    needs: run-decision
+    if: |
+      github.event_name == 'pull_request' ||
+      needs.run-decision.outputs.is-full-run == 'true'
+    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
+    with:
+      default-packages: ""
+      job-name: test-mlx-muse-glimmer-build
+      runner: macos-14-xlarge
+      python-version: "3.12"
+      submodules: recursive
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 90
+      script: |
+        set -eux
+
+        echo "::group::Install ExecuTorch"
+        ${CONDA_RUN} python install_executorch.py > /dev/null
+        echo "::endgroup::"
+
+        ${CONDA_RUN} pip list
+
+        echo "::group::Build Muse Glimmer runners and worker with MLX"
+        ${CONDA_RUN} make muse-glimmer-mlx
+        BIN_DIR=cmake-out/examples/models/muse-glimmer
+        for bin in solo_runner dflash_runner muse_glimmer_worker; do
+          if [ ! -x "${BIN_DIR}/${bin}" ]; then
+            echo "Failed: ${bin} not found at ${BIN_DIR}/${bin}"
+            exit 1
+          fi
+        done
+        # Runtime kernel load falls back to an absolute build-tree path, so a
+        # missing colocated metallib only breaks a relocated binary. Assert the
+        # copy so the runners stay self-contained (needs EXECUTORCH_BUILD_MLX=ON
+        # in the preset, which drives executorch_target_copy_mlx_metallib).
+        if [ ! -f "${BIN_DIR}/mlx.metallib" ]; then
+          echo "Failed: mlx.metallib not copied next to the runners"
+          exit 1
+        fi
+        echo "Success: built solo_runner, dflash_runner and muse_glimmer_worker"
+        echo "::endgroup::"
+
   backend-tester:
     needs: run-decision
     if: |

Makefile

@@ -91,7 +91,7 @@
 #
 # ==============================================================================
 
-.PHONY: voxtral-cuda voxtral-cpu voxtral-metal voxtral-mlx voxtral_realtime-cuda voxtral_realtime-cpu voxtral_realtime-metal voxtral_realtime-mlx voxtral_tts-cpu voxtral_tts-cuda whisper-cuda whisper-cuda-debug whisper-cpu whisper-metal parakeet-cuda parakeet-cuda-debug parakeet-cpu parakeet-metal parakeet-mlx parakeet-vulkan dinov2-cuda dinov2-cuda-debug sortformer-cuda sortformer-cpu silero-vad-cpu llama-cuda llama-cuda-debug llama-cpu lfm_2_5-mlx llava-cpu gemma3-cuda gemma3-cpu gemma4_31b-cuda gemma4_31b-mlx muse-glimmer-cuda qwen3_5_moe-cuda qwen3_5_moe-metal qwen3_5_moe-mlx clean help
+.PHONY: voxtral-cuda voxtral-cpu voxtral-metal voxtral-mlx voxtral_realtime-cuda voxtral_realtime-cpu voxtral_realtime-metal voxtral_realtime-mlx voxtral_tts-cpu voxtral_tts-cuda whisper-cuda whisper-cuda-debug whisper-cpu whisper-metal parakeet-cuda parakeet-cuda-debug parakeet-cpu parakeet-metal parakeet-mlx parakeet-vulkan dinov2-cuda dinov2-cuda-debug sortformer-cuda sortformer-cpu silero-vad-cpu llama-cuda llama-cuda-debug llama-cpu lfm_2_5-mlx llava-cpu gemma3-cuda gemma3-cpu gemma4_31b-cuda gemma4_31b-mlx muse-glimmer-cuda muse-glimmer-mlx qwen3_5_moe-cuda qwen3_5_moe-metal qwen3_5_moe-mlx clean help
 
 help:
 	@echo "This Makefile adds targets to build runners for various models on various backends. Run using \`make <target>\`. Available targets:"
@@ -130,6 +130,7 @@ help:
 	@echo "  gemma4_31b-cuda     - Build Gemma 4 31B runner and worker with CUDA backend"
 	@echo "  gemma4_31b-mlx      - Build Gemma 4 31B runner and worker with MLX backend"
 	@echo "  muse-glimmer-cuda   - Build Muse Glimmer runners and worker with CUDA backend"
+	@echo "  muse-glimmer-mlx    - Build Muse Glimmer runners and worker with MLX backend"
 	@echo "  qwen3_5_moe-cuda    - Build Qwen3.5 MoE runner with CUDA backend"
 	@echo "  qwen3_5_moe-metal   - Build Qwen3.5 MoE runner with Metal backend"
 	@echo "  qwen3_5_moe-mlx     - Build Qwen3.5 MoE runner with MLX backend"
@@ -453,6 +454,17 @@ muse-glimmer-cuda:
 	@echo "  DFlash runner: cmake-out/examples/models/muse-glimmer/dflash_runner"
 	@echo "  Worker:        cmake-out/examples/models/muse-glimmer/muse_glimmer_worker"
 
+muse-glimmer-mlx:
+	@echo "==> Building and installing ExecuTorch with MLX..."
+	cmake --workflow --preset mlx-release
+	@echo "==> Building Muse Glimmer runners and worker with MLX..."
+	cd examples/models/muse-glimmer && cmake --workflow --preset muse-glimmer-mlx
+	@echo ""
+	@echo "✓ Build complete!"
+	@echo "  Solo runner:   cmake-out/examples/models/muse-glimmer/solo_runner"
+	@echo "  DFlash runner: cmake-out/examples/models/muse-glimmer/dflash_runner"
+	@echo "  Worker:        cmake-out/examples/models/muse-glimmer/muse_glimmer_worker"
+
 gemma4_31b-cuda:
 	@echo "==> Building and installing ExecuTorch with CUDA..."
 	cmake --workflow --preset llm-release-cuda

backends/mlx/_memprofile.py

@@ -0,0 +1,188 @@
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+"""Phase-attributed peak-memory profiling for MLX export.
+
+Enabled with ``ET_MLX_MEM_PROFILE=1``; inert otherwise::
+
+    ET_MLX_MEM_PROFILE=1 python -m ...export_dflash --target-gguf ...
+
+    [mem] build: 42.84 GB -> 41.30 GB, -1.54 GB net, high-water 48.26 GB, \
+RAISED PEAK by 5.42 GB, 24.4s
+
+Reports macOS *physical footprint* -- what ``/usr/bin/time -l`` calls "peak
+memory footprint" and Activity Monitor shows as "Memory". RSS is not a usable
+proxy: a 16 GB model export peaked at 65 GB footprint versus 37 GB RSS.
+
+Only phases tagged ``RAISED PEAK`` can lower the number ``/usr/bin/time -l``
+reports; everything else is churn under the existing high-water mark.
+
+The watermark comes from the kernel rather than a sampling thread, which would
+miss short-lived spikes -- a 2 GB allocation freed within 100 ms went unseen at
+a 50 ms sampling interval.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import ctypes.util
+import os
+import time
+from contextlib import contextmanager
+from typing import Optional
+
+from executorch.backends.mlx._logging import logger
+
+_ENV_VAR = "ET_MLX_MEM_PROFILE"
+
+# proc_pid_rusage(pid, flavor, rusage_info_t *buffer), from <libproc.h>.
+# Offsets into struct rusage_info_v4 (<sys/resource.h>): a 16-byte ri_uuid
+# followed by uint64 counters. Sanity-checked at runtime by _self_check().
+_RUSAGE_INFO_V4 = 4
+_OFF_PHYS_FOOTPRINT = 72
+_OFF_PROC_START_ABSTIME = 80
+_OFF_PROC_EXIT_ABSTIME = 88
+_OFF_LIFETIME_MAX_PHYS_FOOTPRINT = 240
+_RUSAGE_BUF_BYTES = 512  # generous; larger than any rusage_info_v4
+
+# Bounds for the self-check: any live process is above the floor, and no real
+# footprint approaches the ceiling.
+_MIN_PLAUSIBLE = 1 << 20  # 1 MiB
+_MAX_PLAUSIBLE = 1 << 50  # 1 PiB
+
+_libc = None
+_usable: Optional[bool] = None
+_depth = 0
+
+
+def _read_rusage() -> Optional[bytes]:
+    global _libc
+    try:
+        if _libc is None:
+            _libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
+        buf = (ctypes.c_uint8 * _RUSAGE_BUF_BYTES)()
+        if _libc.proc_pid_rusage(os.getpid(), _RUSAGE_INFO_V4, ctypes.byref(buf)) != 0:
+            return None
+        return bytes(buf)
+    except (AttributeError, OSError, ValueError):
+        return None
+
+
+def _field(raw: Optional[bytes], offset: int) -> Optional[int]:
+    if raw is None:
+        return None
+    return int.from_bytes(raw[offset : offset + 8], "little")
+
+
+def _self_check() -> bool:
+    """Smoke-test the hardcoded rusage offsets before trusting what they return.
+
+    The struct layout is pinned by the flavor argument (Apple adds rusage_info_v5
+    rather than reordering v4), so this is a guard against being wrong about the
+    layout, not against it changing underneath us.
+
+    ri_proc_exit_abstime is the anchor: it is necessarily zero for the process
+    asking, so a non-zero read means the offsets are not landing where we think.
+    The bounds checks on the two footprint fields are weaker -- a misread can
+    still satisfy them by coincidence, since neighbouring counters hold numbers
+    of a similar magnitude -- so they catch gross errors only.
+    """
+    raw = _read_rusage()
+    current = _field(raw, _OFF_PHYS_FOOTPRINT)
+    peak = _field(raw, _OFF_LIFETIME_MAX_PHYS_FOOTPRINT)
+    started = _field(raw, _OFF_PROC_START_ABSTIME)
+    exited = _field(raw, _OFF_PROC_EXIT_ABSTIME)
+
+    if current is None:
+        reason = "proc_pid_rusage unavailable (non-macOS?)"
+    elif exited != 0 or not started:
+        reason = (
+            f"struct layout unrecognized (proc_start={started}, proc_exit={exited}; "
+            "expected a non-zero start and a zero exit for a live process)"
+        )
+    elif not _MIN_PLAUSIBLE <= current <= peak <= _MAX_PLAUSIBLE:
+        reason = f"implausible readings (current={current}, lifetime max={peak})"
+    else:
+        return True
+
+    logger.warning(f"[mem] memory profiling disabled: {reason}")
+    return False
+
+
+def enabled() -> bool:
+    """Whether profiling is switched on and the platform counters are trustworthy."""
+    global _usable
+    if os.environ.get(_ENV_VAR, "0") == "0":
+        return False
+    if _usable is None:
+        _usable = _self_check()
+    return _usable
+
+
+def phys_footprint() -> Optional[int]:
+    """Current physical footprint in bytes, or None if unavailable."""
+    return _field(_read_rusage(), _OFF_PHYS_FOOTPRINT)
+
+
+def peak_footprint() -> Optional[int]:
+    """Process lifetime maximum physical footprint in bytes.
+
+    Matches the "peak memory footprint" line from ``/usr/bin/time -l``.
+    """
+    return _field(_read_rusage(), _OFF_LIFETIME_MAX_PHYS_FOOTPRINT)
+
+
+def _gb(n: Optional[int]) -> str:
+    return "?" if n is None else f"{n / (1 << 30):.2f} GB"
+
+
+def _delta(before: Optional[int], after: Optional[int]) -> Optional[int]:
+    return None if (before is None or after is None) else after - before
+
+
+@contextmanager
+def mem_phase(name: str):
+    """Log footprint across `name`, attributing any new high-water mark to it.
+
+    Nested phases are indented. Does nothing unless ET_MLX_MEM_PROFILE is set.
+    """
+    global _depth
+
+    if not enabled():
+        yield
+        return
+
+    indent = "  " * _depth
+    start, start_peak = phys_footprint(), peak_footprint()
+    started_at = time.perf_counter()
+    _depth += 1
+    try:
+        yield
+    finally:
+        _depth -= 1
+        end, end_peak = phys_footprint(), peak_footprint()
+        net = _delta(start, end)
+        raised = _delta(start_peak, end_peak)
+
+        sign = "+" if net is not None and net >= 0 else ""
+        message = (
+            f"[mem]{indent} {name}: {_gb(start)} -> {_gb(end)}, "
+            f"{sign}{_gb(net)} net, high-water {_gb(end_peak)}"
+        )
+        if raised:
+            message += f", RAISED PEAK by {_gb(raised)}"
+        logger.info(f"{message}, {time.perf_counter() - started_at:.1f}s")
+
+
+def log_footprint(label: str) -> None:
+    """Log a one-off footprint reading."""
+    if enabled():
+        logger.info(
+            f"[mem] {label}: {_gb(phys_footprint())} "
+            f"(high-water {_gb(peak_footprint())})"
+        )

backends/mlx/builder/op_helpers.py

@@ -534,6 +534,28 @@ def emit_quantized_gather(
     )
 
 
+def mlx_qparams_supported(
+    in_features: int, num_groups: int, group_size: int, bits: int
+) -> bool:
+    """Whether to_mlx_qparams + regroup_affine_scales can consume this layout.
+
+    Mirrors their asserts using shape metadata only, so a handler can answer
+    "would I lower this?" without reading (and repacking) the weight:
+
+    * to_mlx_qparams packs a row into whole uint32 words, so
+      ``in_features * bits`` must be a multiple of 32.
+    * regroup_affine_scales repeat-interleaves, which only ever splits a group
+      finer, so the weight's own group must be a whole multiple of the
+      MLX-legal ``group_size``.
+    """
+    if in_features <= 0 or num_groups <= 0 or in_features % num_groups != 0:
+        return False
+    if (in_features * bits) % 32 != 0:
+        return False
+    weight_group_size = in_features // num_groups
+    return weight_group_size >= group_size and weight_group_size % group_size == 0
+
+
 def to_mlx_qparams(
     qdata: torch.Tensor,
     scale: torch.Tensor,