I did not know about checkpointing. It looks promising. Thank you



 mrshenli:

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?


Sorry, I didn’t explain my training well. Each graph has up to 8 nodes. A full forward pass for a single node + the detached messages from the other nodes takes about 8GBs of memory. I go through the nodes using a for loop. I have to call backward() at the end of each iteration to calculate the loss, so it does get accumulated at each iteration. I do a single optimizer.step() after the loop to update the parameters. I technically do fwd-bwd x [1 - 8], into step().
From my understanding of DDP, the gradient reduction across multiple processes starts with the backward() call. That’s why I thought it wouldn’t work if different processes called it different amount of times. I did not know about no_sync() either. I’ll have to read about it. Thanks.