Discrepancy between loading models with meta tensors and normal load_from_state_dict

username: patrickvonplaten
cooked: <p>The PyTorch <a href="https://pytorch.org/docs/stable/generated/torch.Tensor.is_meta.html" class="inline-onebox" rel="noopener nofollow ugc">torch.Tensor.is_meta — PyTorch 1.13 documentation</a> “meta” device was introduced a while back to load models efficiently into memory. More precisely it allows one to not have to allocate memory before loading the actual model weights, <em>i.e.</em>:</p>
<pre><code class="lang-python">import torch
layer = torch.nn.Linear(100, 100)
</code></pre>
<p>allocates memory while</p>
<pre><code class="lang-python">layer = torch.nn.Linear(100, 100, device="meta")
</code></pre>
<p>doesn’t. This is extremely important for the ML community because language models (<a href="https://huggingface.co/EleutherAI/gpt-j-6B" rel="noopener nofollow ugc">GPT-J</a>, <a href="https://huggingface.co/docs/transformers/model_doc/bloom" rel="noopener nofollow ugc">BLOOM</a>) have become so huge that even CPU RAM can be a limiting factor.</p>
<p>Another super nice additional gain of <code>device="meta"</code> is that it doesn’t run any weight initialization code which can save up to 95% of time when loading models.</p>
<p>Hence, loading models with the meta device has already been widely adopted by the community:</p>
<ul>
<li>Accelerate offers this functionality: <a href="https://huggingface.co/docs/accelerate/usage_guides/big_modeling#how-the-process-works-working-with-code" class="inline-onebox" rel="noopener nofollow ugc">Handling big models</a>
</li>
<li>
<a href="https://github.com/huggingface/transformers" rel="noopener nofollow ugc">Transformers</a> and <a href="https://github.com/huggingface/diffusers" rel="noopener nofollow ugc">Diffusers</a> have adopted it (diffusers has made it even the default way of loading models)</li>
<li>Galactica uses it via accelerate: <a href="https://github.com/paperswithcode/galai/blob/471a22d152079baa5823b7bdf18aa3ad0b47b39e/galai/model.py#L112" class="inline-onebox" rel="noopener nofollow ugc">galai/model.py at 471a22d152079baa5823b7bdf18aa3ad0b47b39e · paperswithcode/galai · GitHub</a>
</li>
<li>…</li>
</ul>
<p>The problem is that there is currently a bit of a discrepancy of <strong>dtype casting</strong> between loading a model with <code>device="meta"</code> and without. Consider the following code snippet:</p>
<pre><code class="lang-python">import os
import torch
from torch import nn
import tempfile


print("1. save &amp; load raw dict in fp16 precision")
with tempfile.TemporaryDirectory() as tmpdirname:
    layer = nn.Linear(100, 100)
    torch.save(layer.state_dict(), os.path.join(tmpdirname, "fp32.bin"))
    fp16_dict = torch.load("fp16.bin")

with tempfile.TemporaryDirectory() as tmpdirname:
    layer = nn.Linear(100, 100).half()
    torch.save(layer.state_dict(), "fp16.bin")
    fp32_dict = torch.load("fp32.bin")

print("2. load weights into fp32 dtype")
print(25 * "-")
layer_fp32 = nn.Linear(100, 100)
layer_fp32.load_state_dict(fp32_dict)
print("Fp32 into Fp32 =&gt;", next(iter(layer_fp32.parameters())).dtype)  # =&gt; fp32

layer_fp32 = nn.Linear(100, 100)
layer_fp32.load_state_dict(fp16_dict)
print("Fp16 into Fp32 =&gt;", next(iter(layer_fp32.parameters())).dtype)  # =&gt; fp32

print("3. load weights into fp16 dtype")
print(25 * "-")
layer_fp16 = nn.Linear(100, 100, dtype=torch.float16)  # =&gt; fp16
layer_fp16.load_state_dict(fp32_dict)
print("Fp32 into Fp16 =&gt;", next(iter(layer_fp16.parameters())).dtype)

