export.py

# 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.

"""Export Gemma 4 31B-IT to ExecuTorch (.pte + .ptd).

Two methods are exported and lowered together so they share KV-cache buffers:
  - "decode":  T=1, static shape, returns the next sampled token.
  - "prefill": T>=2, dynamic shape, returns the next sampled token.

Three input paths:
  --prequantized <dir>      Load a quantized checkpoint (from quantize_and_save.py)
                            and pack for the target backend. No re-quantization.
  --gguf <file>             Load a GGUF file (e.g., Q4_K_M from the community).
  --model-dir <hf>          Load bf16 checkpoint, quantize, pack, and export
                            in one shot.

Backends:
  --backend cuda            (default) CUDA via tinygemm INT4 + CudaPartitioner.
  --backend mlx             Apple Silicon via MLXPartitioner (single method,
                            dynamic seq_len, host-side sampling).
"""

import argparse
import json
import os

import torch
import torch.nn as nn
from executorch.examples.models.gemma4_31b.model import (
    Gemma4_31B,
    Gemma4_31BConfig,
    materialize_runtime_buffers,
)


# ---------------------------------------------------------------------------
# Load paths


def _build_and_pack(
    config: Gemma4_31BConfig,
    path: str,
    backend: str,
    *,
    key_map=None,
    tie_map=None,
) -> Gemma4_31B:
    """Build a meta-device model, stream ``path`` into it, and finalize.

    Shared by ``load_prequantized_model`` (safetensors) and ``load_gguf_model``
    (GGUF). ``load_checkpoint`` picks the reader by ``path`` and asserts every
    parameter is materialized; we then fill runtime buffers left on meta (RoPE
    tables, KV caches, scalar constants) and switch to eval. Pass ``key_map`` /
    ``tie_map`` for formats whose tensor names differ from the model's or that
    tie weights (GGUF).
    """
    from executorch.extension.llm.export.load import load_checkpoint

    # Validate the backend before any file I/O so an unknown backend fails fast
    # rather than erroring later on a missing/large checkpoint.
    if backend not in _SUPPORTED_BACKENDS:
        raise ValueError(
            f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}."
        )

    print("Building model on meta device...")
    with torch.device("meta"):
        model = Gemma4_31B(config)

    load_checkpoint(
        path,
        model,
        key_map=key_map,
        tie_map=tie_map,
        dtype=torch.bfloat16,
    )
    materialize_runtime_buffers(model, dtype=torch.bfloat16)
    _apply_backend_pass(model, backend)
    model.eval()

    print(f"Model: {config.num_hidden_layers} layers, hidden={config.hidden_size}")
    return model


def load_prequantized_model(
    prequantized_dir: str,
    max_seq_len: int = 4096,
    backend: str = "cuda",
) -> tuple[Gemma4_31B, Gemma4_31BConfig]:
    """Load a quantized safetensors checkpoint and pack for the target backend."""
    config = Gemma4_31BConfig.from_hf_config(
        os.path.join(prequantized_dir, "config.json")
    )
    config.max_seq_len = max_seq_len

    safetensors_path = os.path.join(prequantized_dir, "model.safetensors")
    print(f"Loading quantized checkpoint from {safetensors_path}...")
    model = _build_and_pack(config, safetensors_path, backend)
    return model, config


def load_gguf_model(
    gguf_path: str,
    max_seq_len: int = 4096,
    backend: str = "cuda",
    config: Gemma4_31BConfig | None = None,
) -> tuple[Gemma4_31B, Gemma4_31BConfig]:
    """Load a GGUF file and pack for the target backend.

    ``key_map`` (``gguf_to_model_key``) remaps GGUF tensor names to model FQNs.
    GGUF ties the token embedding and lm_head into one tensor; ``tie_map`` fans
    it out to ``lm_head`` -- untied on CUDA (``clone=True``: ``lm_head`` keeps a
    packed matmul weight while ``pack_embedding_for_cuda`` decodes a gatherable
    int8 copy) and kept tied on MLX (``clone=False``: both share one quantized
    tensor). ``config`` defaults to the full Gemma 4 31B config; pass a smaller
    one (e.g. in tests) to load a GGUF for a tiny model.
    """
    from executorch.examples.models.gemma4_31b.gguf_loader import gguf_to_model_key

    if config is None:
        config = Gemma4_31BConfig(max_seq_len=max_seq_len)

    tie_map = {"embed_tokens.weight": ("lm_head.weight", backend == "cuda")}
    print(f"Loading GGUF from {gguf_path}...")
    model = _build_and_pack(
        config,
        gguf_path,
        backend,
        key_map=gguf_to_model_key,
        tie_map=tie_map,
    )
    return model, config


