hos-b:

Model parallelism will probably also not work because even though the torch calls are async, I have to call backward() at the end of the for loop, creating a blocking bottleneck. The computation graph would otherwise get too large for a single GPU.


I see. Have you tried checkpointing some part of the autograd graph? This would drop the the activations and autograd graph in the forward and recompute them again in the backward. It’s like paying more compute to reduce memory footprint.
https://pytorch.org/docs/stable/checkpoint.html

I have to call backward() at the end of the for loop

Curious, is this because you need to compute the global loss across all iterations, so that you cannot run multiple fwd-bwd-fwd-bwd to accumulate grads (DDP can support this with no_sync() context manager) and then run optimizer.step() once to update parameters?