Kernel dies on loss.backward()

jamesbriggs: <p>Really cool to see MPS support in PyTorch! I have been testing it and it all works great until it comes trying to fine-tune a BERT model from HF. I have a simple training loop that looks like:</p>
<pre><code class="lang-python">model = BertForSequenceClassification('bert-base-uncased')

model.train()

optim = torch.optim.Adam(model.parameters(), lr=5e-5)

model.to(device)

loop = tqdm(loader, leave=True)
for batch in loop:
    batch_mps = {
        'input_ids': batch['input_ids'].to(device),
        'attention_mask': batch['attention_mask'].to(device),
        'labels': batch['labels'].to(device)
    }
    optim.zero_grad()
    outputs = model(**batch_mps)
    loss = outputs[0]
    loss.backward()
    optim.step()
</code></pre>
<p>As soon as it hits the <code>loss.backward()</code> step the kernel dies. I tried minimizing the <code>batch_size</code> (although I figure with unified memory it can handle the same batch size as CPU?), but no luck.</p>
<p>It works as expected on CPU, am I missing something?</p>

mxahan: <p>Can you please check your memory of GPU or CPU during the execution?</p>

jamesbriggs: <p>Thanks! The problem was simply that the batches, despite being incredibly small are still too large for my first gen M1. It begins to manage with a batch_size of 1. Thanks for your help!</p>

philipturner: <p><a class="mention" href="/u/alband">@albanD</a> could you modify the backend so that it fails gracefully when a <code>MTLBuffer</code> is initialized to <code>nil</code> because the computer ran out of memory?</p>

albanD: <p>Hi,</p>
<p>There is already a check there for the <code>MTLBuffer</code>: <a href="https://github.com/pytorch/pytorch/blob/4428218945e797cfc71a93dbb2d165535ea5a85b/aten/src/ATen/mps/MPSAllocator.mm#L54-L56" class="inline-onebox">pytorch/MPSAllocator.mm at 4428218945e797cfc71a93dbb2d165535ea5a85b · pytorch/pytorch · GitHub</a></p>
<p>But the problem with shared memory is that any CPU allocation might start failing as well. And these we don’t always control and can lead to hard crash.</p>

philipturner: <p>Tracing the code you cited, I saw something interesting. It’s not really relevant to this thread. I mentioned it in <a href="https://github.com/pytorch/pytorch/issues/77799#issuecomment-1132237526" class="inline-onebox" rel="noopener nofollow ugc">MPS device appears much slower than CPU on M1 Mac Pro · Issue #77799 · pytorch/pytorch · GitHub</a>, but let’s keep discussion on this forums thread for now.</p>
<p>It seems that you indeed use heap-backed memory, something I thought of myself to allow for zero-cost allocation: <a href="https://github.com/pytorch/pytorch/blob/09be44de7b56495bcb5ad1d47376200cbb853097/aten/src/ATen/mps/MPSAllocator.h#L81" class="inline-onebox" rel="noopener nofollow ugc">pytorch/MPSAllocator.h at 09be44de7b56495bcb5ad1d47376200cbb853097 · pytorch/pytorch · GitHub</a>. Could you go a little more into detail for how you came up with that allocation heuristic? Is there any parallel in the CUDA backend? Did you try an exponentially increasing idea such as the one described in <a href="https://github.com/AnarchoSystems/DeepSwift/issues/1#issuecomment-1129891796" class="inline-onebox" rel="noopener nofollow ugc">Sharing ideas about our work. · Issue #1 · AnarchoSystems/DeepSwift · GitHub</a>? It seems that you tried to make a heap be the same size as the buffer. That would make sense if you repeatedly recreate tensors of the same size, so you have heaps hanging around from previous allocations to reuse. But how do you know when enough pre-existing heaps is enough and start deleting previously allocated heaps before you run out of memory?</p>
<p>I haven’t tested my theory yet. If I could reuse your algorithm and the time you spent investigating this performance problem, that would be a big help for my personal ML project. I would make a documentation comment giving credit to PyTorch for coming up with the idea first.</p>