def load_and_quantize(
    model_dir: str,
    recipe_name: str,
    max_seq_len: int = 4096,
    backend: str = "cuda",
) -> tuple[Gemma4_31B, Gemma4_31BConfig]:
    """Load bf16 checkpoint, quantize, pack — one shot."""
    from executorch.examples.models.gemma4_31b.quantize_and_save import _RECIPES
    from executorch.extension.llm.export.load import assign_state_dict
    from executorch.extension.llm.export.quant import quantize_model

    recipe = _RECIPES[recipe_name]

    print("Loading checkpoint (lazy, shard-by-shard)...")
    model, config = Gemma4_31B.from_hf_checkpoint(model_dir, max_seq_len=max_seq_len)

    if model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr():
        print("Untying embed_tokens / lm_head...")
        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())

    print(f"Quantizing with recipe '{recipe_name}'...")
    state_dict = quantize_model(model, recipe, verbose=True)

    print(f"Packing for {backend}...")
    with torch.device("meta"):
        model = Gemma4_31B(config)
    assign_state_dict(model, state_dict)
    _apply_backend_pass(model, backend)
    model.eval()

    print(f"Model: {config.num_hidden_layers} layers, hidden={config.hidden_size}")
    return model, config


# ---------------------------------------------------------------------------
# Backend dispatch helpers


_SUPPORTED_BACKENDS = ("cuda", "mlx")


def _apply_backend_pass(model: nn.Module, backend: str) -> None:
    """Run the backend's terminal layout-conversion pass over a loaded model.

    CUDA repacks quantized weights into the coalesced/planar decode layouts; MLX
    (gemma4 is dense — no switch linears) needs no extra pass, the portable forms
    export directly.
    """
    if backend == "mlx":
        return
    if backend == "cuda":
        from executorch.examples.models.gemma4_31b.cuda_packers import (
            convert_quantized_tensors_for_cuda,
        )

        convert_quantized_tensors_for_cuda(model)
        return
    raise ValueError(
        f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}."
    )


# ---------------------------------------------------------------------------
# Export + lower


def _mutable_buffer_metadata(model: nn.Module) -> str:
    mutable = [name for name, _ in model.named_buffers() if ".kv_cache." in name]
    return json.dumps({"version": 1, "mutable_buffers": mutable})


def export_and_lower(
    model: Gemma4_31B,
    config: Gemma4_31BConfig,
    output_dir: str,
    backend: str = "cuda",
    use_turboquant: bool = False,
    sample: bool = False,
) -> None:
    """Export and lower the model to ExecuTorch for the given backend."""
    if backend == "cuda":
        _export_cuda(model, config, output_dir, use_turboquant=use_turboquant)
    elif backend == "mlx":
        _export_mlx(
            model, config, output_dir, use_turboquant=use_turboquant, sample=sample
        )
    else:
        raise ValueError(
            f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}."
        )


