[bugfix] OSS no reduce loss (#133)

* bugfix
* adjust default non-regression loss, not all_reduced now

.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 13.7 --reference_memory 4390 --reference_loss 0.595
+        python benchmarks/oss.py --check_regression --world_size 4 --reference_speed 13.7 --reference_memory 4390 --reference_loss 0.152
 
 run_oss_gloo: &run_oss_gloo
 - run:

benchmarks/oss.py

@@ -124,8 +124,6 @@ def train(
                 loss /= world_size
                 loss.backward()
 
-                dist.all_reduce(loss, op=dist.ReduceOp.SUM)
-
                 if use_sdp:
                     ddp.reduce()  # Send the gradients to the appropriate shards
 

docs/source/tutorials/oss.rst

@@ -42,7 +42,6 @@ Let's suppose that your trainer looks like
                 loss = loss_fn(outputs, target)
                 loss /= world_size
                 loss.backward()
-                torch.distributed.all_reduce(loss, op=torch.distributed.ReduceOp.SUM)
                 optimizer.step()
 
 
@@ -90,7 +89,6 @@ Then sharding the optimizer state is merely a matter of wrapping your optimizer
                 loss = loss_fn(outputs, target)
                 loss /= world_size
                 loss.backward()
-                torch.distributed.all_reduce(loss, op=torch.distributed.ReduceOp.SUM)
                 optimizer.step()