Using Tensor subclasses for chunked loss and backprop?

cbcase: <p>We are looking at training transformers with both long sequences and large vocabs. When you do so, there is memory pressure specifically around the loss, since you materialize the full logits tensor of size <code>(seqlen, vocab)</code> (assuming batch=1) – and furthermore, you end up with a few tensors of this size, due to the usual cast-to-fp32 as well as the activation grads for those two tensors.</p>
<p>One option (discussed in <a href="https://github.com/pytorch/pytorch/issues/124480" class="inline-onebox" rel="noopener nofollow ugc">Fused Linear and Cross-Entropy Loss `torch.nn.functional.linear_cross_entropy` · Issue #124480 · pytorch/pytorch · GitHub</a>) is to fuse the final linear projection and the loss function. A simpler alternative is to chunk the last hidden activations (<em>before</em> the final linear projection) and compute the loss per-chunk. Doing so requires breaking the autograd graph but is otherwise straightforward. In almost-pytorch pseudocode:</p>
<pre data-code-wrap="python"><code class="lang-python"># h: Tensor of size (seqlen, hidden)  -- these are final hidden activations before last linear
# final_linear: Module for the last linear that projects hidden -&gt; vocab
# targets: Tensor of size (seqlen) with the true labels
h_detached = h.detach().requires_grad_()
total_loss = torch.zeros([], device="cuda")
for start, end in chunks:
    logits = final_linear(h_detached[start:end, ...]).float()
    loss = loss_fn(logits, targets[start:end])
    total_loss += loss.detach()
    loss.backward()
# Now continue backprop where you left off
h.backward(h_detached.grad)
# Note that we _haven't_ normalized the loss by len(chunks), which you likely want to do.
</code></pre>
<p>This works great as far as it goes. What is tricky is that most real codebases have code that broadly looks like:</p>
<pre><code class="lang-auto">loss = loss_fn(model(x), y)
loss = loss * some_factor  # or divide, or mutate in some other way
loss.backward()
</code></pre>
<p>where <code>loss</code> is a scalar Tensor. This relies on a single, unbroken autograd graph from the loss to all the leaves.</p>
<p>My sense is that you could map the code above onto the pattern below (“<code>loss</code> is a scalar Tensor”) using a Tensor subclass. That subclass tracks the extra state (the Tensor <code>h</code> to continue backprop from and the Tensor <code>h_detached.grad</code> that represents so-far-accumulated backprop up to that point) and hooks into backprop to make <code>loss.backward()</code> work. But I can’t figure out from the docs where to begin / which pattern of Tensor subclassing this maps onto. Tagging <a class="mention" href="/u/alband">@albanD</a> for tensor subclass value. Thanks!</p>

albanD: <p>Hey!</p>
<blockquote>
<p>where <code>loss</code> is a scalar Tensor. This relies on a single, unbroken autograd graph from the loss to all the leaves.</p>
</blockquote>
<p>Note that the backward is linear. So doing this multiplication by a factor within the loop above will actually give you the same result. This only works of course if you don’t need the full loss to compute this factor!</p>
<p>I’m afraid I don’t recall exactly where the overhead is here but I’m not sure a subclass will help you.<br>
I would suggest a custom autograd.Function though so that you can control both the fw and bw and hold onto the pieces you need.<br>
But you most likely will, at best, get to the same behavior as here <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>
<p>If you need the full loss before starting to backprop and want to keep the low saved buffers, I think you will have to recompute the forward of the loss twice. But if you’re ok with that, you custom Function can just compute the full loss as usual in the forward() (which runs in nograd mode). And then in the backward, you recompute and do small backwards per-slice.</p>