def _export_cuda(
    model: Gemma4_31B,
    config: Gemma4_31BConfig,
    output_dir: str,
    use_turboquant: bool = False,
) -> None:
    import gc

    import torch._inductor.config as inductor_config
    from executorch.backends.cuda.cuda_backend import CudaBackend
    from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
    from executorch.exir import (
        EdgeCompileConfig,
        ExecutorchBackendConfig,
        to_edge_transform_and_lower,
    )
    from executorch.exir.backend.compile_spec_schema import CompileSpec
    from executorch.exir.passes import MemoryPlanningPass
    from executorch.exir.passes.propagate_device_pass import PropagateDeviceConfig
    from torch.export import Dim, export

    inductor_config.coordinate_descent_tuning = False
    inductor_config.aot_inductor.compile_wrapper_opt_level = "O0"

    # Register the quantized F.linear dispatch (decode GEMMs → Triton ops)
    import executorch.backends.cuda.quantize_op_dispatch  # noqa: F401

    materialize_runtime_buffers(model, dtype=torch.bfloat16)
    mutable_buffer_metadata = _mutable_buffer_metadata(model)

    from executorch.examples.models.gemma4_31b.cuda_source_transformations import (
        cuda_source_transformations,
    )

    cuda_source_transformations(model, use_turboquant=use_turboquant)

    # Int4Tensor weights are used directly — no format conversion.
    # F.linear dispatches decode (static M<=4) to the Triton INT4 bucket kernels.
    # Both decode and prefill share the same nibble-packed weights.

    # Prefill: dequant + cuBLAS (optimal for large M).
    max_prefill = min(config.max_seq_len - 1, config.sliding_window * 2)
    seq_dim = Dim("seq_len", min=5, max=max_prefill)
    print(f"Exporting prefill (T in [2, {max_prefill}])...")
    with torch.no_grad():
        prefill_ep = export(
            model,
            (
                torch.zeros((1, max_prefill), dtype=torch.long),
                torch.arange(max_prefill, dtype=torch.long),
                torch.tensor([1.0], dtype=torch.float32),
            ),
            dynamic_shapes=({1: seq_dim}, {0: seq_dim}, None),
            strict=True,
        )

    # Decode (T=1): same Int4Tensor weights, same format. No transform needed.
    print("Exporting decode (T=1)...")
    with torch.no_grad():
        decode_ep = export(
            model,
            (
                torch.tensor([[0]], dtype=torch.long),
                torch.tensor([0], dtype=torch.long),
                torch.tensor([1.0], dtype=torch.float32),
            ),
            strict=True,
        )

    del model
    gc.collect()

    print("Lowering to ExecuTorch with CUDA backend...")
    et_prog = to_edge_transform_and_lower(
        {"decode": decode_ep, "prefill": prefill_ep},
        partitioner={
            "decode": [
                CudaPartitioner(
                    [
                        CudaBackend.generate_method_name_compile_spec("decode"),
                        CompileSpec("low_memory_mode", b"ON"),
                    ]
                )
            ],
            "prefill": [
                CudaPartitioner(
                    [
                        CudaBackend.generate_method_name_compile_spec("prefill"),
                        CompileSpec("low_memory_mode", b"ON"),
                    ]
                )
            ],
        },
        compile_config=EdgeCompileConfig(
            _check_ir_validity=False,
            _skip_dim_order=True,
        ),
        constant_methods={
            "get_max_seq_len": config.max_seq_len,
            "get_vocab_size": config.vocab_size,
            "get_n_layers": config.num_hidden_layers,
            "get_max_prefill_chunk": max_prefill,
            "get_min_prefill_chunk": 5,
            "get_mutable_buffer_metadata": mutable_buffer_metadata,
            "use_kv_cache": True,
            "use_sdpa_with_kv_cache": False,
            "enable_dynamic_shape": True,
        },
    )
    del decode_ep, prefill_ep
    gc.collect()

    et_program = et_prog.to_executorch(
        config=ExecutorchBackendConfig(
            extract_delegate_segments=True,
            do_quant_fusion_and_const_prop=True,
            memory_planning_pass=MemoryPlanningPass(
                alloc_graph_input=False,
            ),
            emit_mutable_buffer_names=True,
            # Keep method inputs/outputs device-resident so the CUDA backend
            # does not insert boundary H2D/D2H copies: the runner stages inputs
            # in CUDA memory and reads the sampled token back with a single
            # small D2H. CUDA-only (no effect on the MLX path).
            propagate_device_config=PropagateDeviceConfig(
                skip_h2d_for_method_inputs=True,
                skip_d2h_for_method_outputs=True,
            ),
        ),
    )

    del et_prog
    gc.collect()

    os.makedirs(output_dir, exist_ok=True)
    pte_path = os.path.join(output_dir, "model.pte")
    print(f"Saving to {pte_path}...")
    with open(pte_path, "wb") as f:
        et_program.write_to_file(f)
    print(f"  {os.path.getsize(pte_path) / 1024**2:.1f} MB")

    if et_program._tensor_data:
        et_program.write_tensor_data_to_file(output_dir)
        print(f"  Saved tensor data (.ptd) to {output_dir}/")
    print("Done.")


