Model being occupying 2.5x the size - loading and running

aerialmyth: <p>I am trying to load and run the <code>meta-llama/Meta-Llama-3-8B</code> on a 16 GB Nvidia RTX 4060 Ti.</p>
<p>I am using the following <code>BitsAndBytesConfig</code> to load the model in 4-bit and use <code>torch.bfloat16</code> for compute.</p>
<pre><code class="lang-auto">bnb_config = BitsAndBytesConfig(
    # load_in_8bit=True,
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModel.from_pretrained(
    'meta-llama/Meta-Llama-3-8B',
    quantization_config=bnb_config,
    device_map='auto'
)
</code></pre>
<p>Just loading the model, the <code>nvidia-smi</code> output shows almost about <strong>4.5 GB</strong> of GPU memory being used.</p>
<p>Then when I run the following command, the GPU memory immediately shoots up to <strong>15.5 GB</strong>.</p>
<pre><code class="lang-auto">inp_sample = next(iter(train_loader))
logits = model(ex_sample['input_ids'])
</code></pre>
<p>I think I am missing something but just running one sample of size <code>[2, 512]</code> through this <strong>Meta-Llama-3-8B</strong> increases the memory by this much.</p>
<p>Added a couple of screenshots for reference. Would love any pointers on this one.<br>
<code>Just Loading the model</code><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/e/de65aa58b2af7e75fdeed7e79e25b1b63a3029e1.png" data-download-href="https://discuss.pytorch.org/uploads/default/de65aa58b2af7e75fdeed7e79e25b1b63a3029e1" title="Screenshot from 2024-06-29 15-04-00"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/e/de65aa58b2af7e75fdeed7e79e25b1b63a3029e1.png" alt="Screenshot from 2024-06-29 15-04-00" data-base62-sha1="vJpOz7Dky9YZya9e80LEThqd0t3" width="690" height="184" data-dominant-color="441B37"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-06-29 15-04-00</span><span class="informations">798×213 22.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I want to try fine tuning the model after but that seems impossible now if the an 8B model 4-bit model runs a 16GB GPU out of memory.</p>

aerialmyth: <p>Adding the second screenshot here because the help forum doesn’t allow new users have more than one embedded items in the post.</p>
<p><code>After running one inference step through the model</code><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/4/a4d5bfd57ac461643a805fa2f1a2af10120b3569.png" data-download-href="https://discuss.pytorch.org/uploads/default/a4d5bfd57ac461643a805fa2f1a2af10120b3569" title="Screenshot from 2024-06-29 15-04-47"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/4/a4d5bfd57ac461643a805fa2f1a2af10120b3569.png" alt="Screenshot from 2024-06-29 15-04-47" data-base62-sha1="nwcgbnVfhugQaxDfjxSi4tO5EH7" width="690" height="184" data-dominant-color="441B37"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-06-29 15-04-47</span><span class="informations">798×213 22.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>Besides the model parameters and buffers, the model execution will also use memory to store intermediate forward activations needed to compute the gradients. You could use e.g. <a href="https://discuss.pytorch.org/t/number-of-activations-without-loading-the-model/202419/2">this until.</a> to estimate the memory usage.</p>
<p>If you don’t want to compute the gradients at all, you should wrap the forward pass into a <code>torch.no_grad()</code> context, if that’s not already done.</p>