Skip to content

Commit c600e2e

Browse files
MoECalibrationModule bugfix (#2023)
SUMMARY: Without this bugfix, `examples/awq/qwen3_moe_example.py` errors out with stack trace: ``` Replacing MoE modules for calibration: 0%| | 0/48 [00:00<?, ?it/s] Traceback (most recent call last): File "/home/brian-dellabetta/projects/llm-compressor/examples/awq/qwen3_moe_example.py", line 65, in <module> oneshot( File "/home/brian-dellabetta/projects/llm-compressor/src/llmcompressor/entrypoints/oneshot.py", line 348, in oneshot one_shot() File "/home/brian-dellabetta/projects/llm-compressor/src/llmcompressor/entrypoints/oneshot.py", line 172, in __call__ self.apply_recipe_modifiers( File "/home/brian-dellabetta/projects/llm-compressor/src/llmcompressor/entrypoints/oneshot.py", line 216, in apply_recipe_modifiers with moe_calibration_context( File "/home/brian-dellabetta/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/contextlib.py", line 135, in __enter__ return next(self.gen) File "/home/brian-dellabetta/projects/llm-compressor/src/llmcompressor/modeling/moe_context.py", line 106, in moe_calibration_context replacement = MoECalibrationModule.load_from_registry( TypeError: RegistryMixin.load_from_registry() takes 2 positional arguments but 4 were given ``` This resolves the bug (but subsequently hits another error #2022 ) TEST PLAN: "please outline how the changes were tested" Signed-off-by: Brian Dellabetta <[email protected]> Co-authored-by: Dipika Sikka <[email protected]>
1 parent bf4c477 commit c600e2e

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

src/llmcompressor/modeling/moe_context.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,8 @@ def moe_calibration_context(
105105
):
106106
replacement = MoECalibrationModule.load_from_registry(
107107
class_name,
108-
module,
109-
model.config,
108+
original=module,
109+
config=model.config,
110110
calibrate_all_experts=calibrate_all_experts,
111111
)
112112
model.set_submodule(name, replacement)

0 commit comments

Comments
 (0)