Sholto
Theme

Themes inspired by:

Blog · Research infrastructure ·

When the batch doesn't fit: recursive splitting under CUDA OOM

A training loop that catches a CUDA out-of-memory error, halves whatever did not fit and tries again, and what going back to read it turned up: an except clause that catches far more than OOM, and a queue that quietly does not accumulate the gradient it looks like it should.

  • Interactive
  • pytorch
  • cuda
  • out-of-memory
  • training
  • research-infrastructure

An image tensor has a fixed shape and a predictable memory cost: a 224x224x3 batch of eight costs the same whether the pictures are a cat or a wall. A point cloud does not. One ScanNet room handed to the convolution I was benchmarking might be 20,000 points; the next might be 300,000, and there is no way to know which until you have already tried to run it through the model. A batch of four scenes that trained fine on Tuesday can run out of GPU memory on Wednesday because the fifth epoch happened to draw a bigger room. The lazy fix is to pick a batch size conservative enough to survive the worst room in the dataset and waste memory on every batch that is not that room. The fix I actually wrote, Trainer.recursive_train in the same experiment bed as the convolution post above, catches the failure instead: when a batch does not fit, split it and try the pieces separately, recursing until either everything fits or a piece is too small to be worth the trouble.

A queue, not really a recursion

The method is called recursive_train, but the “recursion” in it is a plain while loop over a list used as a queue, data_objs.pop(0) off the front, data_objs.append(...) onto the back:

def recursive_train(self, data):
    data_objs = [(data, True)]
    ret = []
    r_loss = None
    while len(data_objs) > 0:
        data, split_on_x = data_objs.pop(0)
        try:
            ...

Nothing here actually calls itself. A batch that needs three splits to fit just means the while loop runs a few extra times, popping whatever the previous iteration appended. It is recursive in the sense that matters, “the same procedure applied again to a smaller piece of the same problem”, just implemented with a list rather than the call stack, which is also the more sensible choice: three nested Python stack frames per split would work fine for a point cloud, but there is no reason to pay for them.

Each entry in the queue carries a split_on_x flag alongside the data, and that flag is what decides how a failure gets handled:

            except Exception as e:
                print(e)
                split_on_batch = (data.batch.cpu().detach().min() != data.batch.cpu().detach().max()).item()
                str_spl = 'batch' if split_on_batch else 'x' if split_on_x else 'y'
                print(f"When training, ran out of memory. Splitting on '{str_spl}' and retrying")
                split_axis = 0 if split_on_x else 1

data.batch is PyTorch Geometric’s standard bookkeeping tensor, one entry per point, saying which scene in the batch that point came from. If it still holds more than one distinct value, this fragment is a whole batch of several scenes stacked together, and the cheap thing to try first is splitting by scene. Only once a fragment is down to a single scene does it fall back to bisecting that scene’s points in space.

Splitting the batch, then splitting the scene

Batch splitting iterates the unique scene ids and slices every tensor field whose leading shape matches the per-point mask, resetting each new fragment’s batch column to a single id since it is now, by construction, one scene:

                if split_on_batch:
                    v_hold = {}
                    for batch_nr in np.unique(data.batch).tolist():
                        msk = data.batch == batch_nr
                        for k in data.keys:
                            d_k = data[k]
                            if isinstance(d_k, torch.Tensor) and d_k.shape[:len(msk.shape)] == msk.shape:
                                v_hold[k] = d_k[msk]
                            else:
                                v_hold[k] = d_k
                        v_hold['batch'][:] = 0
                        data_objs.append((Batch(**v_hold), split_on_x))

Every scene gets requeued here, however small, no floor applied. The floor only shows up once there is nothing left to split by scene, because a whole scene, unlike an arbitrary fragment of one, is never so small that skipping it is the right call: it is a full, labelled training example, and dropping it loses a real piece of the dataset for that epoch. Once a fragment holds one scene, the code instead bisects it spatially, at the midpoint of its bounding box on one axis:

                else:
                    max_p = data.pos[..., split_axis].max()
                    min_p = data.pos[..., split_axis].min()
                    mid_p = (max_p + min_p) / 2.0
                    mask = data.pos[..., split_axis] < mid_p
                    v1, v2 = {}, {}
                    for k in data.keys:
                        d_k = data[k]
                        if isinstance(d_k, torch.Tensor) and d_k.shape[:len(mask.shape)] == mask.shape:
                            v1[k] = d_k[mask]
                            v2[k] = d_k[torch.logical_not(mask)]
                        else:
                            v1[k] = d_k
                            v2[k] = d_k

                    if v1["pos"].size(-2) > 500:
                        data_objs.append((Batch(**v1), not split_on_x))
                    else:
                        print(f"Segment shape has {v1['pos'].size(-2)} points (<500). Skipping")
                    if v2["pos"].size(-2) > 500:
                        data_objs.append((Batch(**v2), not split_on_x))
                    else:
                        print(f"Segment shape has {v2['pos'].size(-2)} points (<500). Skipping")

