diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index ac5f2ba4748..a15268e9751 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -565,6 +565,19 @@ def build_args_parser() -> argparse.ArgumentParser: help="Use SpinQuant for better quantization performance. Only support cuda and native.", ) + parser.add_argument( + "--use_moe_quantized_op", + action="store_true", + default=False, + help=( + "Replace eager MoE feed-forward modules with the " + "`llama::quantized_moe_ffn` portable-runtime custom op (INT4 " + "weights, INT8 dyn-quant activations via torchao). On aarch64 " + "with ENABLE_QUANTIZED_MOE_FFN the optimized torchao kernel is " + "used; otherwise a portable reference fallback runs." + ), + ) + parser.add_argument( "-qat", "--use_qat", @@ -809,6 +822,7 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager: use_torchao_kernels_linear=llm_config.backend.torchao.use_torchao_kernels_linear, use_torchao_kernels_tied_embedding=llm_config.backend.torchao.use_torchao_kernels_tied_embedding, quantize_with_hqq=llm_config.quantization.use_hqq, + use_moe_quantized_op=llm_config.model.use_moe_quantized_op, ) ) @@ -1712,6 +1726,7 @@ def _get_source_transforms( # noqa use_torchao_kernels_linear: bool = False, use_torchao_kernels_tied_embedding: bool = False, quantize_with_hqq: bool = True, + use_moe_quantized_op: bool = False, ) -> List[Callable[[torch.nn.Module], torch.nn.Module]]: """ Return a list of functions that transform a graph. @@ -1769,6 +1784,17 @@ def _get_source_transforms( # noqa transforms.append(inject_fast_hadamard_transform_native_for_spin_quant) + if use_moe_quantized_op: + from .source_transformation.moe import replace_moe_with_quantized_op + + transforms.append( + partial( + replace_moe_with_quantized_op, + group_size=group_size or 32, + weight_nbit=4, + ) + ) + if embedding_quantize: """ When this option is selected, it finds all embedding layers and transforms diff --git a/extension/llm/custom_ops/BUCK b/extension/llm/custom_ops/BUCK index fbb3f02e9d1..2032d6d353c 100644 --- a/extension/llm/custom_ops/BUCK +++ b/extension/llm/custom_ops/BUCK @@ -117,6 +117,9 @@ fbcode_target(_kind = runtime.python_test, preload_deps = [ ":custom_ops_aot_lib_mkl_noomp", ":custom_ops_aot_py", + # The Python source transform calls + # `torch.ops.torchao._pack_8bit_act_4bit_weight` to pack each + # expert's INT4 weights, so preload the torchao AOT library too. "//pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_aten", ], deps = [ @@ -125,5 +128,7 @@ fbcode_target(_kind = runtime.python_test, "//executorch/examples/models/llama:llama_transformer", "//executorch/examples/models/llama:transformer_modules", "//executorch/examples/models/llama:source_transformation", + "//executorch/examples/models/llama:export_library", + "//executorch/extension/llm/export/config:llm_config", ], ) diff --git a/extension/llm/custom_ops/test_quantized_moe.py b/extension/llm/custom_ops/test_quantized_moe.py index 18d7d6023c1..4b296debb5d 100644 --- a/extension/llm/custom_ops/test_quantized_moe.py +++ b/extension/llm/custom_ops/test_quantized_moe.py @@ -596,3 +596,110 @@ def test_no_shared_expert_is_none(self) -> None: wrapper.m = moe replace_moe_with_quantized_op(wrapper, group_size=32, weight_nbit=4) self.assertIsNone(wrapper.m.shared_expert) + + +class TestExportPipelineWiring(unittest.TestCase): + """The export pipeline correctly includes the MoE transform.""" + + def test_get_source_transforms_includes_moe_when_enabled(self) -> None: + from functools import partial + + from executorch.examples.models.llama.export_llama_lib import ( + _get_source_transforms, + ) + + transforms = _get_source_transforms( + dtype_override=torch.float32, use_moe_quantized_op=True + ) + moe_transforms = [ + t + for t in transforms + if isinstance(t, partial) + and t.func.__name__ == "replace_moe_with_quantized_op" + ] + self.assertEqual(len(moe_transforms), 1) + self.assertEqual(moe_transforms[0].keywords["group_size"], 32) + self.assertEqual(moe_transforms[0].keywords["weight_nbit"], 4) + + def test_get_source_transforms_excludes_moe_when_disabled(self) -> None: + from functools import partial + + from executorch.examples.models.llama.export_llama_lib import ( + _get_source_transforms, + ) + + transforms = _get_source_transforms( + dtype_override=torch.float32, use_moe_quantized_op=False + ) + moe_transforms = [ + t + for t in transforms + if isinstance(t, partial) + and hasattr(t.func, "__name__") + and t.func.__name__ == "replace_moe_with_quantized_op" + ] + self.assertEqual(len(moe_transforms), 0) + + def test_get_source_transforms_passes_custom_group_size(self) -> None: + from functools import partial + + from executorch.examples.models.llama.export_llama_lib import ( + _get_source_transforms, + ) + + transforms = _get_source_transforms( + dtype_override=torch.float32, + use_moe_quantized_op=True, + group_size=64, + ) + moe_transforms = [ + t + for t in transforms + if isinstance(t, partial) + and t.func.__name__ == "replace_moe_with_quantized_op" + ] + self.assertEqual(len(moe_transforms), 1) + self.assertEqual(moe_transforms[0].keywords["group_size"], 64) + + def test_sentinel_op_is_registered(self) -> None: + self.assertTrue(hasattr(torch.ops.llama, "_quantized_moe_ffn_active")) + self.assertTrue(torch.ops.llama._quantized_moe_ffn_active()) + + +class TestLlmConfigMoeFlag(unittest.TestCase): + """llm_config wires the --use_moe_quantized_op flag correctly.""" + + def test_from_args_sets_flag(self) -> None: + import argparse + + from executorch.extension.llm.export.config.llm_config import LlmConfig + + args = argparse.Namespace(use_moe_quantized_op=True) + config = LlmConfig.from_args(args) + self.assertTrue(config.model.use_moe_quantized_op) + + def test_default_is_false(self) -> None: + from executorch.extension.llm.export.config.llm_config import ModelConfig + + self.assertFalse(ModelConfig().use_moe_quantized_op) + + def test_from_args_missing_field_defaults_false(self) -> None: + import argparse + + from executorch.extension.llm.export.config.llm_config import LlmConfig + + args = argparse.Namespace() + config = LlmConfig.from_args(args) + self.assertFalse(config.model.use_moe_quantized_op) + + def test_argparser_flag_true(self) -> None: + from executorch.examples.models.llama.export_llama_lib import build_args_parser + + args = build_args_parser().parse_args(["--use_moe_quantized_op"]) + self.assertTrue(args.use_moe_quantized_op) + + def test_argparser_flag_default_false(self) -> None: + from executorch.examples.models.llama.export_llama_lib import build_args_parser + + args = build_args_parser().parse_args([]) + self.assertFalse(args.use_moe_quantized_op) diff --git a/extension/llm/export/config/llm_config.py b/extension/llm/export/config/llm_config.py index 76f2f5a92c6..f1f93ba28c6 100644 --- a/extension/llm/export/config/llm_config.py +++ b/extension/llm/export/config/llm_config.py @@ -205,6 +205,9 @@ class ModelConfig: use_kv_cache: bool = False quantize_kv_cache: bool = False local_global_attention: Optional[List[int]] = None + # Replace eager MOEFeedForward modules with the + # `llama::quantized_moe_ffn` portable-runtime custom op. + use_moe_quantized_op: bool = False def __post_init__(self): self._validate_attention_sink() @@ -728,6 +731,8 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig": # noqa: C901 llm_config.model.quantize_kv_cache = args.quantize_kv_cache if hasattr(args, "local_global_attention"): llm_config.model.local_global_attention = args.local_global_attention + if hasattr(args, "use_moe_quantized_op"): + llm_config.model.use_moe_quantized_op = args.use_moe_quantized_op # ExportConfig if hasattr(args, "max_seq_length"):