|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""Export a supported Hugging Face decoder-only LM through the native backend. |
| 8 | +
|
| 9 | +Builds the model with ``AutoModelForCausalLM`` from a Hugging Face model |
| 10 | +directory (or only its config, with random weights), wraps it with the cache |
| 11 | +adapter from ``transformers.integrations.executorch``, and exports the |
| 12 | +``(input_ids [1, T], cache_position [T])`` signature. Reports which operators |
| 13 | +the native partitioner leaves behind. |
| 14 | +""" |
| 15 | + |
| 16 | +import argparse |
| 17 | +import collections |
| 18 | +import logging |
| 19 | +import operator |
| 20 | + |
| 21 | +import torch |
| 22 | +import torch.nn as nn |
| 23 | +from executorch.backends.native import get_default_compile_config |
| 24 | +from executorch.backends.native.partitioner import NativePartitioner |
| 25 | +from executorch.backends.native.passes import get_default_passes |
| 26 | +from executorch.exir import to_edge_transform_and_lower |
| 27 | +from executorch.exir.lowered_backend_module import ( |
| 28 | + executorch_call_delegate, |
| 29 | + get_lowered_submodules, |
| 30 | +) |
| 31 | +from executorch.exir.native import to_native |
| 32 | +from torch.export import Dim, export |
| 33 | +from torchao.quantization import ( |
| 34 | + Int8DynamicActivationIntxWeightConfig, |
| 35 | + IntxWeightOnlyConfig, |
| 36 | + quantize_, |
| 37 | +) |
| 38 | +from torchao.quantization.granularity import PerGroup |
| 39 | +from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel |
| 40 | +from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM |
| 41 | + |
| 42 | + |
| 43 | +def _load(path: str, layers: int, random_weights: bool) -> PreTrainedModel: |
| 44 | + config = AutoConfig.from_pretrained(path) |
| 45 | + text_config = config.get_text_config() |
| 46 | + if layers: |
| 47 | + text_config.num_hidden_layers = layers |
| 48 | + if getattr(text_config, "layer_types", None): |
| 49 | + text_config.layer_types = text_config.layer_types[:layers] |
| 50 | + kwargs = {"torch_dtype": torch.float32, "attn_implementation": "sdpa"} |
| 51 | + if random_weights: |
| 52 | + model = AutoModelForCausalLM.from_config(config, **kwargs) |
| 53 | + else: |
| 54 | + model = AutoModelForCausalLM.from_pretrained(path, config=config, **kwargs) |
| 55 | + model.generation_config.use_cache = True |
| 56 | + model.generation_config.cache_implementation = "static" |
| 57 | + return model.eval() |
| 58 | + |
| 59 | + |
| 60 | +class LastTokenLM(nn.Module): |
| 61 | + """Use the Transformers ExecuTorch cache wrapper and return last-position logits.""" |
| 62 | + |
| 63 | + def __init__(self, model: PreTrainedModel, max_cache_len: int) -> None: |
| 64 | + super().__init__() |
| 65 | + # TODO(PTN): attention exports as masked, non-causal GQA aten.sdpa over |
| 66 | + # the whole static cache; engines must pattern-match it back to causal |
| 67 | + # SDPA. |
| 68 | + self.cached = TorchExportableModuleForDecoderOnlyLM( |
| 69 | + model, batch_size=1, max_cache_len=max_cache_len |
| 70 | + ).model |
| 71 | + self.cache = ( |
| 72 | + self.cached.static_cache |
| 73 | + if hasattr(self.cached, "static_cache") |
| 74 | + else self.cached.cache |
| 75 | + ) |
| 76 | + |
| 77 | + def forward( |
| 78 | + self, input_ids: torch.Tensor, cache_position: torch.Tensor |
| 79 | + ) -> torch.Tensor: |
| 80 | + return self.cached.model( |
| 81 | + input_ids=input_ids, |
| 82 | + cache_position=cache_position, |
| 83 | + attention_mask=None, |
| 84 | + past_key_values=self.cache, |
| 85 | + use_cache=True, |
| 86 | + logits_to_keep=1, |
| 87 | + ).logits |
| 88 | + |
| 89 | + |
| 90 | +# TEMPORARY: use torchao-based 8da4w quantization until PTN export has its own |
| 91 | +# quantization flow. |
| 92 | +def _quantize_8da4w( |
| 93 | + model: PreTrainedModel, group_size: int, embedding_group_size: int |
| 94 | +) -> None: |
| 95 | + quantize_( |
| 96 | + model, |
| 97 | + IntxWeightOnlyConfig( |
| 98 | + weight_dtype=torch.int4, |
| 99 | + granularity=PerGroup(embedding_group_size), |
| 100 | + intx_choose_qparams_algorithm="hqq_scale_only", |
| 101 | + ), |
| 102 | + filter_fn=lambda m, _: isinstance(m, nn.Embedding), |
| 103 | + ) |
| 104 | + quantize_( |
| 105 | + model, |
| 106 | + Int8DynamicActivationIntxWeightConfig( |
| 107 | + weight_dtype=torch.int4, |
| 108 | + weight_granularity=PerGroup(group_size), |
| 109 | + intx_choose_qparams_algorithm="hqq_scale_only", |
| 110 | + ), |
| 111 | + filter_fn=lambda m, _: isinstance(m, nn.Linear) |
| 112 | + and m.in_features % group_size == 0, |
| 113 | + ) |
| 114 | + |
| 115 | + |
| 116 | +def _histogram(gm: torch.fx.GraphModule) -> collections.Counter: |
| 117 | + return collections.Counter( |
| 118 | + str(n.target) |
| 119 | + for n in gm.graph.nodes |
| 120 | + if n.op == "call_function" |
| 121 | + and n.target not in (executorch_call_delegate, operator.getitem) |
| 122 | + ) |
| 123 | + |
| 124 | + |
| 125 | +def main() -> None: |
| 126 | + parser = argparse.ArgumentParser() |
| 127 | + parser.add_argument( |
| 128 | + "--model", required=True, help="Hugging Face model directory or hub ID." |
| 129 | + ) |
| 130 | + parser.add_argument( |
| 131 | + "--random_weights", |
| 132 | + action="store_true", |
| 133 | + help="Use only the model's config, with random weights.", |
| 134 | + ) |
| 135 | + parser.add_argument("--layers", type=int, default=0) |
| 136 | + parser.add_argument("--max_seq_length", type=int, default=128) |
| 137 | + parser.add_argument("--max_context_length", type=int, default=0) |
| 138 | + parser.add_argument("--output", default="") |
| 139 | + parser.add_argument("--non_strict", action="store_true") |
| 140 | + parser.add_argument("--quantize", choices=["none", "8da4w"], default="none") |
| 141 | + parser.add_argument("--group_size", type=int, default=128) |
| 142 | + parser.add_argument("--embedding_group_size", type=int, default=32) |
| 143 | + args = parser.parse_args() |
| 144 | + cache_len = args.max_context_length or args.max_seq_length |
| 145 | + if cache_len < args.max_seq_length - 1: |
| 146 | + parser.error("the context must hold a full prompt of max_seq_length - 1 tokens") |
| 147 | + logging.basicConfig(level=logging.INFO) |
| 148 | + torch.manual_seed(0) |
| 149 | + |
| 150 | + model = _load(args.model, args.layers, args.random_weights) |
| 151 | + if args.quantize == "8da4w": |
| 152 | + _quantize_8da4w(model, args.group_size, args.embedding_group_size) |
| 153 | + wrapper = LastTokenLM(model, cache_len).eval() |
| 154 | + |
| 155 | + input_ids = torch.zeros((1, 6), dtype=torch.long) |
| 156 | + cache_position = torch.arange(input_ids.shape[1], dtype=torch.long) |
| 157 | + seq = Dim("seq", min=1, max=args.max_seq_length - 1) |
| 158 | + with torch.no_grad(): |
| 159 | + ep = export( |
| 160 | + wrapper, |
| 161 | + (input_ids, cache_position), |
| 162 | + dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}}, |
| 163 | + strict=not args.non_strict, |
| 164 | + ) |
| 165 | + print("== exported aten ops ==", flush=True) |
| 166 | + for k, v in sorted(_histogram(ep.graph_module).items()): |
| 167 | + print(f" {v:4d} {k}") |
| 168 | + |
| 169 | + edge = to_edge_transform_and_lower( |
| 170 | + ep, |
| 171 | + transform_passes=get_default_passes(), |
| 172 | + partitioner=[ |
| 173 | + NativePartitioner(external_constants_tag=None, _serialize_as_ptn=True) |
| 174 | + ], |
| 175 | + compile_config=get_default_compile_config(), |
| 176 | + ) |
| 177 | + gm = edge.exported_program().graph_module |
| 178 | + lowered = get_lowered_submodules(gm) |
| 179 | + print(f"== {len(lowered)} native delegate(s) ==") |
| 180 | + delegated = collections.Counter() |
| 181 | + for _, module, _ in lowered: |
| 182 | + delegated += _histogram(module.original_module.graph_module) |
| 183 | + for k, v in sorted(delegated.items()): |
| 184 | + print(f" {v:4d} {k}") |
| 185 | + print("== residual (not delegated) ==") |
| 186 | + for k, v in sorted(_histogram(gm).items()): |
| 187 | + print(f" {v:4d} {k}") |
| 188 | + |
| 189 | + if args.output: |
| 190 | + to_native(ep).save(args.output) |
| 191 | + print(f"saved {args.output}") |
0 commit comments