kleingeo:

Should it not be reduced with gradient_as_bucket_view? Or am I just confused with the parameter?


The memory should be reduced when you specify gradient_as_bucket_view. After initializing DDP, can you run a few iterations and then look at the memory consumption? Also, typically looking at nvidia-smi is not a good indication of memory usage due to the caching allocator (see CUDA semantics — PyTorch 1.10.1 documentation). I’d suggest using torch.cuda.memory_allocated — PyTorch 1.10.1 documentation or torch.cuda.max_memory_allocated — PyTorch 1.10.1 documentation to track this.