Repository: huggingface/transformers
File: src/transformers/models/gpt_oss/modular_gpt_oss.py (Python)

Code before:
        "attentions": GptOssAttention,
    }

    @torch.no_grad()
    def _init_weights(self, module):
        PreTrainedModel._init_weights(self, module)
        std = self.config.initializer_range
        if isinstance(module, GptOssExperts):
            init.normal_(module.gate_up_proj, mean=0.0, std=std)
            init.zeros_(module.gate_up_proj_bias)
            init.normal_(module.down_proj, mean=0.0, std=std)
            init.zeros_(module.down_proj_bias)
        elif isinstance(module, GptOssAttention):
            init.normal_(module.sinks, mean=0.0, std=std)
        elif isinstance(module, GptOssTopKRouter):
            init.normal_(module.weight, mean=0.0, std=std)
            init.normal_(module.bias, mean=0.0, std=std)


class GptOssModel(MixtralModel):
    _no_split_modules = ["GptOssDecoderLayer"]

    def __init__(self, config: GptOssConfig):
        super().__init__(config)

        if config._attn_implementation in ["flash_attention_2", "flash_attention_3"]:
            raise ValueError(
                f"GPT-OSS models do not support {config._attn_implementation} because they utilize "
                "attention sinks, which are not currently supported by the standard Flash Attention kernels. "
                "Please use 'eager' implementation (attn_implementation='eager') or a custom kernel if available."
            )

            
    @check_model_inputs
    @auto_docstring
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> MoeModelOutputWithPast:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

Review comment (suggestion):
2 things that might not work as expected
- This only checks during init, but we can also update afterwards with e.g. `set_attn_implementation`. We probably cannot avoid checking during runtime (e,g, forward).
- There are more FA implementations especially with the kernels ones. Afaik, we only have the vllm kernel working so we should check along `if "flash" in config._attn_implementation and config._attn_implementation != "vllm-kernel": raise ValueError(...)` (don't remember the name of the kernel). 

Code after:
    @torch.no_grad()
    def _init_weights(self, module):
        PreTrainedModel._init_weights(self, module)
        std = self.config.initializer_range
        if isinstance(module, GptOssExperts):
            init.normal_(module.gate_up_proj, mean=0.0, std=std)
            init.zeros_(module.gate_up_proj_bias)
            init.normal_(module.down_proj, mean=0.0, std=std)
            init.zeros_(module.down_proj_bias)
        elif isinstance(module, GptOssAttention):
            init.normal_(module.sinks, mean=0.0, std=std)
        elif isinstance(module, GptOssTopKRouter):
            init.normal_(module.weight, mean=0.0, std=std)
            init.normal_(module.bias, mean=0.0, std=std)


class GptOssModel(MixtralModel):
    _no_split_modules = ["GptOssDecoderLayer"]

    @check_model_inputs
    @auto_docstring
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> MoeModelOutputWithPast:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        # It may already have been prepared by e.g. `generate`
        if not isinstance(causal_mask_mapping := attention_mask, dict):