albanD: <p>Hi,</p>
<p>The allocator that we use for MPS is based on the CUDA Caching allocator that we already had.<br>
It is not as fully featured yet but you can find some high level doc at <a href="https://pytorch.org/docs/master/notes/cuda.html#memory-management" class="inline-onebox">CUDA semantics — PyTorch master documentation</a> and more detailed comments in the code itself <a href="https://github.com/pytorch/pytorch/blob/e2eb7a1edccfe9788edc50b6e8cddd62c2afff7a/c10/cuda/CUDACachingAllocator.cpp#L32" class="inline-onebox">pytorch/CUDACachingAllocator.cpp at e2eb7a1edccfe9788edc50b6e8cddd62c2afff7a · pytorch/pytorch · GitHub</a></p>
<p>Note that this is a little bit copy-pasted right now for the MPS side and it will be refactored for the two to be closer once the MPS version is stable.</p>

philipturner: <p>I’ve fully translated the PyTorch MPS allocator from Objective-C to Swift, and it’s working quite well for me. I have many other unrelated optimizations, which remove intermediate tensors and reduce the need for zero-cost allocation. That makes the heap allocator quite overkill for my framework. (1) You’re barely even allocating memory and (2) when you are, it’s instantaneous.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/philipturner/metal-experiment-1/blob/22a178547af5caaa37287a239b72f225bc0e1f66/Sources/MetalExperiment1/Core/HeapAllocator.swift">
  <header class="source">

      <a href="https://github.com/philipturner/metal-experiment-1/blob/22a178547af5caaa37287a239b72f225bc0e1f66/Sources/MetalExperiment1/Core/HeapAllocator.swift" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/philipturner/metal-experiment-1/blob/22a178547af5caaa37287a239b72f225bc0e1f66/Sources/MetalExperiment1/Core/HeapAllocator.swift" target="_blank" rel="noopener nofollow ugc">philipturner/metal-experiment-1/blob/22a178547af5caaa37287a239b72f225bc0e1f66/Sources/MetalExperiment1/Core/HeapAllocator.swift</a></h4>


      <pre><code class="lang-swift">//
//  HeapAllocator.swift
//  
//
//  Created by Philip Turner on 7/10/22.
//

import Metal

// Implementation based on PyTorch's MPS allocator:
// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/mps/MPSAllocator.h

fileprivate let megabyte = 1024 * 1024
fileprivate let kMaxSmallAlloc = megabyte
fileprivate let kMinLargeAlloc = 10 * megabyte
fileprivate let kSmallHeap = 8 * megabyte
fileprivate let kLargeHeap = 32 * megabyte
fileprivate let kRoundLarge = 2 * megabyte

protocol AllocatorBlockProtocol {
</code></pre>



  This file has been truncated. <a href="https://github.com/philipturner/metal-experiment-1/blob/22a178547af5caaa37287a239b72f225bc0e1f66/Sources/MetalExperiment1/Core/HeapAllocator.swift" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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

<p>I have intensely stress-tested this, and it works wonderfully. Great job! I would suggest a few cleanups, like removing the code for “splits” inside a heap. That code is unutilized and seems copied from the CUDA allocator, which must have manual heap placement. I removed it from my translation. However, it should be useful in my OpenCL backend for S4TF, which requires manual placement of sub-allocations within a larger “heap”. I’ll have to look at your CUDA allocator when I get around to my OpenCL backend.</p>
<p>Another optimization is something that happens when you exceed system memory limits. Unlike PyTorch, I wait to encode/submit operations until the GPU is almost starved of work. That means I could have ~100 operations queued up. If each op allocates a sizable chunk of memory, that could grow astronomical. To solve this problem, I flush the operation queue when memory allocation exceeds system RAM, waiting until more memory can be released. This is one step further than your allocator’s automatic purging of cached <code>BufferBlock</code>’s.</p>
<p>When you exceed system RAM size and flushing the operation queue doesn’t reduce total allocated memory, this switches to a mode labeled <code>permitExceedingSystemRAM</code>. The framework knows that you’re allocating absurd amounts of memory, and stops flushing the operation queue after every tensor materialization. In other words, the optimization permits a reasonable amount of CPU-side encoding performance in such a situation.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/philipturner/metal-experiment-1/blob/41c0f10a054b7db9bc4dc94a604be3f1658e7a82/Sources/MetalExperiment1/Core/Memory.swift#L229-L281">
  <header class="source">

