Skip to content

Commit 9789eea

Browse files
authored
[ExecuTorch][llm] Wire --use_moe_quantized_op through export pipeline and LLM config (pytorch#21540)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: pytorch#21126 by @digantdesai ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/digantdesai/76/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/76/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/digantdesai/75/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/76/orig Differential Revision: [D102381996](https://our.internmc.facebook.com/intern/diff/D102381996/) @diff-train-skip-merge Co-authored-by: Digant Desai <digantdesai@meta.com>
1 parent 607c4c1 commit 9789eea

4 files changed

Lines changed: 143 additions & 0 deletions

File tree

‎examples/models/llama/export_llama_lib.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -565,6 +565,19 @@ def build_args_parser() -> argparse.ArgumentParser:
565565
help="Use SpinQuant for better quantization performance. Only support cuda and native.",
566566
)
567567

568+
parser.add_argument(
569+
"--use_moe_quantized_op",
570+
action="store_true",
571+
default=False,
572+
help=(
573+
"Replace eager MoE feed-forward modules with the "
574+
"`llama::quantized_moe_ffn` portable-runtime custom op (INT4 "
575+
"weights, INT8 dyn-quant activations via torchao). On aarch64 "
576+
"with ENABLE_QUANTIZED_MOE_FFN the optimized torchao kernel is "
577+
"used; otherwise a portable reference fallback runs."
578+
),
579+
)
580+
568581
parser.add_argument(
569582
"-qat",
570583
"--use_qat",
@@ -809,6 +822,7 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager:
809822
use_torchao_kernels_linear=llm_config.backend.torchao.use_torchao_kernels_linear,
810823
use_torchao_kernels_tied_embedding=llm_config.backend.torchao.use_torchao_kernels_tied_embedding,
811824
quantize_with_hqq=llm_config.quantization.use_hqq,
825+
use_moe_quantized_op=llm_config.model.use_moe_quantized_op,
812826
)
813827
)
814828

@@ -1712,6 +1726,7 @@ def _get_source_transforms( # noqa
17121726
use_torchao_kernels_linear: bool = False,
17131727
use_torchao_kernels_tied_embedding: bool = False,
17141728
quantize_with_hqq: bool = True,
1729+
use_moe_quantized_op: bool = False,
17151730
) -> List[Callable[[torch.nn.Module], torch.nn.Module]]:
17161731
"""
17171732
Return a list of functions that transform a graph.
@@ -1769,6 +1784,17 @@ def _get_source_transforms( # noqa
17691784

17701785
transforms.append(inject_fast_hadamard_transform_native_for_spin_quant)
17711786

1787+
if use_moe_quantized_op:
1788+
from .source_transformation.moe import replace_moe_with_quantized_op
1789+
1790+
transforms.append(
1791+
partial(
1792+
replace_moe_with_quantized_op,
1793+
group_size=group_size or 32,
1794+
weight_nbit=4,
1795+
)
1796+
)
1797+
17721798
if embedding_quantize:
17731799
"""
17741800
When this option is selected, it finds all embedding layers and transforms