Two things here are easy to misread on a skim, and I misread both of them the first time I went back through this for the post. It is a midpoint-of-the-bounding-box split, not a median-of-the-points split, so the two halves are equal in space, not necessarily in point count, a room with a dense cluster of points on one side of the midline hands back two very uneven fragments. And split_axis = 0 if split_on_x else 1 is not choosing the longest axis, it is just alternating between axis 0 and axis 1 on every spatial split, x, then y, then x again, however many times a stubborn scene needs bisecting. It never touches axis 2. I remembered writing “split on whichever axis has the most spread” and the code does not do that; it does something simpler that happens to work almost as well for a room-shaped point cloud, where the vertical extent is usually the smallest of the three anyway.

The except that catches too much

The line directly above the exception handler I have been quoting is a comment:

            # except RuntimeError as e:
            except Exception as e:

Somewhere in this repo’s history I had except RuntimeError, the narrower exception PyTorch actually raises for a CUDA allocation failure, and widened it to catch anything. I do not remember making that change on purpose, and going back through the surrounding commits does not tell me why either, only that the commented line is what came before. The effect is that this handler cannot tell a model that ran out of memory apart from a model that hit a shape mismatch, a NaN in a loss term, or any other bug entirely unrelated to how big the batch was. All three get the same response: print the exception, split whatever was being processed, retry the pieces. A genuine bug in the model does not get fixed by having fewer points to look at, so it keeps failing, keeps getting split, and keeps getting smaller, right down to the 500-point floor, at which point the last fragment is silently skipped and training moves on to the next batch as if nothing had happened. There is no traceback, because there never was an unhandled exception; there is just a training run that quietly gets a little slower and a little noisier every time that bug fires, with nothing in the logs pointing at why.

Try it: a memory budget and a queue that does not know why it failed

InteractiveGPU memory bisector

With JavaScript on, this becomes a simulated GPU memory budget and a queue of synthetic point-cloud scenes running the exact decision tree above: batch-split first if a fragment still holds more than one scene, bisect spatially once it does not, drop anything under 500 points. Step through the recursion tree one attempt at a time, watch retries, wasted forward passes and wall-clock overhead accumulate, switch the size distribution from uniform to heavy-tailed to see why this trick earns its keep exactly when a dataset has a few huge outliers in it, and flip the catch mode from “OOM only” to “everything” with a fake bug injected into one scene to watch the failure mode above happen for yourself: the same fragment, crashing loudly in one mode and dissolving into smaller batches with no error at all in the other.

What a retry actually costs

“Split the batch and retry” sounds like a free lunch, and it mostly is, but not entirely, and it is worth being precise about where the cost actually sits.

The obvious comparison is gradient accumulation: run several forward and backward passes over smaller chunks without stepping the optimiser, sum their gradients, and step once at the end, which is the standard way to simulate a larger batch than memory allows. Recursive splitting looks like it is doing this, several forward-backward passes happen inside one call to recursive_train before train_epoch calls self.optimiser.step(), but going back through the loop closely, it is not. self.optimiser.zero_grad() sits inside the try block, at the top of every iteration:

            try:
                if self.only_x_device:
                    data['x'] = data['x'].to(self.device)
                else:
                    data = data.to(self.device)
                self.optimiser.zero_grad()

which means every fragment that gets processed clears whatever gradient the previous fragment’s backward() had just accumulated, before computing its own. By the time the queue empties and train_epoch finally calls step(), only the gradient from whichever fragment happened to run last is still there to apply. The metrics for every fragment along the way get logged, ret is reassigned each pass and every after_batch callback fires, so the loss curve looks like it saw the whole original batch, but the weight update that follows only ever reflects the last piece the queue popped. That is a genuinely different trade than gradient accumulation makes: accumulation spends extra forward-backward passes to keep the same effective batch size, while this spends them and still trains on whatever the last split happened to be, smaller than what was asked for, and not particularly well chosen.

