[minor] OSS: bring DDP in the benchmark (#130)

More realistic benchmarks, comparing apples to apples. DDP/OSS+DDP/OSS+SDP

.circleci/config.yml

@@ -100,7 +100,7 @@ run_oss_benchmark: &run_oss_benchmark
   - run:
       name: Run OSS Benchmark
       command: |
-        python benchmarks/oss.py --check_regression --world_size 4 --reference_speed 21.2 --reference_memory 4220 --reference_loss 0.63
+        python benchmarks/oss.py --check_regression --world_size 4 --reference_speed 13.7 --reference_memory 4390 --reference_loss 0.595
 
 run_oss_gloo: &run_oss_gloo
 - run:

benchmarks/oss.py

@@ -12,6 +12,7 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 from torchvision.datasets import FakeData
 from torchvision.models import resnet101
@@ -40,7 +41,9 @@ def get_problem(rank, data_size, batch_size):
         }
 
     dataloader = DataLoader(
-        dataset=FakeData(transform=ToTensor(), size=data_size), batch_size=batch_size, collate_fn=collate
+        dataset=FakeData(transform=ToTensor(), size=data_size, random_offset=rank),
+        batch_size=batch_size,
+        collate_fn=collate,
     )
     loss_fn = nn.CrossEntropyLoss()
     return model, dataloader, loss_fn
@@ -85,12 +88,13 @@ def train(
             optimizer=OPTIM,
             optimizer_params={"lr": 1e-4, "momentum": 0.9},
             world_size=world_size,
-            broadcast_buffers=False,
+            broadcast_buffers=True,
         )
         ddp.train()
         optimizer = ddp.optimizer
         model = ddp
     else:
+        model = DDP(model, device_ids=[rank], find_unused_parameters=True)  # type: ignore
         optimizer = (
             OSS(params=model.parameters(), optim=OPTIM, lr=1e-4, momentum=0.9)
             if use_oss
@@ -216,7 +220,7 @@ if __name__ == "__main__":
         )
 
     if args.optim_type == OptimType.oss or args.optim_type == OptimType.everyone:
-        print("\nBenchmark OSS")
+        print("\nBenchmark OSS with DDP")
         mp.spawn(
             train,
             args=(
@@ -237,7 +241,7 @@ if __name__ == "__main__":
         )
 
     if args.optim_type == OptimType.oss_sdp or args.optim_type == OptimType.everyone:
-        print("\nBenchmark OSS DDP")
+        print("\nBenchmark OSS with SDP")
         mp.spawn(
             train,
             args=(
@@ -248,7 +252,7 @@ if __name__ == "__main__":
                 backend,
                 True,  # OSS
                 True,  # SDP
-                args.check_regression,
+                False,  # FIXME: @lefaudeux - SDP should give the same results
                 -1,  # Not checking SDP for speed regression for now, still slower than OSS
                 args.reference_memory,
                 args.reference_loss,