up

Author: metascroy

backends/apple/mlx/README.md

@@ -8,7 +8,7 @@
 
 """MLX backend for ExecuTorch - executes models on Apple Silicon using MLX."""
 
-# Import custom_ops module to register custom ATen ops before anything else
+# Import custom_ops module to register custom ATen ops (rope, etc.)
 from executorch.backends.apple.mlx import custom_ops as _custom_ops  # noqa: F401
 from executorch.backends.apple.mlx.partitioner import MLXPartitioner
 

backends/apple/mlx/__init__.py

@@ -14,7 +14,6 @@
 can execute efficiently but may not have direct PyTorch equivalents.
 
 The ops are registered using torch.library and include:
-- rms_norm: RMSNorm normalization
 - rope: Rotary Position Embedding (single tensor)
 """
 
@@ -24,37 +23,6 @@
 from torch import Tensor
 
 
-# =============================================================================
-# rms_norm: RMSNorm normalization
-# =============================================================================
-
-
-@torch.library.custom_op("mlx::rms_norm", mutates_args=())
-def rms_norm(x: Tensor, weight: Tensor, eps: float = 1e-5) -> Tensor:
-    """
-    RMSNorm normalization.
-
-    Args:
-        x: Input tensor of shape (..., hidden_dim)
-        weight: Weight tensor of shape (hidden_dim,)
-        eps: Small constant for numerical stability
-
-    Returns:
-        Normalized tensor of the same shape as x
-    """
-    x_f = x.to(torch.float32)
-    var = x_f.pow(2).mean(dim=-1, keepdim=True)
-    y = x_f * torch.rsqrt(var + eps)
-    y = y.to(x.dtype)
-    return y * weight.to(x.dtype)
-
-
-@torch.library.register_fake("mlx::rms_norm")
-def rms_norm_fake(x: Tensor, weight: Tensor, eps: float = 1e-5) -> Tensor:
-    """Fake implementation for tracing."""
-    return x.new_empty(x.shape)
-
-
 # =============================================================================
 # rope: Rotary Position Embedding (single tensor)
 # =============================================================================

backends/apple/mlx/custom_ops.py

@@ -7,7 +7,7 @@ This example demonstrates how to export and run Llama models using the MLX deleg
 - **Export**: Convert HuggingFace Llama models to ExecutorCh format with MLX delegate
 - **Quantization**: Optional INT4/INT8 weight quantization via TorchAO
 - **KV Cache**: Efficient KV cache implementation for autoregressive generation
-- **Custom Ops**: Uses `mlx::rms_norm` and `mlx::apply_rope` for optimal MLX execution
+- **Custom Ops**: Uses `mlx::apply_rope` for optimal MLX execution
 - **Pybindings**: Run inference using ExecutorCh Python bindings
 
 ## Requirements
@@ -89,7 +89,7 @@ python -m executorch.backends.apple.mlx.examples.llama.run_llama \
 
 The example uses a custom model wrapper (`LlamaWithFunctionalKV`) that:
 
-1. **Replaces RMSNorm** with `torch.ops.mlx.rms_norm` - a custom op that maps directly to MLX's efficient RMSNorm implementation
+1. **Replaces RMSNorm** with `torch.nn.functional.rms_norm` - which maps to MLX's efficient RMSNorm implementation via the aten.rms_norm handler
 
 2. **Replaces Attention** with `KVCacheAttention` which:
    - Uses `torch.ops.mlx.apply_rope` for rotary position embeddings

backends/apple/mlx/examples/llama/README.md

@@ -33,7 +33,7 @@
 if TYPE_CHECKING:
     from transformers import AutoModelForCausalLM
 
-# Import custom MLX ops for rms_norm and apply_rope
+# Import MLX ops to register handlers
 import executorch.backends.apple.mlx.ops  # noqa: F401
 import torch
 import torch.nn as nn
@@ -51,12 +51,18 @@
 
 
 # =============================================================================
-# Custom RMSNorm using MLX op
+# Custom RMSNorm using aten op
 # =============================================================================
 
 
 class CustomRMSNorm(nn.Module):
-    """RMSNorm using the custom mlx::rms_norm op for efficient MLX execution."""
+    """RMSNorm using torch.nn.functional.rms_norm for efficient MLX execution.
+
+    This replaces the HuggingFace LlamaRMSNorm (which uses manual variance + rsqrt
+    computation) with the aten rms_norm op. The MLX backend has a handler that maps
+    aten.rms_norm directly to MLX's fast::rms_norm, so this gives us fused execution
+    without needing a custom op.
+    """
 
     def __init__(self, base_rms: nn.Module):
         super().__init__()
@@ -66,7 +72,7 @@ def __init__(self, base_rms: nn.Module):
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return torch.ops.mlx.rms_norm(x, self.weight, self.eps)
+        return F.rms_norm(x, (self.weight.shape[0],), self.weight, self.eps)
 
 
 # =============================================================================
@@ -290,7 +296,7 @@ def forward(self, hidden_states: torch.Tensor, pos_int: int) -> torch.Tensor:
 class LlamaWithFunctionalKV(nn.Module):
     """
     Wrapper around HuggingFace Llama that:
