moe_mamba_pruning_mixin#

Classes

MoELayerDescriptor

Descriptor for MoE activation scoring targets.

MoEPruningMixIn

MambaLayerDescriptor

MambaPruningMixIn

class MambaLayerDescriptor#

Bases: LayerDescriptor

__init__(target_name='mixer.in_proj', mamba_prefix_name='backbone.layers.{layer_idx}.mixer', in_proj_name='in_proj', out_proj_name='out_proj')#
Parameters:

  • target_name (str)

  • mamba_prefix_name (str)

  • in_proj_name (str)

  • out_proj_name (str)

Return type:

None

canonical_score_key(method, block_idx, fallback)#
Parameters:

  • method (str)

  • block_idx (int | None)

  • fallback (str)

Return type:

str

in_proj_name: str = 'in_proj'#
mamba_prefix(layer_idx)#
Parameters:

layer_idx (int)

Return type:

str

mamba_prefix_name: str = 'backbone.layers.{layer_idx}.mixer'#
module_name_regex()#
Return type:

str

out_proj_name: str = 'out_proj'#
target_name: str = 'mixer.in_proj'#
class MambaPruningMixIn#

Bases: PruningMixIn

__init__(layer_descriptor)#
Parameters:

layer_descriptor (MambaLayerDescriptor)

supported_hooks()#
Return type:

List[Type[ForwardHook]]

class MoELayerDescriptor#

Bases: LayerDescriptor

Descriptor for MoE activation scoring targets.

target_name can be a concrete suffix or a regex: pattern. The canonical key helpers always return checkpoint/AnyModel names under moe_prefix_name, while the hook target can be a native AutoModel module path.

__init__(target_name='mixer', moe_prefix_name='backbone.layers.{layer_idx}.mixer', gate_name='gate', experts_name='experts', shared_experts_name='shared_experts', latent_fc1_name='fc1_latent_proj', latent_fc2_name='fc2_latent_proj', require_attrs=(), expert_down_proj_regex='\\.experts\\.(\\d+)\\.down_proj$')#
Parameters:

  • target_name (str)

  • moe_prefix_name (str)

  • gate_name (str)

  • experts_name (str)

  • shared_experts_name (str)

  • latent_fc1_name (str)

  • latent_fc2_name (str)

  • require_attrs (tuple[str, ...])

  • expert_down_proj_regex (str)

Return type:

None

canonical_score_key(method, block_idx, fallback)#
Parameters:

  • method (str)

  • block_idx (int | None)

  • fallback (str)

Return type:

str

expert_down_proj_regex: str = '\\.experts\\.(\\d+)\\.down_proj$'#
expert_idx_from_module_name(module_name)#
Parameters:

module_name (str)

Return type:

int | None

experts_name: str = 'experts'#
experts_prefix(layer_idx)#
Parameters:

layer_idx (int)

Return type:

str

gate_name: str = 'gate'#
gate_prefix(layer_idx)#
Parameters:

layer_idx (int)

Return type:

str

get_modules_names_to_hook(model)#
Return type:

List[Tuple[int, str]]

latent_fc1_name: str = 'fc1_latent_proj'#
latent_fc2_name: str = 'fc2_latent_proj'#
module_name_regex()#
Return type:

str

moe_prefix(layer_idx)#
Parameters:

layer_idx (int)

Return type:

str

moe_prefix_name: str = 'backbone.layers.{layer_idx}.mixer'#
require_attrs: tuple[str, ...] = ()#
shared_experts_name: str = 'shared_experts'#
shared_experts_prefix(layer_idx)#
Parameters:

layer_idx (int)

Return type:

str

target_name: str = 'mixer'#
class MoEPruningMixIn#

Bases: PruningMixIn

__init__(layer_descriptor)#
Parameters:

layer_descriptor (MoELayerDescriptor)

supported_hooks()#
Return type:

List[Type[ForwardHook]]