layer_fp16 = nn.Linear(100, 100, dtype=torch.float16)  # =&gt; fp16
layer_fp16.load_state_dict(fp16_dict)
print("Fp16 into Fp16 =&gt;", next(iter(layer_fp16.parameters())).dtype)

# =&gt; Makes sense! Weights are always upcasted / downcasted to "expected random init weights"
# =&gt; Now let's check out meta

print("4. load weights into meta")
print(25 * "-")
layer_meta = nn.Linear(100, 100, device="meta")
# Let's load weights as before
layer_meta.load_state_dict(fp32_dict)  # this is essentially a no-op since weights aren't filled; why? is this supposed to stay?
print("Show loaded weigths when using state_dict: ", next(iter(layer_meta.parameters())))

# Ok doesn't work, let's instead set the weights as params (think that's how it's done in accelerate)
for key, value in fp16_dict.items():
    param_cls = type(layer_meta._parameters[key])
    new_value = param_cls(value, **layer_meta._parameters[key].__dict__)
    layer_meta._parameters[key] = new_value
print("Fp16 into meta =&gt;", next(iter(layer_meta.parameters())).dtype) # =&gt; fp16

layer_meta = nn.Linear(100, 100, device="meta")
for key, value in fp32_dict.items():
    param_cls = type(layer_meta._parameters[key])
    new_value = param_cls(value, **layer_meta._parameters[key].__dict__)
    layer_meta._parameters[key] = new_value
print("Fp32 into meta =&gt;", next(iter(layer_meta.parameters())).dtype) # =&gt; fp32
</code></pre>
<p>If you execute this code, you can see that the <strong>dtype is always casted to the dtype expected by the randomly initialized weights</strong> =&gt; this makes sense. The problem is when you make use of <code>device="meta"</code> the are no randomly initialized weights really so what API / behavior should be adopted here?</p>
<p>Some possible solutions that don’t seem to work currently:</p>
<ul>
<li>1.) Actually allow to load a state_dict into a module that has <code>device="meta"</code> weights. E.g. this codesnippet <code>layer_meta.load_state_dict(fp32_dict) </code> is currently a no-op - is the plan to change this? When doing so should maybe the <code>dtype</code> of the “meta” weight also define the dtype of the loaded weights? To be more precise when doing:</li>
</ul>
<pre><code class="lang-auto">layer_meta = nn.Linear(100, 100, device="meta")
</code></pre>
<p>the module <code>layer_meta</code> actually has a dtype which defaults to <code>fp32</code> just like “normal” modules. The problem however is that <code>load_state_dict</code> on this module doesn’t work libraries currently cannot adopt this approach. Also, would it make sense to do this? If a module <strong>doesn’t initialize random weights</strong> are things like <code>module.dtype</code> meaningful? Because the <code>dtype</code> really is an attribute of the attached weights and there are no weights by definition, so does it make sense to actually use the <code>dtype</code> of a “meta” tensor?</p>
<ul>
<li>2.) The loaded weights define what <code>dtype</code> the model should have. If I instantiate a module that allocates no weights, no memory with <code>torch.meta</code> and then load weights in fp16 into such a module, should the module maybe just take this dtype? This seems to make sense to me, <strong>but</strong> the big drawback is that it would be a bit inconsistent with the existing API and people will see different results depending on whether the model was initialized with <code>device="meta"</code> or not</li>
<li>3.) Not do anything since <code>device="meta"</code> is still pretty new and has no official docs yet really (but as said before it’s already widely used by the community)</li>
</ul>
<p>To throw in another pointer that might be helpful. In JAX/Flax the weights are not attached to the module and only forwarded to the model. Therefore only the loaded weights decide the dtype of the model. There is no problematic question of whether the “random model weights” should be upcasted to the “loaded model weights” or the other way around.</p>