Class DogePreTrainedModel

src/transformers/models/doge/modeling_doge.py:503–531

@auto_docstring
class DogePreTrainedModel(PreTrainedModel):
    config: DogeConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["DogeDecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn = False
    _supports_sdpa = True
    _supports_flex_attn = True
    _can_compile_fullgraph = False
    _supports_attention_backend = True
    _can_record_outputs = {
        "router_logits": OutputRecorder(DogeCDMoE, index=1),
        "hidden_states": DogeDecoderLayer,
        "attentions": DogeAttention,
    }

    @torch.no_grad()
    def _init_weights(self, module):
        """Initialize the weights"""
        super()._init_weights(module)
        if isinstance(module, DogeAttention):
            if hasattr(module, "A"):
                init.zeros_(module.A)
        elif isinstance(module, DogeDecoderLayer):
            if hasattr(module, "input_residual"):
                init.ones_(module.input_residual)
            if hasattr(module, "post_attention_residual"):
                init.ones_(module.post_attention_residual)


@auto_docstring