def _export_mlx(
    model: Gemma4_31B,
    config: Gemma4_31BConfig,
    output_dir: str,
    use_turboquant: bool = False,
    sample: bool = False,
) -> None:
    """Export to .pte via torch.export + MLX backend.

    Unlike CUDA (which exports separate decode/prefill methods with an
    Int4Tensor dispatch override), MLX uses a single method with dynamic
    sequence length.  No quantize_op_dispatch import — IntxUnpackedToInt8Tensor's
    default dispatch produces the ``dequantize_affine → linear`` pattern
    that MLX's QuantizedLinearHandler matches.

    When ``use_turboquant=True``, full-attention layers swap to
    ``MLXTurboQuantKVCache`` for ~3.8× KV cache memory savings. Sliding
    layers are unaffected (already use ``RingBufferKVCache``).
    """
    import gc

    # Register the GGUF dequant op + MLX GGUF pattern handlers so quantized GGUF
    # weights lower to the fused Q6_K kernels / Q4_K quantized matmul.
    import executorch.backends.mlx.custom_kernel_ops.gguf.patterns  # noqa: F401
    import executorch.extension.llm.export.gguf  # noqa: F401
    import executorch.extension.llm.export.int4  # noqa: F401

    from executorch.backends.mlx import MLXPartitioner
    from executorch.backends.mlx.passes import get_default_passes
    from executorch.examples.models.gemma4_31b.mlx_source_transformations import (
        mlx_source_transformations,
    )
    from executorch.exir import (
        EdgeCompileConfig,
        ExecutorchBackendConfig,
        to_edge_transform_and_lower,
    )
    from executorch.exir.passes import MemoryPlanningPass
    from torch.export import Dim, export

    max_prefill = 256

    mlx_source_transformations(
        model,
        dtype=torch.bfloat16,
        use_turboquant=use_turboquant,
        max_write_len=max_prefill,
    )

    materialize_runtime_buffers(model, dtype=torch.bfloat16)

    seq_dim = Dim("seq_len", min=1, max=max_prefill)

    example_tokens = torch.tensor([[0, 1]], dtype=torch.long)
    example_input_pos = torch.tensor([0, 1], dtype=torch.long)
    if sample:
        # forward(tokens, input_pos, temperature, top_k, top_p, seed) -> token id.
        # gemma's MLX forward already returns last-token logits (B, vocab), so
        # SamplingHead is used directly with no per-model wrapper.
        from executorch.backends.mlx.llm.sampling import SamplingHead

        model = SamplingHead(model)
        example_args = (
            example_tokens,
            example_input_pos,
            torch.tensor(1.0, dtype=torch.float32),
            torch.tensor(torch.iinfo(torch.int64).max, dtype=torch.int64),
            torch.tensor(1.0, dtype=torch.float32),
            torch.tensor(0, dtype=torch.int64),
        )
        # SamplingHead.forward takes ``*args``; dynamic_shapes mirrors that single
        # variadic parameter as one nested tuple over the positional inputs.
        dynamic_shapes = (({1: seq_dim}, {0: seq_dim}, None, None, None, None),)
    else:
        example_args = (example_tokens, example_input_pos)
        dynamic_shapes = ({1: seq_dim}, {0: seq_dim})

    print(f"Exporting (T in [1, {max_prefill}])...")
    with torch.no_grad():
        exported = export(
            model,
            example_args,
            dynamic_shapes=dynamic_shapes,
            strict=True,
        )

    del model
    gc.collect()

    print("Lowering to ExecuTorch with MLX backend...")
    et_prog = to_edge_transform_and_lower(
        exported,
        transform_passes=get_default_passes(),
        partitioner=[MLXPartitioner()],
        compile_config=EdgeCompileConfig(
            _check_ir_validity=False,
            _skip_dim_order=True,
        ),
        constant_methods={
            "get_max_seq_len": config.max_seq_len,
            "get_vocab_size": config.vocab_size,
            "get_n_layers": config.num_hidden_layers,
            "get_max_prefill_chunk": max_prefill,
            "use_kv_cache": True,
            "use_sdpa_with_kv_cache": False,
            "enable_dynamic_shape": True,
            "use_sampling": sample,
        },
    )

    del exported
    gc.collect()

    et_program = et_prog.to_executorch(
        config=ExecutorchBackendConfig(
            extract_delegate_segments=True,
            memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
        ),
    )

    del et_prog
    gc.collect()

    os.makedirs(output_dir, exist_ok=True)
    pte_path = os.path.join(output_dir, "model.pte")
    print(f"Saving to {pte_path}...")
    with open(pte_path, "wb") as f:
        et_program.write_to_file(f)
    print(f"  {os.path.getsize(pte_path) / 1024**2:.1f} MB")

    if et_program._tensor_data:
        et_program.write_tensor_data_to_file(output_dir)
        print(f"  Saved tensor data (.ptd) to {output_dir}/")
    print("Done.")


