diff --git a/src/llmcompressor/modeling/moe/fp8_experts.py b/src/llmcompressor/modeling/moe/fp8_experts.py new file mode 100644 index 0000000000..4564bb2176 --- /dev/null +++ b/src/llmcompressor/modeling/moe/fp8_experts.py @@ -0,0 +1,211 @@ +"""REAP support for Transformers native fine-grained FP8 experts.""" + +import torch +import torch.nn as nn +from compressed_tensors import align_module_device +from transformers.integrations.finegrained_fp8 import FP8Experts + +from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag + + +class FP8PrunableExperts(FP8Experts): + """A REAP-aware view of a Transformers :class:`FP8Experts` module. + + Existing FP8 modules are changed to this class in place. This preserves + their packed parameters, offload state, and checkpoint key names. The stock + Transformers forward remains active except while REAP norm collection is + enabled, when an eager path records each unweighted expert output norm. + """ + + _EXPERT_TENSOR_NAMES = ( + "gate_up_proj", + "gate_up_proj_scale_inv", + "up_proj", + "up_proj_scale_inv", + "down_proj", + "down_proj_scale_inv", + "gate_up_proj_activation_scale", + "down_proj_activation_scale", + ) + num_experts: int + _reap_norms: dict[int, torch.Tensor] | None + + def start_reap_norm_collection(self) -> None: + self._reap_norms: dict[int, torch.Tensor] | None = {} + + def stop_reap_norm_collection(self) -> None: + self._reap_norms = None + + def take_reap_norms(self) -> dict[int, torch.Tensor]: + norms = getattr(self, "_reap_norms", None) + if norms is None: + return {} + + self._reap_norms = {} + return norms + + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + if getattr(self, "_reap_norms", None) is None: + return super().forward(hidden_states, top_k_index, top_k_weights) + + # Match Transformers FP8Experts eager math. Fused expert kernels cannot + # expose the unweighted per-expert outputs needed by REAP. + final_hidden_states = torch.zeros_like(hidden_states, dtype=torch.float32) + calibrate_all = get_calibrate_all_experts_flag() + + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot( + top_k_index, num_classes=self.num_experts + 1 + ) + expert_mask = expert_mask.permute(2, 1, 0) + if calibrate_all: + expert_indices = range(self.num_experts) + else: + expert_indices = ( + torch.greater(expert_mask.sum(dim=(-1, -2)), 0) + .nonzero(as_tuple=False) + .view(-1) + ) + + for expert_index in expert_indices: + if int(expert_index) == self.num_experts: + continue + + top_k_pos, token_indices = torch.where(expert_mask[expert_index]) + current_state = ( + hidden_states if calibrate_all else hidden_states[token_indices] + ) + gate_up_activation_scale = ( + self.gate_up_proj_activation_scale[expert_index] + if self.activation_scheme == "static" + else None + ) + projected = self.linear( + current_state, + self.gate_up_proj[expert_index] + if self.has_gate + else self.up_proj[expert_index], + self.gate_up_proj_scale_inv[expert_index] + if self.has_gate + else self.up_proj_scale_inv[expert_index], + activation_scale=gate_up_activation_scale, + ) + projected = ( + self._apply_gate(projected) if self.has_gate else self.act_fn(projected) + ) + down_activation_scale = ( + self.down_proj_activation_scale[expert_index] + if self.activation_scheme == "static" + else None + ) + expert_output = self.linear( + projected, + self.down_proj[expert_index], + self.down_proj_scale_inv[expert_index], + activation_scale=down_activation_scale, + ) + + norms = self._reap_norms + assert norms is not None + with torch.no_grad(): + norms[int(expert_index)] = torch.linalg.norm( + expert_output.float(), dim=-1 + ).reshape(-1) + + if calibrate_all: + expert_output = expert_output[token_indices] + routing_weights = top_k_weights[token_indices, top_k_pos, None] + weighted_output = expert_output * routing_weights.to(expert_output.dtype) + final_hidden_states.index_add_( + 0, + token_indices, + weighted_output.to(final_hidden_states.dtype), + ) + + return final_hidden_states.to(hidden_states.dtype) + + def prune_experts_(self, retained: list[int]) -> None: + """Slice packed FP8 weights and scales along their expert dimension.""" + weight = self.gate_up_proj if self.has_gate else self.up_proj + if weight.dtype != torch.float8_e4m3fn: + raise TypeError( + "REAP native packed expert pruning only supports e4m3 FP8 " + f"weights, got {weight.dtype}" + ) + if not retained: + raise ValueError("REAP must retain at least one expert") + if len(set(retained)) != len(retained): + raise ValueError("REAP retained expert indices must be unique") + if min(retained) < 0 or max(retained) >= self.num_experts: + raise IndexError( + f"REAP retained expert indices must be in [0, {self.num_experts})" + ) + + retained_indices = torch.tensor(retained, dtype=torch.long) + sliced: dict[str, tuple[torch.Tensor, bool | None]] = {} + with align_module_device(self): + for name in self._EXPERT_TENSOR_NAMES: + tensor = getattr(self, name, None) + if tensor is None: + continue + if tensor.ndim == 0 or tensor.shape[0] != self.num_experts: + raise ValueError( + f"Cannot REAP-prune FP8 expert tensor {name}: expected " + f"leading dimension {self.num_experts}, got " + f"{tuple(tensor.shape)}" + ) + indices = retained_indices.to(tensor.device) + is_parameter = name in self._parameters + is_buffer = name in self._buffers + if not is_parameter and not is_buffer: + raise TypeError( + f"Cannot REAP-prune FP8 expert tensor {name}: expected a " + "registered parameter or buffer" + ) + sliced[name] = ( + tensor.detach()[indices].contiguous(), + tensor.requires_grad if is_parameter else None, + ) + + for name, (tensor, requires_grad) in sliced.items(): + if requires_grad is not None: + setattr( + self, + name, + nn.Parameter(tensor, requires_grad=requires_grad), + ) + else: + setattr(self, name, tensor) + + self.num_experts = len(retained) + + +def make_fp8_experts_reap_prunable(module: nn.Module) -> nn.Module: + """Adapt a loaded Transformers FP8 expert module without copying tensors.""" + if isinstance(module, FP8Experts) and not isinstance(module, FP8PrunableExperts): + weight = module.gate_up_proj if module.has_gate else module.up_proj + if weight.dtype != torch.float8_e4m3fn: + raise TypeError( + "REAP native packed expert pruning only supports e4m3 FP8 " + f"weights, got {weight.dtype}" + ) + module.__class__ = FP8PrunableExperts + + # Both Accelerate and compressed-tensors wrap ``forward`` by saving the + # method that was active when offloading was installed. REAP adaptation + # commonly happens after that installation, so changing ``__class__`` + # alone would leave the wrapper calling the stock FP8Experts forward and + # norm collection would silently remain empty. Refresh only known + # wrapper slots; preserve any unrelated custom forward implementation. + old_forward = getattr(module, "_old_forward", None) + if getattr(old_forward, "__func__", None) is FP8Experts.forward: + module._old_forward = FP8PrunableExperts.forward.__get__(module) + + if getattr(module, "_original_forward_func", None) is FP8Experts.forward: + module._original_forward_func = FP8PrunableExperts.forward + return module diff --git a/src/llmcompressor/modeling/moe/helpers.py b/src/llmcompressor/modeling/moe/helpers.py index 41c21f56db..045d1e7bf3 100644 --- a/src/llmcompressor/modeling/moe/helpers.py +++ b/src/llmcompressor/modeling/moe/helpers.py @@ -46,6 +46,39 @@ def __validate__(cls, object: object) -> bool: ) +class ReapPrunableExpertsProtocol(TorchModuleProtocol): + """Expert container operations required by REAP. + + Implementations own norm collection because an expert container may store + experts either as child modules or as packed tensors. + """ + + num_experts: int + + def start_reap_norm_collection(self) -> None: + raise NotImplementedError() + + def stop_reap_norm_collection(self) -> None: + raise NotImplementedError() + + def take_reap_norms(self) -> dict[int, torch.Tensor]: + raise NotImplementedError() + + def prune_experts_(self, retained: list[int]) -> None: + raise NotImplementedError() + + @classmethod + def __validate__(cls, object: object) -> bool: + return ( + isinstance(getattr(object, "num_experts", None), int) + and callable(getattr(object, "forward", None)) + and callable(getattr(object, "start_reap_norm_collection", None)) + and callable(getattr(object, "stop_reap_norm_collection", None)) + and callable(getattr(object, "take_reap_norms", None)) + and callable(getattr(object, "prune_experts_", None)) + ) + + def get_use_experts_implementation_args(experts_cls: type) -> dict[str, bool] | None: """ Get the keyword arguments to the `@use_experts_implementation` decorator which diff --git a/src/llmcompressor/modifiers/pruning/reap/base.py b/src/llmcompressor/modifiers/pruning/reap/base.py index 6d93960f2c..a6290d9314 100644 --- a/src/llmcompressor/modifiers/pruning/reap/base.py +++ b/src/llmcompressor/modifiers/pruning/reap/base.py @@ -14,7 +14,8 @@ from llmcompressor.core import Event, State from llmcompressor.core.session_functions import active_session from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag -from llmcompressor.modeling.moe.linear_experts import ExpertMLP +from llmcompressor.modeling.moe.helpers import ReapPrunableExpertsProtocol +from llmcompressor.modeling.moe.linear_experts import ExpertMLP, LinearExperts2D from llmcompressor.modifiers import Modifier from llmcompressor.modifiers.pruning.reap.utils import ( MoeModelAttrs, @@ -65,6 +66,7 @@ class REAPPruningModifier(Modifier): ) _n_experts_to_drop: int = PrivateAttr(default=0) _n_experts_to_drop_per_group: int | None = PrivateAttr(default=None) + _pruning_complete: bool = PrivateAttr(default=False) _norm_buffers: dict[str, dict[int, torch.Tensor]] = PrivateAttr( default_factory=dict ) @@ -153,6 +155,7 @@ def on_initialize(self, state: State, **kwargs) -> bool: def on_calibration_start(self, state: State, event: Event, **kwargs): model = state.model + self._pruning_complete = False # Ensure that REAP is the only modifier for this calibration pass session = active_session() @@ -180,35 +183,56 @@ def on_calibration_start(self, state: State, event: Event, **kwargs): self._saliency_trackers[layer_name] = REAPSaliencyTracker( self._moe_attrs.num_experts ) - self._norm_buffers[layer_name] = {} - - # One hook per expert to record its per-token output norm and - # store in the layer's norm buffer experts = getattr(module, self._moe_attrs.experts_attr) - expert_list = [ - expert for expert in experts.children() if isinstance(expert, ExpertMLP) - ] # Filter out the activation function submodule - for idx, expert in enumerate(expert_list): - self.register_hook( - expert, partial(self._expert_hook, layer_name, idx), "forward" + if isinstance(experts, LinearExperts2D): + self._norm_buffers[layer_name] = {} + expert_list = [ + expert + for expert in experts.children() + if isinstance(expert, ExpertMLP) + ] + for index, expert in enumerate(expert_list): + self.register_hook( + expert, + partial(self._expert_hook, layer_name, index), + "forward", + ) + elif isinstance(experts, ReapPrunableExpertsProtocol): + experts.start_reap_norm_collection() + else: + raise TypeError( + f"Experts at {layer_name} are neither LinearExperts2D nor a " + "REAP-prunable packed container" ) # Hook for the experts block to capture the router's top-k - # routing decisions and weights. This hook also executes pruning, - # which requires the expert output norms. Therefore, it must be a - # forward hook, so that it runs after the individual expert hooks - # have populated the norm buffer + # routing decisions and weights after the container has recorded + # the corresponding unweighted expert output norms. self.register_hook( experts, partial(self._experts_block_hook, layer_name), "forward" ) def on_calibration_end(self, state: State, event: Event, **kwargs): self.remove_hooks() + for layer_name in self._moe_attrs.moe_layer_names: + module = state.model.get_submodule(layer_name) + experts = getattr(module, self._moe_attrs.experts_attr) + if isinstance(experts, ReapPrunableExpertsProtocol): + experts.stop_reap_norm_collection() + + self._ensure_all_layers_pruned() + self._pruning_complete = True def on_finalize(self, state: State, **kwargs) -> bool: """Finalize the model config to reflect the new number of experts.""" model = state.model + if not self._pruning_complete: + raise RuntimeError( + "REAP calibration did not complete; refusing to update the model " + "config" + ) + self._ensure_all_layers_pruned() new_num_experts = self._moe_attrs.num_experts - self._n_experts_to_drop update_model_config(model, self._moe_attrs, new_num_experts) @@ -218,6 +242,19 @@ def on_finalize(self, state: State, **kwargs) -> bool: return True + def _ensure_all_layers_pruned(self) -> None: + if not self._saliency_trackers: + return + + unpruned = sorted(self._saliency_trackers) + preview = ", ".join(unpruned[:3]) + if len(unpruned) > 3: + preview += f", ... ({len(unpruned)} total)" + raise RuntimeError( + "REAP did not collect saliency and structurally prune every MoE " + f"layer; refusing to update the model config. Unpruned: {preview}" + ) + # -- decision finalization ---------------------------------------------- def on_sequential_epoch_end(self, state: State, event: Event, **kwargs): @@ -240,11 +277,15 @@ def on_sequential_epoch_end(self, state: State, event: Event, **kwargs): len(retained) == expected ), f"Expected {expected} retained experts, got {len(retained)}" + moe_block = model.get_submodule(layer_name) + experts = getattr(moe_block, self._moe_attrs.experts_attr) prune_moe_layer(model, layer_name, retained, self._moe_attrs) # free this layer's accumulators / buffers now del self._saliency_trackers[layer_name] self._norm_buffers.pop(layer_name, None) + if isinstance(experts, ReapPrunableExpertsProtocol): + experts.stop_reap_norm_collection() # -- calibration hooks --------------------------------------------------- @@ -281,9 +322,13 @@ def _experts_block_hook( if tracker is None: return - norm_buffer = self._norm_buffers[layer_name] + if isinstance(module, ReapPrunableExpertsProtocol): + expert_norms = module.take_reap_norms() + else: + expert_norms = self._norm_buffers[layer_name] with torch.no_grad(): - tracker.update(top_k_indices, top_k_weights, norm_buffer) + tracker.update(top_k_indices, top_k_weights, expert_norms) - self._norm_buffers[layer_name] = {} + if isinstance(module, LinearExperts2D): + self._norm_buffers[layer_name] = {} diff --git a/src/llmcompressor/modifiers/pruning/reap/utils.py b/src/llmcompressor/modifiers/pruning/reap/utils.py index c46baecc1c..c98bd120ea 100644 --- a/src/llmcompressor/modifiers/pruning/reap/utils.py +++ b/src/llmcompressor/modifiers/pruning/reap/utils.py @@ -13,7 +13,9 @@ from loguru import logger from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag +from llmcompressor.modeling.moe.fp8_experts import make_fp8_experts_reap_prunable from llmcompressor.modeling.moe.granitemoe import GraniteMoeLinearExperts +from llmcompressor.modeling.moe.helpers import ReapPrunableExpertsProtocol from llmcompressor.modeling.moe.linear_experts import ExpertMLP, LinearExperts2D from llmcompressor.modeling.moe.llama4 import Llama4LinearExperts @@ -42,11 +44,21 @@ class MoeModelAttrs: ROUTER_ATTRS = ["router", "gate"] EXPERTS_ATTRS = ["experts"] -NUM_EXPERTS_CONFIG_KEYS = ["num_experts", "num_local_experts", "moe_num_experts"] +NUM_EXPERTS_CONFIG_KEYS = [ + "n_routed_experts", + "num_experts", + "num_local_experts", + "moe_num_experts", +] TOP_K_CONFIG_KEYS = ["num_experts_per_tok", "top_k", "moe_top_k"] N_GROUP_CONFIG_KEYS = ["n_group"] TOP_K_GROUP_CONFIG_KEYS = ["topk_group", "top_k_group"] -NUM_EXPERTS_MODULE_KEYS = ["num_experts", "n_experts", "n_routed_experts"] +NUM_EXPERTS_MODULE_KEYS = [ + "num_experts", + "n_experts", + "n_routed_experts", + "num_local_experts", +] def get_moe_attrs(model: nn.Module, ignore: list[str]) -> MoeModelAttrs | None: @@ -148,13 +160,11 @@ def get_moe_attrs(model: nn.Module, ignore: list[str]) -> MoeModelAttrs | None: if any(re.search(pattern, name) for pattern in ignore): continue experts = getattr(module, experts_attr) - # REAP currently only supports LinearExperts2D experts, as they receive the - # top_k indices and weights from the router in their forward pass. - # Granite and Llama4 experts diverge from this behavior, so they are - # unsupported for now. - if not isinstance(experts, LinearExperts2D): + experts = make_fp8_experts_reap_prunable(experts) + if not isinstance(experts, (LinearExperts2D, ReapPrunableExpertsProtocol)): logger.warning( - f"Skipping layer {name}: experts module is not LinearExperts2D" + f"Skipping layer {name}: experts module is neither " + "LinearExperts2D nor a REAP-prunable packed container" ) continue if isinstance(experts, GraniteMoeLinearExperts): @@ -171,18 +181,15 @@ def get_moe_attrs(model: nn.Module, ignore: list[str]) -> MoeModelAttrs | None: if not moe_layer_names: raise ValueError( - "Could not find any supported MoE layers with experts in " - "LinearExperts2D format. Make sure the model has MoE layers " + "Could not find any supported MoE layers with REAP-prunable experts. " + "Make sure the model has MoE layers " "(excluding GraniteMoeLinearExperts and Llama4LinearExperts), " "and that the name of its experts module is in EXPERTS_ATTRS " "and it the name of its router module is in ROUTER_ATTRS in " "reap/utils.py" ) - logger.info( - f"Found {len(moe_layer_names)} MoE layers with experts in " - "LinearExperts2D format" - ) + logger.info(f"Found {len(moe_layer_names)} MoE layers with REAP-prunable experts") return MoeModelAttrs( num_experts_config_key=num_experts_config_key, @@ -354,26 +361,34 @@ def prune_moe_layer( router = getattr(moe_block, moe_attrs.router_attr) experts = getattr(moe_block, moe_attrs.experts_attr) - # Preserve non-expert modules (e.g., act_fn in LinearExperts2D) - # These are modules that are not instances of ExpertMLP subclasses - non_expert_modules = {} - for key, module in experts._modules.items(): - if not isinstance(module, ExpertMLP): - non_expert_modules[key] = module - - # Rebuild with retained experts - new_modules = OrderedDict( - ((str(i), experts[pos]) for i, pos in enumerate(retained)) - ) - - # Re-add non-expert modules - new_modules.update(non_expert_modules) - - experts._modules = new_modules - experts.num_experts = len(retained) + if isinstance(experts, ReapPrunableExpertsProtocol): + experts.prune_experts_(retained) + elif isinstance(experts, LinearExperts2D): + # Preserve non-expert modules (e.g., act_fn in LinearExperts2D) + non_expert_modules = { + key: module + for key, module in experts._modules.items() + if not isinstance(module, ExpertMLP) + } + new_modules = OrderedDict( + (str(index), experts[position]) for index, position in enumerate(retained) + ) + new_modules.update(non_expert_modules) + + experts._modules = new_modules + experts.num_experts = len(retained) + else: + raise TypeError( + f"Experts at {layer_name} are neither LinearExperts2D nor a " + "REAP-prunable packed container" + ) _prune_router(router, retained) + # Some architectures keep per-expert router state on the enclosing MoE + # block. Slice it in lockstep with the router logits when present. + _prune_expert_tensor(moe_block, "e_score_correction_bias", retained) + # Update num_experts for any other modules in the layer that may track it for holder in (moe_block, router): for key in NUM_EXPERTS_MODULE_KEYS: @@ -389,31 +404,45 @@ def _prune_router(router: nn.Module, retained: list[int]): with align_module_device(router): retained_t = retained_t.to(router.weight.device) new_weight = router.weight.detach()[retained_t].contiguous() + weight_requires_grad = router.weight.requires_grad new_bias = None + bias_requires_grad = False if getattr(router, "bias", None) is not None: new_bias = router.bias.detach()[retained_t].contiguous() - # group-limited routers (DeepSeek-V3 / GLM4 / GLM-DSA) carry a per-expert - # score-correction bias buffer that must be shrunk in lockstep - correction = getattr(router, "e_score_correction_bias", None) - new_correction = ( - correction.detach()[retained_t].contiguous() - if correction is not None - else None - ) + bias_requires_grad = router.bias.requires_grad # Direct attribute assignment replaces a parameter/buffer with a different # shape and is correct for both offloaded modules (routed through the # OffloadCache, which re-offloads the new shape) and ordinary modules. - router.weight = nn.Parameter(new_weight, requires_grad=router.weight.requires_grad) + router.weight = nn.Parameter(new_weight, requires_grad=weight_requires_grad) if new_bias is not None: - router.bias = nn.Parameter(new_bias, requires_grad=router.bias.requires_grad) - if new_correction is not None: - router.e_score_correction_bias = new_correction + router.bias = nn.Parameter(new_bias, requires_grad=bias_requires_grad) + + _prune_expert_tensor(router, "e_score_correction_bias", retained) if isinstance(getattr(router, "out_features", None), int): router.out_features = len(retained) +def _prune_expert_tensor(holder: nn.Module, name: str, retained: list[int]) -> None: + is_parameter = name in holder._parameters + is_buffer = name in holder._buffers + if not is_parameter and not is_buffer and not hasattr(holder, name): + return + + retained_t = torch.tensor(retained, dtype=torch.long) + with align_module_device(holder): + value = getattr(holder, name) + retained_t = retained_t.to(value.device) + pruned = value.detach()[retained_t].contiguous() + requires_grad = value.requires_grad + + if is_parameter: + setattr(holder, name, nn.Parameter(pruned, requires_grad=requires_grad)) + else: + setattr(holder, name, pruned) + + def update_model_config( model: nn.Module, moe_attrs: MoeModelAttrs, diff --git a/tests/llmcompressor/modifiers/pruning/reap/test_base.py b/tests/llmcompressor/modifiers/pruning/reap/test_base.py index fdf25f0c18..6696081c24 100644 --- a/tests/llmcompressor/modifiers/pruning/reap/test_base.py +++ b/tests/llmcompressor/modifiers/pruning/reap/test_base.py @@ -743,6 +743,29 @@ def test_reap_initialization_zero_drop(): modifier.initialize(state) +@pytest.mark.unit +def test_reap_rejects_unpruned_calibration_and_finalize(): + config = FakeMoEConfig(num_experts=8, num_hidden_layers=2) + model = FakeMoEModel(config) + modifier = REAPPruningModifier(sparsity=0.5) + state = _make_state(model) + + modifier.initialize(state) + with pytest.raises(RuntimeError, match="calibration did not complete"): + modifier.finalize(state) + + modifier.update_event(state, Event(type_=EventType.CALIBRATION_START)) + + with pytest.raises(RuntimeError, match="Unpruned: layers.0, layers.1"): + modifier.update_event(state, Event(type_=EventType.CALIBRATION_END)) + with pytest.raises(RuntimeError, match="refusing to update the model config"): + modifier.finalize(state) + + assert model.config.num_experts == 8 + assert not modifier.ended_ + assert not modifier.finalized_ + + @pytest.mark.unit def test_reap_full_lifecycle(): """End-to-end test: initialize, calibrate, finalize, verify pruning. diff --git a/tests/llmcompressor/modifiers/pruning/reap/test_fp8.py b/tests/llmcompressor/modifiers/pruning/reap/test_fp8.py new file mode 100644 index 0000000000..fd8ffd206b --- /dev/null +++ b/tests/llmcompressor/modifiers/pruning/reap/test_fp8.py @@ -0,0 +1,416 @@ +import json +import types + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate.hooks import AlignDevicesHook, add_hook_to_module +from compressed_tensors.offload.module import offload_module +from safetensors.torch import load_file +from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM +from transformers.integrations.finegrained_fp8 import FP8Experts +from transformers.utils.quantization_config import FineGrainedFP8Config + +from llmcompressor.core import Event, EventType, State +from llmcompressor.modeling.moe.fp8_experts import ( + FP8PrunableExperts, + make_fp8_experts_reap_prunable, +) +from llmcompressor.modeling.moe.helpers import ReapPrunableExpertsProtocol +from llmcompressor.modifiers.pruning.reap import REAPPruningModifier +from llmcompressor.modifiers.pruning.reap.utils import ( + REAPSaliencyTracker, + get_moe_attrs, + prune_moe_layer, + update_model_config, +) + + +def _config(num_experts: int = 4) -> GlmMoeDsaConfig: + return GlmMoeDsaConfig( + hidden_size=4, + moe_intermediate_size=2, + n_routed_experts=num_experts, + num_experts_per_tok=2, + n_group=1, + topk_group=1, + ) + + +def _tiny_glm_config(num_experts: int = 4) -> GlmMoeDsaConfig: + config = GlmMoeDsaConfig( + vocab_size=16, + hidden_size=8, + intermediate_size=16, + moe_intermediate_size=4, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + n_routed_experts=num_experts, + num_experts_per_tok=2, + n_group=1, + topk_group=1, + first_k_dense_replace=0, + indexer_types=["full"], + index_n_heads=2, + index_head_dim=4, + index_topk=2, + kv_lora_rank=4, + q_lora_rank=4, + qk_rope_head_dim=2, + qk_nope_head_dim=2, + v_head_dim=2, + max_position_embeddings=32, + ) + config.quantization_config = FineGrainedFP8Config( + weight_block_size=(2, 2), + activation_scheme="dynamic", + modules_to_not_convert=["self_attn", "shared_experts", "lm_head"], + ) + return config + + +def _cpu_linear( + self, + input: torch.Tensor, + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + activation_scale: torch.Tensor | None = None, +) -> torch.Tensor: + # CPU test substitute for the CUDA FP8 kernel. The expert forward and + # packed FP8 storage remain the same as the Transformers implementation. + del weight_scale_inv, activation_scale + return F.linear(input, weight.float()) + + +def _fp8_experts( + config: GlmMoeDsaConfig, + activation_scheme: str = "dynamic", + has_gate: bool = True, +) -> FP8Experts: + experts = FP8Experts( + config, + block_size=(2, 2), + activation_scheme=activation_scheme, + has_gate=has_gate, + ) + with torch.no_grad(): + torch.manual_seed(0) + input_projection = experts.gate_up_proj if has_gate else experts.up_proj + input_projection.copy_( + torch.randn(input_projection.shape).to(input_projection.dtype) + ) + experts.down_proj.copy_( + torch.randn(experts.down_proj.shape).to(experts.down_proj.dtype) + ) + input_scale = ( + experts.gate_up_proj_scale_inv if has_gate else experts.up_proj_scale_inv + ) + input_scale.fill_(1.0) + experts.down_proj_scale_inv.fill_(1.0) + experts.linear = types.MethodType(_cpu_linear, experts) + return experts + + +class _FP8Router(nn.Module): + def __init__(self, config: GlmMoeDsaConfig): + super().__init__() + self.n_routed_experts = config.n_routed_experts + self.weight = nn.Parameter( + torch.arange( + config.n_routed_experts * config.hidden_size, + dtype=torch.float32, + ).reshape(config.n_routed_experts, config.hidden_size) + ) + self.register_buffer( + "e_score_correction_bias", + torch.arange(config.n_routed_experts, dtype=torch.float32), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.linear(hidden_states, self.weight) + + +class _FP8MoEBlock(nn.Module): + def __init__(self, config: GlmMoeDsaConfig, activation_scheme="dynamic"): + super().__init__() + self.n_routed_experts = config.n_routed_experts + self.top_k = config.num_experts_per_tok + self.gate = _FP8Router(config) + self.experts = _fp8_experts(config, activation_scheme) + self.register_buffer( + "e_score_correction_bias", + torch.arange(config.n_routed_experts, dtype=torch.float32), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + original_shape = hidden_states.shape + hidden_states = hidden_states.reshape(-1, original_shape[-1]) + router_logits = self.gate(hidden_states) + router_probabilities = F.softmax(router_logits, dim=-1) + top_k_weights, top_k_indices = torch.topk( + router_probabilities, k=self.top_k, dim=-1 + ) + return self.experts(hidden_states, top_k_indices, top_k_weights).reshape( + original_shape + ) + + +class _FP8MoEModel(nn.Module): + def __init__(self, config: GlmMoeDsaConfig, activation_scheme="dynamic"): + super().__init__() + self.config = config + self.layers = nn.ModuleList([_FP8MoEBlock(config, activation_scheme)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +@pytest.mark.unit +def test_fp8_experts_collect_unweighted_reap_norms(): + model = _FP8MoEModel(_config()) + get_moe_attrs(model, ignore=[]) + experts = model.layers[0].experts + + assert isinstance(experts, FP8PrunableExperts) + assert isinstance(experts, ReapPrunableExpertsProtocol) + + hidden_states = torch.tensor( + [[1.0, 0.0, -1.0, 0.5], [0.5, 1.0, 0.0, -0.5], [1.0, 1.0, 1.0, 1.0]] + ) + top_k_indices = torch.tensor([[0, 1], [2, 0], [3, 1]]) + top_k_weights = torch.tensor([[0.75, 0.25], [0.6, 0.4], [0.9, 0.1]]) + + experts.start_reap_norm_collection() + experts(hidden_states, top_k_indices, top_k_weights) + norms = experts.take_reap_norms() + + assert set(norms) == {0, 1, 2, 3} + assert {index: len(value) for index, value in norms.items()} == { + 0: 2, + 1: 2, + 2: 1, + 3: 1, + } + + tracker = REAPSaliencyTracker(num_experts=4) + tracker.update(top_k_indices, top_k_weights, norms) + flat_indices = top_k_indices.T.reshape(-1) + flat_weights = top_k_weights.T.reshape(-1) + for expert_index, expert_norms in norms.items(): + expected = (flat_weights[flat_indices == expert_index] * expert_norms).mean() + assert tracker.mean_saliency[expert_index].item() == pytest.approx( + expected.item() + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("activation_scheme", ["dynamic", "static"]) +@pytest.mark.parametrize("has_gate", [True, False]) +def test_reap_norm_collection_preserves_fp8_forward(activation_scheme, has_gate): + experts = _fp8_experts( + _config(), activation_scheme=activation_scheme, has_gate=has_gate + ) + hidden_states = torch.tensor( + [[1.0, 0.0, -1.0, 0.5], [0.5, 1.0, 0.0, -0.5], [1.0, 1.0, 1.0, 1.0]] + ) + top_k_indices = torch.tensor([[0, 1], [2, 0], [3, 1]]) + top_k_weights = torch.tensor([[0.75, 0.25], [0.6, 0.4], [0.9, 0.1]]) + + expected = experts(hidden_states, top_k_indices, top_k_weights) + make_fp8_experts_reap_prunable(experts) + experts.start_reap_norm_collection() + actual = experts(hidden_states, top_k_indices, top_k_weights) + + torch.testing.assert_close(actual, expected) + assert set(experts.take_reap_norms()) == {0, 1, 2, 3} + + +@pytest.mark.unit +def test_fp8_adapter_refreshes_compressed_tensors_offload_forward(): + experts = _fp8_experts(_config()) + offload_module(experts, onload_device="cpu", offload_device="cpu") + original_weight = experts.gate_up_proj.detach().clone() + original_scale = experts.gate_up_proj_scale_inv.detach().clone() + + make_fp8_experts_reap_prunable(experts) + + assert isinstance(experts, FP8PrunableExperts) + assert experts._original_forward_func is FP8PrunableExperts.forward + + hidden_states = torch.tensor([[1.0, 0.0, -1.0, 0.5]]) + top_k_indices = torch.tensor([[0, 1]]) + top_k_weights = torch.tensor([[0.75, 0.25]]) + experts.start_reap_norm_collection() + experts(hidden_states, top_k_indices, top_k_weights) + + assert set(experts.take_reap_norms()) == {0, 1} + + retained = [3, 1] + experts.prune_experts_(retained) + torch.testing.assert_close(experts.gate_up_proj, original_weight[retained]) + torch.testing.assert_close(experts.gate_up_proj_scale_inv, original_scale[retained]) + + +@pytest.mark.unit +def test_fp8_adapter_refreshes_accelerate_offload_forward(): + experts = _fp8_experts(_config()) + add_hook_to_module(experts, AlignDevicesHook(execution_device="cpu")) + + make_fp8_experts_reap_prunable(experts) + + assert isinstance(experts, FP8PrunableExperts) + assert experts._old_forward.__func__ is FP8PrunableExperts.forward + + hidden_states = torch.tensor([[1.0, 0.0, -1.0, 0.5]]) + top_k_indices = torch.tensor([[0, 1]]) + top_k_weights = torch.tensor([[0.75, 0.25]]) + experts.start_reap_norm_collection() + experts(hidden_states, top_k_indices, top_k_weights) + + assert set(experts.take_reap_norms()) == {0, 1} + + +@pytest.mark.unit +def test_prune_fp8_experts_slices_weights_scales_router_and_config(): + model = _FP8MoEModel(_config(), activation_scheme="static") + attrs = get_moe_attrs(model, ignore=[]) + layer = model.layers[0] + experts = layer.experts + retained = [3, 1] + + tensor_names = ( + "gate_up_proj", + "gate_up_proj_scale_inv", + "down_proj", + "down_proj_scale_inv", + "gate_up_proj_activation_scale", + "down_proj_activation_scale", + ) + originals = {name: getattr(experts, name).detach().clone() for name in tensor_names} + original_router_weight = layer.gate.weight.detach().clone() + + prune_moe_layer(model, attrs.moe_layer_names[0], retained, attrs) + update_model_config(model, attrs, len(retained)) + + assert experts.num_experts == len(retained) + for name, original in originals.items(): + pruned = getattr(experts, name) + assert pruned.dtype == original.dtype + torch.testing.assert_close(pruned, original[retained]) + + torch.testing.assert_close(layer.gate.weight, original_router_weight[retained]) + torch.testing.assert_close( + layer.gate.e_score_correction_bias, + torch.tensor(retained, dtype=torch.float32), + ) + torch.testing.assert_close( + layer.e_score_correction_bias, + torch.tensor(retained, dtype=torch.float32), + ) + assert layer.n_routed_experts == len(retained) + assert layer.gate.n_routed_experts == len(retained) + assert model.config.n_routed_experts == len(retained) + assert model.config.num_local_experts == len(retained) + + state_keys = set(experts.state_dict()) + assert state_keys == { + "gate_up_proj", + "gate_up_proj_scale_inv", + "down_proj", + "down_proj_scale_inv", + "gate_up_proj_activation_scale", + "down_proj_activation_scale", + } + assert attrs.num_experts_config_key == "n_routed_experts" + + +@pytest.mark.unit +def test_native_fp8_adapter_rejects_non_fp8_weight_storage(): + model = _FP8MoEModel(_config()) + experts = model.layers[0].experts + experts.gate_up_proj = nn.Parameter(experts.gate_up_proj.detach().float()) + + with pytest.raises(TypeError, match="only supports e4m3 FP8"): + get_moe_attrs(model, ignore=[]) + + +@pytest.mark.unit +def test_fp8_reap_modifier_full_lifecycle(): + model = _FP8MoEModel(_config()) + modifier = REAPPruningModifier(sparsity=0.5) + state = State( + model=model, + teacher_model=None, + optimizer=None, + optim_wrapped=False, + loss=None, + batch_data=None, + ) + + modifier.initialize(state) + modifier.update_event(state, Event(type_=EventType.CALIBRATION_START)) + with torch.no_grad(): + model(torch.randn(2, 3, model.config.hidden_size)) + modifier.update_event(state, Event(type_=EventType.SEQUENTIAL_EPOCH_END)) + modifier.update_event(state, Event(type_=EventType.CALIBRATION_END)) + modifier.finalize(state) + + layer = model.layers[0] + assert layer.experts.num_experts == 2 + assert layer.experts.gate_up_proj.shape[0] == 2 + assert layer.gate.weight.shape[0] == 2 + assert layer.e_score_correction_bias.shape[0] == 2 + assert layer.gate.e_score_correction_bias.shape[0] == 2 + assert model.config.n_routed_experts == 2 + with torch.no_grad(): + output = model(torch.randn(2, 3, model.config.hidden_size)) + assert output.shape == (2, 3, model.config.hidden_size) + assert torch.isfinite(output).all() + + +@pytest.mark.unit +def test_pruned_fp8_glm_saves_and_reloads_with_transformers(tmp_path): + config = _tiny_glm_config() + model = GlmMoeDsaForCausalLM(config) + model.model.layers[0].mlp.experts = _fp8_experts(config) + + attrs = get_moe_attrs(model, ignore=[]) + prune_moe_layer(model, attrs.moe_layer_names[0], [3, 1], attrs) + update_model_config(model, attrs, new_num_experts=2) + model.save_pretrained(tmp_path) + + with open(tmp_path / "config.json") as config_file: + saved_config = json.load(config_file) + assert saved_config["n_routed_experts"] == 2 + assert saved_config["quantization_config"]["quant_method"] == "fp8" + + saved_state = load_file(tmp_path / "model.safetensors") + expert_weights = { + name: tensor + for name, tensor in saved_state.items() + if ".mlp.experts." in name and name.endswith(".weight") + } + expert_scales = { + name: tensor + for name, tensor in saved_state.items() + if ".mlp.experts." in name and name.endswith(".weight_scale_inv") + } + assert len(expert_weights) == 6 + assert len(expert_scales) == 6 + assert all( + tensor.dtype == torch.float8_e4m3fn for tensor in expert_weights.values() + ) + + reloaded, loading_info = GlmMoeDsaForCausalLM.from_pretrained( + tmp_path, output_loading_info=True + ) + assert reloaded.config.n_routed_experts == 2 + assert reloaded.model.layers[0].mlp.experts.gate_up_proj.shape[0] == 2 + assert not [key for key in loading_info["missing_keys"] if ".mlp.experts." in key] + assert not [ + key for key in loading_info["unexpected_keys"] if ".mlp.experts." in key + ]