Avoid unnecessary DDP synchronization when gradient_accumulation_steps > 1

[noamwies]

🚀 Feature

Avoid unnecessary DDP synchronization when gradient_accumulation_steps > 1

Motivation

When training large models the synchronization is costly and the actual speedup from 2 gpus is much lower than 200%

Pitch

We can use DDP no_sync feature to avoid synchronization in steps that doesn't call optimizer_step

[github-actions]

Hi! thanks for your contribution!, great first issue!

[justusschock]

Hi @noamwies ,

Thanks for the suggestion. However, after looking into this, I think that this is not necessary, since under the hood we don't accumulate gradients, but loss values (on each device independently) and thus there are no ddb-syncs, since backward actually is not called before the accumulation is finished.

[noamwies]

Hi @justusschock

Regardless the gradients accumulation mechanism that is unclear to me, I think this feature still necessary because under the hood in each forward pass LightningDistributedDataParallel synchronized the model parameters although they changed only when (self.trainer.batch_idx + 1) % self.trainer.accumulate_grad_batches == 0

[noamwies]

@justusschock

To give concrete example i attached the following simple benchmark on rtx2080 ti GPU's

	1 GPU	2 GPU	2 GPU no sync
iters per GPU per sec	9.41	7.17	9.24

Thus using 2 GPU's give 150% performance as opposite to 196% when I usingno_sync

from argparse import ArgumentParser

from pl_bolts.models import ImageGPT
from pl_bolts.datamodules.binary_mnist_datamodule import BinaryMNISTDataModule
from pytorch_lightning import Trainer
from torchvision import transforms


def parse_args():
    parser = ArgumentParser()
    parser = Trainer.add_argparse_args(parser)
    parser.add_argument('--layers', type=int, default=2)
    parser.add_argument('--heads', type=int, default=2)
    parser.add_argument('--embed_dim', type=int, default=16)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--num_workers', type=int, default=8)
    parser.add_argument('--image_size', type=int, default=16)
    parser.add_argument('--vocab_size', type=int, default=50000)    
    return parser.parse_args()


def main(args):
    mnist_transforms = transforms.Compose([
                transforms.CenterCrop(args.image_size),
                transforms.ToTensor()
            ])
    dm = BinaryMNISTDataModule(data_dir='.', num_workers=args.num_workers)
    dm.train_transforms  = mnist_transforms
    dm.val_transforms = mnist_transforms
    model = ImageGPT(dm, embed_dim=args.embed_dim, heads=args.heads, vocab_size=args.vocab_size,
        layers=args.layers, batch_size=args.batch_size, num_workers=args.num_workers)
    trainer = Trainer.from_argparse_args(args)
    trainer.fit(model)

if __name__ == '__main__':
    main(parse_args())

python main.py --gpus 1  --precision 16 --batch_size 1 --accumulate_grad_batches 32 --layers 6 --embed_dim 512 --heads 8

python main.py --gpus 2 --distributed_backend=ddp  --precision 16 --batch_size 1 --accumulate_grad_batches 32 --layers 6 --embed_dim 512 --heads 8

[justusschock]

Have you tried the solution in the linked PR? For comparison?

[noamwies]

I tested it now and It seem the solution from the PR has minor effect on the runtime.

	1 GPU	2 GPU	2 GPU PR 4146	2 GPU no sync
iter per GPU per sec	9.41	7.17	7.21	9.17

@justusschock
Are you sure the backward is not called before the accumulation is finished ? that doesn't make sense to me since without backward all steps except the last one doesn't contribute to the gradients

[justusschock]

@noamwies They do contribute, since the losses are summed up with computational graph. So backprop is only done once, but for the sum of all losses of the previous steps which is equal to summing up gradients.

cc @awaelchli for verification. He's our DDP master :D

[noamwies]

@justusschock
The thing is that IMO the gradient accumulation mechanism you described doesn't save GPU memory, the reason is that you must saving all steps forward activation until the last step backward.

Finally when all steps activation can fit into the GPU memory I see no reason to use gradient_accumulation_steps > 1

[justusschock]

@noamwies You're right. I misunderstood a small part of the code. I checked again with @awaelchli and can now confirm your understanding. We can fix this directly after #4190 is done. I'll ping you here to test, once I got to it for confirmation :)

[noamwies]

Thanks