gpt_oss_model_descriptor

GPT-OSS model descriptor for AnyModel compression.

Classes

GptOssExpertRemovalLayerDescriptor

GPT-OSS MoE layer descriptor for expert removal.

GptOssKVHeadsLayerDescriptor

GptOssModelDescriptor

Model descriptor for GPT-OSS (pure MoE model).
class GptOssExpertRemovalLayerDescriptor

Bases: ExpertRemovalLayerDescriptor

GPT-OSS MoE layer descriptor for expert removal.

Note: This only works for unquantized models (e.g., test models). Production GPT-OSS models use MXFP4 quantization with fused experts (_blocks, _scales, _bias), which requires a different approach.

Structure: - Router: mlp.router with .weight and .bias - Experts: mlp.experts.{idx}.{gate_up_proj,down_proj} with .weight and .bias

__init__(target_name='mlp', moe_prefix_name='model.layers.{layer_idx}.mlp', expert_prefix_name='experts', router_weights=<factory>, router_biases=<factory>, expert_weights=<factory>, expert_biases=<factory>, is_fused_experts=True, fused_expert_weights=<factory>)
Parameters:

  • target_name (str)

  • moe_prefix_name (str)

  • expert_prefix_name (str)

  • router_weights (List[str])

  • router_biases (List[str])

  • expert_weights (List[str])

  • expert_biases (List[str])

  • is_fused_experts (bool)

  • fused_expert_weights (List[str])

Return type:

None

expert_biases: List[str]

Per-expert bias names relative to expert_prefix (per-expert format).

expert_prefix_name: str = 'experts'

Expert prefix relative to moe_prefix with {expert_idx} placeholder, e.g. experts.{expert_idx}.

expert_weights: List[str]

Per-expert weight names relative to expert_prefix (per-expert format).

fused_expert_weights: List[str]

Fused expert weight names relative to moe_prefix, e.g. ["experts.gate_up_proj", "experts.down_proj"].

get_modules_names_to_hook(model)
Return type:

List[Tuple[int, str]]

is_fused_experts: bool = True

If True, experts are stored as single fused tensors (shape [num_experts, ...]).

moe_prefix_name: str = 'model.layers.{layer_idx}.mlp'

MoE prefix layer name with {layer_idx} placeholder, e.g. model.layers.{layer_idx}.moe.

router_biases: List[str]

Router bias names relative to moe_prefix.

router_weights: List[str]

Router weight names relative to moe_prefix.

target_name: str = 'mlp'

Module name for hook registration; supports regex: prefix.

class GptOssKVHeadsLayerDescriptor

Bases: KVHeadsLayerDescriptor

__init__(o_proj_name='self_attn.o_proj', attn_prefix_name='model.layers.{layer_idx}.self_attn', qkvo_weight_names=<factory>)
Parameters:

  • o_proj_name (str)

  • attn_prefix_name (str)

  • qkvo_weight_names (List[str])

Return type:

None

attn_prefix_name: str = 'model.layers.{layer_idx}.self_attn'
o_proj_name: str = 'self_attn.o_proj'
qkvo_weight_names: List[str]
class GptOssModelDescriptor

Bases: ModelDescriptor

Model descriptor for GPT-OSS (pure MoE model).

classmethod anymodel_arch_info()

Map HF GPT-OSS layers to the vLLM implementation used by AnyModel.

Return type:

dict[str, object]

static attn_no_op_post_init(decoder_layer)

Replace attention sublayers with no-op modules.

classmethod automodel_model_kwargs(config, *, distributed=None)

Use GPT-OSS’s supported non-TE attention backend.

Passing a partial backend mapping for native TP disables the AutoModel class’s backend=None default, so the descriptor must keep its required Flex attention choice explicit.

classmethod block_config_to_layer_overrides(block_config)

Map BlockConfig through the shared structural contract.

Parameters:

block_config (BlockConfig)

classmethod checkpoint_equivalence_tolerances()

Allow measured batch-sensitive MXFP4 drift from a basis permutation.

The LM-loss delta reaches 0.01944 on the eight-sample diagnostic batch even when KL, top-1 agreement, cosine, and normalized hidden MSE all satisfy their independent equivalence gates.

Return type:

dict[str, float]

classmethod create_dummy_block(original_layer, block_index)
Parameters:

  • original_layer (GptOssDecoderLayer)

  • block_index (int)

Return type:

Module

classmethod create_runtime_benchmark_model(runtime_config, block_configs)
Parameters:

block_configs (list[BlockConfig])

static decoder_layer_cls()

Get the decoder layer class for GPT-OSS models.

GPT-OSS is a standard transformers model in recent versions. Import directly from transformers.models.gpt_oss.modeling_gpt_oss.

classmethod embedding_pruning_spec(config, *, widths, alignment)
Parameters:

alignment (int)

static final_norm_name()
classmethod generic_decoder_contract(config)
static init_rotary_embedding(model, runtime)

Initialize rotary embeddings on the correct device.

static input_embedding_name()
static layer_block_name(index)
Parameters:

index (int)

static layer_name_predicates(num_layers)

Define regex patterns for grouping weights into subblocks.

Parameters:

num_layers (int)

Return type:

Dict[str, Pattern]

static mlp_no_op_post_init(decoder_layer)

Replace MLP sublayers with no-op modules.

Note: GPT-OSS MoE layers return (hidden_states, router_scores), so we need to return a tuple of 2 values.

static output_embedding_name()
classmethod patch_layer_config(layer_config, block_config, layer_idx)
Parameters:
Return type:

None

classmethod patch_pipeline_model_part(model_part)

Restore GPT-OSS’s stage-aware native inner forward after PP split.

AutoModel’s generic CausalLM PP forward precomputes HF-style rotary embeddings. GPT-OSS instead computes freqs_cis inside its native inner forward, which already tolerates stage-local missing embeddings and norms. Keep the generic outer LM wrapper but restore the inner class method on each split chunk.

Parameters:

model_part (Module)

Return type:

bool

static pruning_mixins()

Return available pruning mixins for GPT-OSS.

Note: Expert removal works for unquantized models (test models). Production models use MXFP4 quantization which is not yet supported.

Return type:

Dict[str, PruningMixIn]

classmethod puzzletron_capabilities(config)
classmethod runtime_benchmark_base_block_config(runtime_config)
Return type:

BlockConfig

classmethod runtime_benchmark_config_fields(lm_config)

Fields needed to build bounded, structurally faithful GPT-OSS proxies.

Return type:

dict[str, object]

static sorted_teacher_layout_kwargs(_lm_config)

Declare query-head state and fused expert storage owned by GPT-OSS.

Return type:

dict[str, object]