-    1. Replaces RMSNorm with custom mlx::rms_norm op
+    1. Replaces RMSNorm with aten rms_norm (mapped to MLX's fast::rms_norm)
     2. Replaces attention with KVCacheAttention (using mlx::apply_rope)
     3. Provides a trace-friendly forward that takes (token_ids, input_pos)
     """
@@ -299,7 +305,7 @@ def __init__(
         self,
         base: "AutoModelForCausalLM",
         time_axis: int = 1,
-        max_seq_len: int = 4096,
+        max_seq_len: int = 1024,
         rope_base: float = 500000.0,
     ):
         super().__init__()
@@ -383,19 +389,21 @@ def forward(self, token_ids: torch.Tensor, input_pos: torch.Tensor) -> torch.Ten
 def export_llama_to_mlx(
     model_id: str,
     output_path: str,
-    quantize: Optional[str] = None,
-    max_seq_len: int = 4096,
+    max_seq_len: int = 1024,
     dtype: str = "fp32",
+    quantize_linear: Optional[str] = None,
+    quantize_embeddings: Optional[str] = None,
 ) -> None:
     """
     Export a Llama model to MLX delegate.
 
     Args:
         model_id: HuggingFace model ID
         output_path: Path to save the .pte file
-        quantize: Quantization method ("int4", "int8", or None)
         max_seq_len: Maximum sequence length for KV cache
         dtype: Model dtype ("fp32", "fp16", "bf16")
+        quantize_linear: Quantization method for linear layers ("int4", "int8", or None)
+        quantize_embeddings: Quantization method for embedding layers ("int4", "int8", or None)
     """
     from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -416,47 +424,51 @@ def export_llama_to_mlx(
     model.eval()
 
     # Apply quantization if requested
-    if quantize:
-        logger.info(f"Applying {quantize} quantization...")
+    if quantize_linear or quantize_embeddings:
+        logger.info("Applying quantization with TorchAO...")
         try:
             from torchao.quantization.granularity import PerGroup
             from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
 
-            if quantize == "int4":
-                # Quantize embeddings with group size 32
-                quantize_(
-                    model,
-                    IntxWeightOnlyConfig(
-                        weight_dtype=torch.int4, granularity=PerGroup(32)
-                    ),
-                    lambda m, fqn: isinstance(m, torch.nn.Embedding),
+            # Quantize embedding layers
+            if quantize_embeddings:
+                embed_dtype = (
+                    torch.int4 if quantize_embeddings == "int4" else torch.int8
                 )
-                # Quantize linear layers with group size 64
-                quantize_(
-                    model,
-                    IntxWeightOnlyConfig(
-                        weight_dtype=torch.int4, granularity=PerGroup(64)
-                    ),
+                embed_group_size = 32 if quantize_embeddings == "int4" else 128
+                logger.info(
+                    f"Quantizing embedding layers with {quantize_embeddings} (group size {embed_group_size})..."
                 )
-            elif quantize == "int8":
                 quantize_(
                     model,
                     IntxWeightOnlyConfig(
-                        weight_dtype=torch.int8, granularity=PerGroup(32)
+                        weight_dtype=embed_dtype,
+                        granularity=PerGroup(embed_group_size),
                     ),
-                    lambda m, fqn: isinstance(m, torch.nn.Embedding),
+                    filter_fn=lambda m, fqn: isinstance(m, torch.nn.Embedding),
+                )
+
+            # Quantize linear layers
+            if quantize_linear:
+                linear_dtype = torch.int4 if quantize_linear == "int4" else torch.int8
+                linear_group_size = 32 if quantize_linear == "int4" else 128
+                logger.info(
+                    f"Quantizing linear layers with {quantize_linear} (group size {linear_group_size})..."
                 )
                 quantize_(
                     model,
                     IntxWeightOnlyConfig(
-                        weight_dtype=torch.int8, granularity=PerGroup(64)
+                        weight_dtype=linear_dtype,
+                        granularity=PerGroup(linear_group_size),
                     ),
+                    filter_fn=lambda m, fqn: isinstance(m, torch.nn.Linear),
                 )
-            else:
-                logger.warning(f"Unknown quantization method: {quantize}")
 
             # Tie lm_head weights to embedding after quantization
-            model.model.lm_head.weight = model.model.model.embed_tokens.weight
+            if quantize_embeddings:
+                model.model.lm_head.weight = model.model.model.embed_tokens.weight
+
+            logger.info("Applied quantization successfully")
         except ImportError:
             logger.error("TorchAO not installed. Run: pip install torchao")
             raise
@@ -484,11 +496,8 @@ def export_llama_to_mlx(
     import executorch.exir as exir
     from executorch.backends.apple.mlx import MLXPartitioner
     from executorch.exir import EdgeCompileConfig
-    from executorch.exir.backend.backend_details import CompileSpec
     from executorch.exir.capture._config import ExecutorchBackendConfig
 
-    compile_specs = [CompileSpec("use_fp16", bytes([False]))]
-
     # Allow repeat_interleave and sdpa ops - they will be handled by MLX backend
     edge_config = EdgeCompileConfig(
         _core_aten_ops_exception_list=[
@@ -499,7 +508,7 @@ def export_llama_to_mlx(
 
     edge_program = exir.to_edge_transform_and_lower(
         ep,
-        partitioner=[MLXPartitioner(compile_specs=compile_specs)],
+        partitioner=[MLXPartitioner()],
         compile_config=edge_config,
     )
 
@@ -538,35 +547,43 @@ def main():
         default="llama_mlx.pte",
         help="Output .pte file path",
     )
-    parser.add_argument(
-        "--quantize",
-        type=str,
-        choices=["int4", "int8"],
-        default=None,
-        help="Quantization method",
-    )
     parser.add_argument(
         "--max-seq-len",
         type=int,
-        default=4096,
+        default=1024,
         help="Maximum sequence length for KV cache",
     )
     parser.add_argument(
         "--dtype",
         type=str,
         choices=["fp32", "fp16", "bf16"],
-        default="fp32",
+        default="bf16",
         help="Model dtype (fp32, fp16, bf16)",
     )
+    parser.add_argument(
+        "--quantize-linear",
+        type=str,
+        choices=["int4", "int8"],
+        default=None,
+        help="Quantization method for linear layers",
+    )
+    parser.add_argument(
+        "--quantize-embeddings",
+        type=str,
+        choices=["int4", "int8"],
+        default=None,
+        help="Quantization method for embedding layers",
+    )
 
     args = parser.parse_args()
 
     export_llama_to_mlx(
         model_id=args.model_id,
         output_path=args.output,
-        quantize=args.quantize,
         max_seq_len=args.max_seq_len,
         dtype=args.dtype,
+        quantize_linear=args.quantize_linear,
+        quantize_embeddings=args.quantize_embeddings,
     )