‎extension/llm/custom_ops/BUCK‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,9 @@ fbcode_target(_kind = runtime.python_test,
138138
preload_deps = [
139139
":custom_ops_aot_lib_mkl_noomp",
140140
":custom_ops_aot_py",
141+
# The Python source transform calls
142+
# `torch.ops.torchao._pack_8bit_act_4bit_weight` to pack each
143+
# expert's INT4 weights, so preload the torchao AOT library too.
141144
"//pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_aten",
142145
],
143146
deps = [
@@ -146,5 +149,7 @@ fbcode_target(_kind = runtime.python_test,
146149
"//executorch/examples/models/llama:llama_transformer",
147150
"//executorch/examples/models/llama:transformer_modules",
148151
"//executorch/examples/models/llama:source_transformation",
152+
"//executorch/examples/models/llama:export_library",
153+
"//executorch/extension/llm/export/config:llm_config",
149154
],
150155
)

‎extension/llm/custom_ops/test_quantized_moe.py‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -596,3 +596,110 @@ def test_no_shared_expert_is_none(self) -> None:
596596
wrapper.m = moe
597597
replace_moe_with_quantized_op(wrapper, group_size=32, weight_nbit=4)
598598
self.assertIsNone(wrapper.m.shared_expert)
599+
600+
601+
class TestExportPipelineWiring(unittest.TestCase):
602+
"""The export pipeline correctly includes the MoE transform."""
603+
604+
def test_get_source_transforms_includes_moe_when_enabled(self) -> None:
605+
from functools import partial
606+
607+
from executorch.examples.models.llama.export_llama_lib import (
608+
_get_source_transforms,
609+
)
610+
611+
transforms = _get_source_transforms(
612+
dtype_override=torch.float32, use_moe_quantized_op=True
613+
)
614+
moe_transforms = [
615+
t
616+
for t in transforms
617+
if isinstance(t, partial)
618+
and t.func.__name__ == "replace_moe_with_quantized_op"
619+
]
620+
self.assertEqual(len(moe_transforms), 1)
621+
self.assertEqual(moe_transforms[0].keywords["group_size"], 32)
622+
self.assertEqual(moe_transforms[0].keywords["weight_nbit"], 4)
623+
624+
def test_get_source_transforms_excludes_moe_when_disabled(self) -> None:
625+
from functools import partial
626+
627+
from executorch.examples.models.llama.export_llama_lib import (
628+
_get_source_transforms,
629+
)
630+
631+
transforms = _get_source_transforms(
632+
dtype_override=torch.float32, use_moe_quantized_op=False
633+
)
634+
moe_transforms = [
635+
t
636+
for t in transforms
637+
if isinstance(t, partial)
638+
and hasattr(t.func, "__name__")
639+
and t.func.__name__ == "replace_moe_with_quantized_op"
640+
]
641+
self.assertEqual(len(moe_transforms), 0)
642+
643+
def test_get_source_transforms_passes_custom_group_size(self) -> None:
644+
from functools import partial
645+
646+
from executorch.examples.models.llama.export_llama_lib import (
647+
_get_source_transforms,
648+
)
649+
650+
transforms = _get_source_transforms(
651+
dtype_override=torch.float32,
652+
use_moe_quantized_op=True,
653+
group_size=64,
654+
)
655+
moe_transforms = [
656+
t
657+
for t in transforms
658+
if isinstance(t, partial)
659+
and t.func.__name__ == "replace_moe_with_quantized_op"
660+
]
661+
self.assertEqual(len(moe_transforms), 1)
662+
self.assertEqual(moe_transforms[0].keywords["group_size"], 64)
663+
664+
def test_sentinel_op_is_registered(self) -> None:
665+
self.assertTrue(hasattr(torch.ops.llama, "_quantized_moe_ffn_active"))
666+
self.assertTrue(torch.ops.llama._quantized_moe_ffn_active())
667+
668+
669+
class TestLlmConfigMoeFlag(unittest.TestCase):
670+
"""llm_config wires the --use_moe_quantized_op flag correctly."""
671+
672+
def test_from_args_sets_flag(self) -> None:
673+
import argparse
674+
675+
from executorch.extension.llm.export.config.llm_config import LlmConfig
676+
677+
args = argparse.Namespace(use_moe_quantized_op=True)
678+
config = LlmConfig.from_args(args)
679+
self.assertTrue(config.model.use_moe_quantized_op)
680+
681+
def test_default_is_false(self) -> None:
682+
from executorch.extension.llm.export.config.llm_config import ModelConfig
683+
684+
self.assertFalse(ModelConfig().use_moe_quantized_op)
685+
686+
def test_from_args_missing_field_defaults_false(self) -> None:
687+
import argparse
688+
689+
from executorch.extension.llm.export.config.llm_config import LlmConfig
690+
691+
args = argparse.Namespace()
692+
config = LlmConfig.from_args(args)
693+
self.assertFalse(config.model.use_moe_quantized_op)
694+
695+
def test_argparser_flag_true(self) -> None:
696+
from executorch.examples.models.llama.export_llama_lib import build_args_parser
697+
698+
args = build_args_parser().parse_args(["--use_moe_quantized_op"])
699+
self.assertTrue(args.use_moe_quantized_op)
700+
701+
def test_argparser_flag_default_false(self) -> None:
702+
from executorch.examples.models.llama.export_llama_lib import build_args_parser
703+
704+
args = build_args_parser().parse_args([])
705+
self.assertFalse(args.use_moe_quantized_op)

‎extension/llm/export/config/llm_config.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,9 @@ class ModelConfig:
205205
use_kv_cache: bool = False
206206
quantize_kv_cache: bool = False
207207
local_global_attention: Optional[List[int]] = None
208+
# Replace eager MOEFeedForward modules with the
209+
# `llama::quantized_moe_ffn` portable-runtime custom op.
210+
use_moe_quantized_op: bool = False
208211

209212
def __post_init__(self):
210213
self._validate_attention_sink()
@@ -728,6 +731,8 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig": # noqa: C901
728731
llm_config.model.quantize_kv_cache = args.quantize_kv_cache
729732
if hasattr(args, "local_global_attention"):
730733
llm_config.model.local_global_attention = args.local_global_attention
734+
if hasattr(args, "use_moe_quantized_op"):
735+
llm_config.model.use_moe_quantized_op = args.use_moe_quantized_op
731736

732737
# ExportConfig
733738
if hasattr(args, "max_seq_length"):

0 commit comments

Comments
 (0)