What does nn.utils.skip_init() really do?

vince62s: <p>I am working with the Linear4bit class from bitsandbytes here:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L204">
  <header class="source">

      <a href="https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L204" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L204" target="_blank" rel="noopener nofollow ugc">TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L204</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="194" style="counter-reset: li-counter 193 ;">
          <li>                self.quant_state.to(device)</li>
          <li></li>
          <li>            new_param = Params4bit(super().to(device=device, dtype=dtype, non_blocking=non_blocking),</li>
          <li>                                   requires_grad=self.requires_grad, quant_state=self.quant_state,</li>
          <li>                                   blocksize=self.blocksize, compress_statistics=self.compress_statistics,</li>
          <li>                                   quant_type=self.quant_type)</li>
          <li></li>
          <li>            return new_param</li>
          <li></li>
          <li></li>
          <li class="selected">class Linear4bit(nn.Linear):</li>
          <li></li>
          <li>    def __init__(self, input_features, output_features, bias=True, compute_dtype=None, compress_statistics=True, quant_type='fp4', device=None):</li>
          <li>        super().__init__(input_features, output_features, bias, device)</li>
          <li>        self.weight = Params4bit(self.weight.data, requires_grad=False, compress_statistics=compress_statistics, quant_type=quant_type)</li>
          <li>        # self.persistent_buffers = []  # TODO consider as way to save quant state</li>
          <li>        self.compute_dtype = compute_dtype</li>
          <li>        self.compute_type_is_set = False</li>
          <li></li>
          <li>    def set_compute_type(self, x):</li>
          <li>        if x.dtype in [torch.float32, torch.bfloat16]:</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Using this snippet</p>
<pre><code class="lang-auto">import bitsandbytes as bnb
import torch.nn as nn
import torch

model = nn.utils.skip_init(
            bnb.nn.Linear4bit,
            16,
            16,
            bias=None,
            quant_type='fp4',
            device=torch.device("cpu")
        )

print(model.weight)

model.weight.data = torch.randint(0, 100, [16,16])

model.cuda()

input = torch.randint(0, 16, [16]).float().cuda()

print(model(input))
</code></pre>
<p>triggers an error</p>
<p><code>AttributeError: 'Parameter' object has no attribute 'quant_state'</code></p>
<p>because the statement model.cuda() does NOT go through this:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L167">
  <header class="source">

      <a href="https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L167" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L167" target="_blank" rel="noopener nofollow ugc">TimDettmers/bitsandbytes/blob/main/bitsandbytes/nn/modules.py#L167</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="157" style="counter-reset: li-counter 156 ;">
          <li>@classmethod</li>
          <li>def from_prequantized(cls, data: torch.Tensor, quantized_stats: Dict[str, Any], requires_grad: bool = False, device='cuda', **kwargs) -&gt; "Params4bit":</li>
          <li>    self = torch.Tensor._make_subclass(cls, data.to(device))</li>
          <li>    self.requires_grad = requires_grad</li>
          <li>    self.quant_state = QuantState.from_dict(qs_dict=quantized_stats, device=device)</li>
          <li>    self.blocksize = self.quant_state.blocksize</li>
          <li>    self.compress_statistics = self.quant_state.nested</li>
          <li>    self.quant_type = self.quant_state.quant_type</li>
          <li>    return self</li>
          <li></li>
          <li class="selected">def cuda(self, device):</li>
          <li>    w = self.data.contiguous().half().cuda(device)</li>
          <li>    w_4bit, quant_state = bnb.functional.quantize_4bit(w, blocksize=self.blocksize, compress_statistics=self.compress_statistics, quant_type=self.quant_type)</li>
          <li>    self.data = w_4bit</li>
          <li>    self.quant_state = quant_state</li>
          <li></li>
          <li>    return self</li>
          <li></li>
          <li>@overload</li>
          <li>def to(self: T, device: Optional[Union[int, device]] = ..., dtype: Optional[Union[dtype, str]] = ..., non_blocking: bool = ...,) -&gt; T:</li>
          <li>    ...</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>which does not happen if I do not use skip_init()<br>
I thought the use of skip_init was just to avoid the uniform/normal init of parameters inside nn.Linear<br>
but I have the feeling it does something else.</p>

vince62s: <p>here is the reason (on the doc)</p>
<ol start="2">
<li>The module must not perform any computation on parameters in its constructor except initialization (i.e. functions from <code>torch.nn.init</code>).</li>
</ol>
<p>so if we use skip_init() we need to manually perform init tasks in Linear4bit().</p>