# ---------------------------------------------------------------------------
# CLI


def main() -> None:
    from executorch.examples.models.gemma4_31b.quantize_and_save import _RECIPES

    parser = argparse.ArgumentParser(description="Export Gemma 4 31B-IT to ExecuTorch.")
    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument(
        "--model-dir",
        default=None,
        help="HuggingFace model dir. Triggers load + quantize + export.",
    )
    src.add_argument(
        "--prequantized",
        default=None,
        help="Path to a quantized checkpoint directory. Skips quantization.",
    )
    src.add_argument(
        "--gguf",
        default=None,
        help="Path to a GGUF file (e.g., gemma-4-31B-it-Q4_K_M.gguf).",
    )
    parser.add_argument(
        "--output-dir",
        default="./gemma4_31b_exports",
        help="Output directory for model.pte / model.ptd.",
    )
    parser.add_argument(
        "--max-seq-len",
        type=int,
        default=4096,
        help="KV cache size.",
    )
    parser.add_argument(
        "--quant-recipe",
        default="default",
        choices=list(_RECIPES),
        help="Quantization recipe (only with --model-dir).",
    )
    parser.add_argument(
        "--backend",
        default="cuda",
        choices=list(_SUPPORTED_BACKENDS),
        help="Target backend for export.",
    )
    parser.add_argument(
        "--turboquant",
        action="store_true",
        help="Use TurboQuant TQ4 KV cache compression. ~3.8× cache memory "
        "savings; applies only to full-attention (non-sliding) layers — "
        "sliding layers keep their default cache. Supported on both "
        "--backend mlx and --backend cuda.",
    )
    parser.add_argument(
        "--sample",
        action="store_true",
        help="MLX only: sample the next token on-device (Gumbel-max with "
        "temperature/top_p/seed runtime inputs) instead of returning logits "
        "for host-side sampling.",
    )
    args = parser.parse_args()

    if args.backend == "cuda" and not torch.cuda.is_available():
        parser.error("CUDA is required for the cuda backend.")

    if args.sample and args.backend != "mlx":
        parser.error("--sample is only supported with --backend mlx")

    if args.prequantized:
        model, config = load_prequantized_model(
            args.prequantized,
            max_seq_len=args.max_seq_len,
            backend=args.backend,
        )
    elif args.gguf:
        model, config = load_gguf_model(
            args.gguf, max_seq_len=args.max_seq_len, backend=args.backend
        )
    else:
        model, config = load_and_quantize(
            args.model_dir,
            args.quant_recipe,
            max_seq_len=args.max_seq_len,
            backend=args.backend,
        )

    export_and_lower(
        model,
        config,
        args.output_dir,
        backend=args.backend,
        use_turboquant=args.turboquant,
        sample=args.sample,
    )


if __name__ == "__main__":
    main()