Keeping only part of model parameters on GPU

Noveen_Sachdeva: <p>Hi everyone,</p>
<p>I have a specific use-case wherein I want to dynamically move model parameters to GPU from CPU, and once the weight updates are done (on the GPU), transfer it back again on the CPU.</p>
<p>To be specific, I am dealing with a classification problem, where the number of classes is huge. Let’s say we have a linear setup, i.e we want to learn a weight matrix of size <code>[num_classes x feature_size]</code>. The problem is that we can’t afford to keep this matrix completely on the GPU. Hence, what I wanted is to initialize everything on CPU and keep loading parts of the matrix onto GPU and update correspondingly.</p>
<p>With this ideology, I’ll have to declare the weight parameters on CPU (along with their optimizers) and when I want to load part of the matrix on GPU, I’ll have to call <code>.cuda()</code>, which, as the docs suggest (<a href="https://pytorch.org/docs/stable/nn.html?highlight=cpu#torch.nn.Module.cuda" rel="nofollow noopener">LINK</a>) needs to be called before optimizer construction. Anyways, on trying this approach, as expected the model weights don’t update.</p>
<p>Any help would be greatly appreciated.</p>
<p>Thanks,<br>
Noveen Sachdeva</p>

albanD: <p>Hi,</p>
<p>Indeed, you will need to keep the <code>nn.Parameter</code> on the CPU and only perform computation on the GPU in a differentiable manner. That way, it will automatically accumulate back onto the CPU parameters.<br>
Something like:</p>
<pre><code class="lang-python">full_params = torch.rand(num_classes, feature_size, device="cpu")

loss = 0.
for slice in slices:
  param_subset = full_params[slice]
  param_subset = param_subset.cuda()
  # Compute that part of the loss
  loss += local_loss
loss.backward()
</code></pre>
<p>Does that match what you want?</p>

Noveen_Sachdeva: <p>Thanks for the prompt response, <a class="mention" href="/u/alband">@albanD</a>.</p>
<p>As for your solution, I have some concerns. Namely:</p>
<ul>
<li>Wouldn’t doing <code>loss += local_loss</code> keep storing the computation graph for each slice? Isn’t this is infeasible because it would store all the intermediary variables/nodes? Just a side note, in my use-case, mathematically I can afford to do a backward call every time I compute the loss function for a slice. So I don’t need to keep adding the loss and can call <code>.backward()</code> for each slice.</li>
<li>As for doing <code>param_subset = param_subset.cuda()</code>, you load the slice on the GPU, but never take it off/to the CPU. Should it matter? Does PyTorch automatically delete GPU variables once they go out of scope?</li>
<li>Also you didn’t talk about optimizers here. The main problem as per my experiments lies in the opimizer. As per the docs (<a href="https://pytorch.org/docs/stable/nn.html?highlight=cpu#torch.nn.Module.cuda" rel="nofollow noopener">Linked</a> again), I need to keep the parameters that are being optimized on the same hardware while declaring the optimizer class and calling <code>optimizer.step()</code>. I have attached the main part of my code for reference.</li>
</ul>
<pre><code class="lang-auto">weight_matrices = []; bias_matrices = []; optimizers = []
for slice in slices:
    weight_matrices.append(
        nn.Parameter(
            torch.FloatTensor(
                num_features, slice_classes
            ).uniform_()
        )
    )
    bias_matrices.append(
        nn.Parameter(
            torch.FloatTensor(
                1, slice_classes
            ).uniform_()
        )
    )
    optimizers.append(
        torch.optim.Adam(
            [ weight_matrices[-1], bias_matrices[-1] ], 
            lr = 0.01
        )
    )

# Notice I've initialized everything on the CPU

for x, y in train_data:
    # x (some feature vector batch), y (indices batch) already on GPU
    for slice in slices:
        w = weight_matrices[slice]
        bias = bias_matrices[slice]
        optimizer = optimizers[slice]

        # Transferring parameters to GPU
        w = w.cuda()
        bias = bias.cuda()

        # Forward Pass
        scores = (w_t @ x) + bias

        # Backward Pass
        # Some pointwise loss
        # Independent of scores for classes outside the slice
        loss = slice_loss(scores, y) 
        loss.backward()
        optimizer.step()

        # Do I need to manually call .cpu() on w and bias?
</code></pre>
<p>Running this code doesn’t update the weight and bias parameters. I feel this could be because of two potential problems:</p>
<ul>
<li>Doing <code>w = w.cuda()</code> and <code>bias = bias.cuda()</code> creates two non-leaf variables which doesn’t pass the gradients, and hence, doesn’t update <code>w</code> and <code>bias</code>. (See <a href="https://discuss.pytorch.org/t/variable-grad-is-always-none-when-extending-autograd/12187">LINK</a> for existing discussion thread)</li>
<li>I’ve initialized <code>w</code> and <code>bias</code> (for all slices) on CPU and I’m loading them one-by-one on GPU. Should this be a problem for the optimizer?</li>
</ul>
<p>Apologies for the detailed question. Thank you for your time, appreciate it!</p>
<p>-Noveen.</p>

albanD: <blockquote>
<p>Does PyTorch automatically delete GPU variables once they go out of scope?</p>
</blockquote>
<p>Yes ! You would get a lot of memory leaks otherwise <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>
<blockquote>
<p>nn.Parameter</p>
</blockquote>
<p>Note that <code>nn.Parameter</code> are <em>only</em> important if your work inside a <code>nn.Module</code> and you want them to be found by <code>mod.parameters()</code>. Otherwise, there is no need to have them.</p>
<p>Here is an updated example with optimizer:</p>
<pre><code class="lang-auto">full_params = torch.rand(num_classes, feature_size, requires_grad=True, device="cpu")
opt = optim.SGD([full_params,], lr=1)

def acc_grads(full_params, slice):
  param_subset = full_params[slice]
  param_subset = param_subset.cuda()
  # Compute that part of the loss
  local_loss = your_fn(param_subset)
  # Accumulate into full_params.grad
  local_loss.backward()
  # Use .item() to drop gradient history and get a python number
  return local_loss.item()

loss = 0.
# reset the gradients
opt.zero_grad()
for slice in slices:
  # Nothing outside `acc_grad` is on GPU. So the all the GPU memory will be freed
  # after it returns (use torch.cuda.memory_allocated() to see the used memory)
  partial_loss = acc_grad(full_params, slice)
  loss += partial_loss
opt.step()
print("Did a step for a total loss of {}".format(loss))
</code></pre>
<p>Note that with this, the gradients will only live on the CPU and the optimizer will work with CPU Tensors.<br>
Also because we send the Tensor to cuda <em>after</em> slicing it, during the backward, the full gradient Tensor will only exist on the CPU.</p>
<p>Does that match what you want better?</p>