[MoE Calibration] Simplify MoE calibration interface

[sairampillai]

Introduce standardized MoE calibration interface and deprecate legacy replace_modules_for_calibration

Summary

Implements a simplified, decorator-based registration system for MoE model calibration using a single MoECalibrationModule base class, making MoE model integration easier and deprecates the legacy replace_modules_for_calibration function.

Problem

MoE model calibration currently requires module replacement logic scattered across replace_modules_for_calibration and manual context management. This makes contributing new MoE model support difficult and error-prone. Additionally, each model required custom replacement functions with duplicated boilerplate code.

Relevant Issues

Fixes #1829

Solution

MoECalibrationModule abstract base class implementation

• Only two required methods: from_original() classmethod and optional restore()
• is_permanent flag to specify if module replacement is to be restored using restore()
• Clear contract: permanent modules stay in calibration form, non-permanent modules get restored after context exit

Decorator-Based Registration: @register_moe_calibration("ModuleName") decorator

• Automatic registration in MOE_CALIBRATION_MODULES registry
• Models self-register when their module is imported

New Model Integration: Adding MoE support requires only:

@register_moe_calibration("YourMoEModule")
class CalibrationYourMoE(MoECalibrationModule):
    is_permanent = True  # or False
    
    @classmethod
    def from_original(cls, original, config, calibrate_all_experts=True):
        return cls(config, original, calibrate_all_experts)

Dataset Arguments: New: moe_calibrate_all_experts: bool = True - Controls whether all experts see all tokens during calibration

• True (default): All experts receive all tokens for proper quantization statistics
• False: Normal routing behavior (only routed experts are used)
• Used by both oneshot() and DatasetArguments
• Automatically passed to moe_calibration_context by pipelines

Automatic Context Management: moe_calibration_context integrated into pipelines

• Wraps calibration automatically in oneshot.py
• Handles module replacement and restoration transparently
• No manual context management required by users

Backward Compatibility: Deprecation of replace_modules_for_calibration with warnings

• Legacy function preserved for compatibility
• Clear migration path documented in deprecation message

Test Plan

• ✅ Unit tests for contextual MoE calibration with automatic module restoration
• ✅ Unit tests for permanent MoE calibration persistence
• ✅ Integration tests with Qwen3, Llama4, and DeepSeek V3 models
• ✅ Verification that all experts receive data during calibration
• ✅ Deprecation warning verification for legacy functions

Testing

• ✅ All unit tests pass
• ✅ Calibration types working correctly
• ✅ Model structure correctly modified and restored inside/outside contexts
• ✅ Linting and type checking pass
• ✅ Backward compatibility verified with deprecation warnings

Migration Guide

Before:

# Required defining MoEModelConfig entries, handling context manually
from llmcompressor.modeling.prepare import replace_modules_for_calibration
model = replace_modules_for_calibration(model, calibrate_all_experts=True)

After:

# Automatic - just use moe_calibration_context
from llmcompressor.modeling import moe_calibration_context

with moe_calibration_context(model, calibrate_all_experts=True):
    # Run calibration - modules replaced automatically
    for batch in dataloader:
        model(**batch)
# Modules restored automatically (if not permanent)

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[sairampillai]

@kylesayrs @dsikka Few clarifications:

• I pushed a couple of commits without signing, how do you suggest I fix that?
• I have deprecated the calibrate_moe_context parameter, do we want to plan how to phase it out?
• I have tested using unit tests but without GPU (gpu poor), can you point me to best ways to test this change end-to-end?

[brian-dellabetta]

@sairampillai , regarding DCO, you can ignore that. We can sign it via github once reviewed/approved

[dsikka]

@kylesayrs

[kylesayrs]

This looks good, but I worry that this implementation uses more abstraction than is necessary. I like the idea of "contextual" vs "permanent" changes, and we should definitely log which one is being used to the user.

Please consider simplifying to a single mapping dictionary, and a single ABC class to handle the from_original and restore functions. Don't be afraid to remove/ refactor existing code!

[kylesayrs]

Hey @sairampillai! Are you still interested in contributing to this PR? If not, please let me know and I can assign someone to pick up where you left off!

[sairampillai]

@kylesayrs I am working on the updates, I will push an update soon for review!

[kylesayrs]

Looks great so far, thanks for following up!

[kylesayrs]

Looks awesome! Is this ready to be tested?

[kylesayrs]

Something like this is also implemented via the RegistryMixin, but we can standardize that in a follow up as well.

[kylesayrs]

Your registry is slightly different, let's leave this for a follow up

[brian-dellabetta]

Thanks for the contribution! I think you can remove the from_original class methods and just use the constructors directly

[brian-dellabetta]

do we want to keep this note for all examples? Might be cleaner without them, what do people think?

[sairampillai]

I felt it helpful to have the note in varied examples since this would be a breaking change once we deprecate older methods. Open to recommendations.

[brian-dellabetta]

Why not just use the constructor? We can probably remove the from_original class method

Suggested change

	return CalibrationDeepseekV3MoE.from_original(
	original=module,
	config=config,
	calibrate_all_experts=calibrate_all_experts,
	return CalibrationDeepseekV3MoE(
	original=module,
	config=config,
	calibrate_all_experts=calibrate_all_experts,

[brian-dellabetta]

Same here, should just be able to use the constructor directly, no?