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