There is a second, quieter cost specific to point-cloud models, which lean on BatchNorm1d throughout this codebase to normalise per-point features. A halved fragment is not a random subsample of the room, it is one spatial half of it, a bounding-box bisection, so a BatchNorm1d layer computing its running statistics over that half sees a biased sample: whatever is architecturally different about “the near wall” versus “the far wall” of a room, rather than a smaller but still representative slice of the whole scene. Splitting doesn’t just make batch statistics noisier, which would even out on average over enough batches, it makes them spatially correlated in a way that ordinary batch sampling never is.

Both of those are why the 500-point floor matters more than it looks like it should. Without it, a sufficiently lopsided bounding-box split, one dense cluster on one side of the midline, empty space on the other, can chase a handful of stray points down through several levels of recursion, each one a full forward pass spent on a fragment that was never going to teach the model anything. BatchNorm1d itself is degenerate on a handful of points anyway; there is no meaningful variance to normalise by. The floor turns an unbounded-in-principle recursion into one that always terminates in a small, fixed number of extra steps, at the cost of occasionally losing a genuinely tiny sliver of a scene for that epoch, which is the right trade.

The last piece is the one the “OOM only versus catch everything” toggle above cannot fully capture, because it has to pick a deterministic memory budget to be legible at all: a real CUDA out-of-memory error is not a reliable signal to retry against. PyTorch’s caching allocator holds onto freed blocks for reuse rather than returning them to the driver immediately, so after a failure the allocator’s internal state is not guaranteed to be back to a clean baseline, fragmentation, or another tensor elsewhere in the same process still growing, can make an identical retry fail again even on a batch that is genuinely smaller than the one that just OOM’d. The code already does the two things you are supposed to do about this: move every CUDA-resident field back to the host before retrying,

                for key, item in data:
                    if torch.is_tensor(item) and item.is_cuda:
                        item_cpu = item.to('cpu')
                        del item
                        data[key] = item_cpu

                torch.cuda.empty_cache()

and only then ask the allocator to try again. That is the honest limit of what “catch and retry” can promise: it gives the allocator its best chance to recover, not a guarantee that it will.

Mixed precision along for the ride, and a profiler that never watched this path

One more thing rides inside that same try block. train.py opens with an optional import:

try:
    from apex import amp
except ImportError:
    amp = None

and Trainer.__init__ degrades gracefully if it fails:

        if amp is None:
            logger.warn("Could not find apex. Defaulting to off.")
            amp_off = True

NVIDIA’s Apex is a submodule reference, lib3rdparty/apex, that I never fetched into this checkout, so amp is None here and every run in this environment trains in full precision, loss.backward() rather than amp.scale_loss(loss, self.optimiser).backward(). Both branches sit inside the same try, which means a mixed-precision NaN, the specific failure mode Apex’s loss scaling exists to catch and recover from on its own, would also be swallowed by the same broad except above and handed to the batch-splitter as if it were a memory problem.

I also built a small profiler for this training loop, TrainTimingMetrics in expbed/util/metrics.py, wrapping codetiming’s Timer and bucketing named sub-timers by a prefix convention:

        transform_times = self.filter_times('transf.', '.call', elapsed_time/100.0)
        self.log_writer.add_scalars('transf_time', transform_times, global_step=step)

        data_load_times = self.filter_times('Data.', '.next', elapsed_time/100.0)
        self.log_writer.add_scalars('data_load_time', data_load_times, global_step=step)

Anything timed with a name starting transf. or Data. gets pulled out and logged to TensorBoard as its own series, which is how I watched data loading and per-transform cost during real training runs. Nothing in this repo ever wraps recursive_train itself in a named timer, so I have no logged number for what any of the retries above actually cost in wall clock on real hardware, which is exactly the gap the widget’s overhead counter is standing in for.

Recursive splitting is a small pattern, “when a job might not fit, halve it and try again, give up gracefully once the pieces are too small to matter”, and it generalises a long way past point clouds: any workload with input sizes that vary enough that no single static batch size is both safe and efficient can use the same shape of recovery. The two things worth carrying with it are the ones this post found on the way back through the code: catch only the failure you actually know how to recover from, and check, rather than assume, that a queue of retries is still accumulating what you think it is.