Design for refactoring module replacement

[yiliu30]

Design for #899

• Requsting developers provide sufficient metadata for replacement
• Supporting in-tree and out-of-tree module registration

Usage

from auto_round.modelling.replace_modules import ReplacementModuleBase

class SequentialGPTOSSMoE(ReplacementModuleBase):
    """
    Replaces GPT-OSS fused-expert MoE with per-expert `GPTOssSingleExpert` modules.
    Copies weights from fused tensors and reuses the original router and optional shared_expert.
    """
    def __init__(self, original: GptOssMLP, config: GptOssConfig):
        super().__init__()
        ....


    @classmethod
    def original_module_class(cls) -> str:
        """Return the class name of the module this replaces."""
        return "GptOssMLP"

    @classmethod
    def from_original(
        cls,
        original: GptOssMLP,
        config: GptOssConfig,
        **kwargs,
    ) -> "SequentialGPTOSSMoE":
        """Create an instance from the original module."""
        return cls(original, config)

• Impl

import torch
from abc import ABC
from typing import Dict, Type
from tqdm import tqdm
from transformers import PreTrainedModel
from loguru import logger
from abc import abstractmethod

class ReplacementModuleBase(ABC, torch.nn.Module):
    """
    Abstract base class for module replacement during calibration phase.
    
    Replacement modules replace original modules to ensure all components
    receive data for proper quantization statistics.
    
    Subclasses must:
    1. Implement `original_module_class()` to return the target module class name
    2. Implement `__init__()` with signature:
       (self, original, config, calibrate_all_experts=True)
    """
    
    # Registry: module class name -> replacement module class
    _replacement_registry: Dict[str, Type['ReplacementModuleBase']] = {}
    
    
    def __init_subclass__(cls, **kwargs):
        """Automatically register subclasses in the replacement registry."""
        super().__init_subclass__(**kwargs)
        
        # Only register if it's a concrete implementation (not ABC)
        if not getattr(cls, '__abstractmethods__', None):
            if cls.original_module_class() is None:
                raise TypeError(
                    f"{cls.__name__} must implement 'original_module_class()' class method "
                    "to return the name of the module class it replaces"
                )
            
            if cls.original_module_class() in cls._replacement_registry:
                existing = cls._replacement_registry[cls.original_module_class()]
                raise ValueError(
                    f"Module '{cls.original_module_class()}' already registered to "
                    f"{existing.__name__}. Cannot register {cls.__name__}."
                )
            
            cls._replacement_registry[cls.original_module_class()] = cls
            print(f"Registered {cls.__name__} for replacing {cls.original_module_class()}")
    
    @classmethod
    def get_replacement_class(cls, module_class_name: str) -> Type['ReplacementModuleBase']:
        """Get replacement class for a given module class name."""
        return cls._replacement_registry.get(module_class_name)
    
    @classmethod
    def is_registered(cls, module_class_name: str) -> bool:
        """Check if a module class has a replacement implementation."""
        return module_class_name in cls._replacement_registry
    
    @classmethod
    def get_registered_modules(cls) -> list:
        """Get list of all registered module class names."""
        return list(cls._replacement_registry.keys())
    
    @classmethod
    @abstractmethod
    def original_module_class(cls) -> str:
        """Return the class name of the module this replaces."""
        pass
    
    @classmethod
    @abstractmethod
    def from_original(
        cls,
        original: torch.nn.Module,
        config,
        # calibrate_all_experts: bool = Tr
    ) -> 'ReplacementModuleBase':
        """Create replacement module from original module."""
        pass



def apply_replacements(
    model: PreTrainedModel,
) -> dict:
    """
    Function to apply module replacements to a model.

    This scans all modules in the model and replaces any registered modules with their
    replacement equivalents. Non-permanent modules are tracked for later restoration.

    The model is modified in-place, so the same model object should be used.

    Args:
        model: The model to apply module replacement to (modified in-place).

    Returns:
        dict: A dictionary of replaced modules, where the keys are module names and
              the values are tuples of (original_module, replacement_module).
    """
    replaced = {}

    # Step 1: Collect all modules that need replacement
    logger.info("Scanning for modules to replace")
    modules_to_replace = []
    for name, module in model.named_modules():
        class_name = (
            module.original_module_class()
            if hasattr(module, "original_module_class") and callable(module.original_module_class)
            else module.__class__.__name__
        )
        if ReplacementModuleBase.is_registered(class_name):
            modules_to_replace.append((name, module, class_name))

    # Step 2: Replace modules
    if modules_to_replace:
        logger.info(f"Found {len(modules_to_replace)} modules to replace")
        for name, module, class_name in tqdm(
            modules_to_replace, desc="Replacing modules"
        ):
            replacement_cls = ReplacementModuleBase.get_replacement_class(class_name)
            replacement = replacement_cls(
                module,
                model.config,
            )
            model.set_submodule(name, replacement)
            replaced[name] = (module, replacement)

    # Log what was replaced
    if replaced:
        logger.info(f"Replaced {len(replaced)} modules")


    return model

• The idea was adapted from llm-compressor. https://github.com/vllm-project/llm-compressor/blob/main/src/llmcompressor/modeling/moe_context.py

@wenhuach21 @n1ck-guo Please help review that design, thx.
cc @thuang6

[yiliu30]

Take refactoring gpt-oss replacement as an example a4a4118

[thuang6]

Need is somehow general, especially for MoE related models. let's include more people in review @wenhuach21, @n1ck-guo, @WeiweiZhang1, @mengniwang95, @xinhe3. Pls give your feedback so that we can move to next step, like raising PR

[wenhuach21]

Consider adding a function to verify the model's configuration or related information. Relying solely on the module name may not provide sufficient robustness.

[wenhuach21]

We also need to consider how to apply this in AutoRound inference on the Transformers backend, where the replacement must occur before loading the weights. But we could assume that the wights is aligned to the replaced module

[yiliu30]

I see, the replacement shouldn’t break the loading behavior here #1067 (comment)

[xin3he]

Looks good to me.

[n1ck-guo]

LGTM

[yiliu30]

#1153