      <a href="https://github.com/philipturner/metal-experiment-1/blob/41c0f10a054b7db9bc4dc94a604be3f1658e7a82/Sources/MetalExperiment1/Core/Memory.swift#L229-L281" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/philipturner/metal-experiment-1/blob/41c0f10a054b7db9bc4dc94a604be3f1658e7a82/Sources/MetalExperiment1/Core/Memory.swift#L229-L281" target="_blank" rel="noopener nofollow ugc">philipturner/metal-experiment-1/blob/41c0f10a054b7db9bc4dc94a604be3f1658e7a82/Sources/MetalExperiment1/Core/Memory.swift#L229-L281</a></h4>



    <pre class="onebox"><code class="lang-swift">
      <ol class="start lines" start="229" style="counter-reset: li-counter 228 ;">
          <li>@inline(never)</li>
          <li>private func actuallyMaterialize(checkingMemoryBounds: Bool = true) throws {</li>
          <li>  if checkingMemoryBounds {</li>
          <li>    let device = Context.global.device</li>
          <li>    let allocatedSize = HeapAllocator.totalAllocatedMemory</li>
          <li>    if Context.global.permitExceedingSystemRAM {</li>
          <li>      // Give it some wiggle room to remain in `permitExceedingSystemRAM` mode. Maximum buffer</li>
          <li>      // length should be &gt;50% system memory size. If it's hovering above the system RAM size</li>
          <li>      // because all that memory needs to exist, it won't suddenly deallocate upon flushing the</li>
          <li>      // command stream. In that case, flushing the command stream constantly as it oscillates</li>
          <li>      // above and below a certain threshold would seriously degrade performance. But if it jumped</li>
          <li>      // the threshold because my backend queued up too many commands, most of the memory would</li>
          <li>      // quickly deallocate.</li>
          <li>      //</li>
          <li>      // On iOS, the threshold is different because `MTLDevice.recommendedMaxWorkingSetSize` does</li>
          <li>      // not exist.</li>
          <li>      #if os(macOS)</li>
          <li>      let threshold = device.maxBufferLength</li>
          <li>      #else</li>
          <li>      let threshold = (device.maxBufferLength * 7) / 8</li>
      </ol>
    </code></pre>


  This file has been truncated. <a href="https://github.com/philipturner/metal-experiment-1/blob/41c0f10a054b7db9bc4dc94a604be3f1658e7a82/Sources/MetalExperiment1/Core/Memory.swift#L229-L281" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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

<p>This would be a good optimization for PyTorch when someone’s trying to test the limits of memory allocation. You would need a way to halt execution until all submitted GPU commands are completed. Furthermore, I learned something interesting while doing this. On unified memory architectures (M1, maybe Intel), GPU memory pages to the disk when it runs out of space. I tested this by allocating and writing to 64 GB of Metal memory (my Mac has 32 GB RAM), and there were no runtime errors*. If people could utilize this virtual memory paging on M1 GPUs with PyTorch, it would allow 100’s of gigabytes of memory. CUDA cards are limited to 10 GB these days, but my Apple GPU has as much memory as my SSD hard drive!</p>
<blockquote>
<p>*When you set the buffer’s storage mode to <code>.storageModePrivate</code>, there are runtime errors. Metal prevents you from exceeding system RAM. This means discrete GPUs can’t page their memory to the disk.</p>
</blockquote>
<p>Also, in <code>MPSHeapAllocatorImpl.Malloc</code>, why does it check whether a buffer is <code>&lt;</code> the device’s max buffer size, but not <code>&lt;=</code>? That would mean if the largest possible buffer size was 16 GB, PyTorch users could only allocate 15.9999999 GB. In my prototype backend, someone can allocate all 16 GB if they want.</p>