Repository: huggingface/transformers
File: src/transformers/models/mamba/modeling_mamba.py (Python)

Code before:
            out_channels=self.intermediate_size,
            bias=config.use_conv_bias,
            kernel_size=config.conv_kernel,
            groups=self.intermediate_size,
            padding=config.conv_kernel - 1,
        )

        self.activation = config.hidden_act
        self.act = ACT2FN[config.hidden_act]

        self.use_mambapy = config.use_mambapy

        # projection of the input hidden states
        self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=config.use_bias)
        # selective projection used to make dt, B and C input dependent
        self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False)
        # time step projection (discretization)
        self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True)

        dt = torch.exp(
            torch.rand(self.intermediate_size) * (math.log(config.time_step_max) - math.log(config.time_step_min))
            + math.log(config.time_step_min)
        ).clamp(min=config.time_step_floor)
        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        init.copy_(self.dt_proj.bias, inv_dt)

        # S4D real initialization. These are not discretized!
        # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
        A = torch.arange(1, self.ssm_state_size + 1, dtype=torch.float32)[None, :]
        A = A.expand(self.intermediate_size, -1).contiguous()

        self.A_log = nn.Parameter(torch.log(A))
        self.D = nn.Parameter(torch.ones(self.intermediate_size))
        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
        self.use_bias = config.use_bias

        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
        causal_conv1d = lazy_load_kernel("causal-conv1d")
        causal_conv1d_update, causal_conv1d_fn = (
            (causal_conv1d.causal_conv1d_update, causal_conv1d.causal_conv1d_fn)
            if causal_conv1d is not None
            else (None, None)
        )
        global mamba_ssm, selective_state_update, selective_scan_fn, mamba_inner_fn
        mamba_ssm = lazy_load_kernel("mamba-ssm")
        selective_state_update, selective_scan_fn, mamba_inner_fn = (
            (mamba_ssm.selective_state_update, mamba_ssm.selective_scan_fn, mamba_ssm.mamba_inner_fn)
            if mamba_ssm is not None
            else (None, None, None)
        )

Review comment (suggestion):
Indeed, as @ArthurZucker said, this will be overwritten in most case. We start by initializing the models on meta device, so all weights defined in `__init__` do not get instantiated (only their shape etc).
Then, if the weight is missing, or if we want a fully randomly instantiated model, the proper way is to it in `_init_weights` currently, which will be called later on.
So here this should go inside this block of the `_init_weights`:

```python
def _init_weights(self, module):
        """Initialize the weights."""
        std = self.config.initializer_range
        if isinstance(module, MambaMixer):
```

Code after:
            out_channels=self.intermediate_size,
            bias=config.use_conv_bias,
            kernel_size=config.conv_kernel,
            groups=self.intermediate_size,
            padding=config.conv_kernel - 1,
        )

        self.activation = config.hidden_act
        self.act = ACT2FN[config.hidden_act]

        self.use_mambapy = config.use_mambapy

        # projection of the input hidden states
        self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=config.use_bias)
        # selective projection used to make dt, B and C input dependent
        self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False)
        # time step projection (discretization)
        self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True)

        # S4D real initialization. These are not discretized!
        # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
        self.A_log = nn.Parameter(torch.empty(self.intermediate_size, self.ssm_state_size))
        self.D = nn.Parameter(torch.empty(self.intermediate_size))
        if initialize_mixer_weights and self.dt_proj.weight.device.type != "meta":
            self.init_mamba_weights()
        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
        self.use_bias = config.use_bias

        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
        causal_conv1d = lazy_load_kernel("causal-conv1d")
        causal_conv1d_update, causal_conv1d_fn = (
            (causal_conv1d.causal_conv1d_update, causal_conv1d.causal_conv1d_fn)
            if causal_conv1d is not None
            else (None, None)
        )
        global mamba_ssm, selective_state_update, selective_scan_fn, mamba_inner_fn
        mamba_ssm = lazy_load_kernel("mamba-ssm")
        selective_state_update, selective_scan_fn, mamba_inner_fn = (
            (mamba_ssm.selective_state_update, mamba_ssm.selective_scan_fn, mamba_ssm.mamba_inner_fn)
            if mamba_ssm is not None
            else (None, None, None)
        )

        self.warn_slow_implementation()

    @torch.no_grad()
    def init_mamba_weights(self):
        A = torch.arange(1, self.ssm_state_size + 1, dtype=torch.float32, device=self.A_log.device)[None, :]
        A = A.expand(self.intermediate_size, -1).contiguous()
        init.copy_(self.A_log, torch.log(A))
        init.ones_(self.D)