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]]