Repository: facebookresearch/fairscale

Issue:
oss_ddp will assert in forward when running in eval mode
## 🐛 Bug

In training, we only forward which will cause the need_reduction check to fail. This can be fixed setting need_reduction in a backward hook instead of forward.

## Command

## To Reproduce

Steps to reproduce the behavior:

<!-- If you were running a command, post the exact command that you were running -->

1.
2.
3.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from the
environment collection script from PyTorch
(or fill out the checklist below manually).

You can run the script with:
```
# For security purposes, please check the contents of collect_env.py before running it.
python -m torch.utils.collect_env
```

 - PyTorch Version (e.g., 1.0):
 - OS (e.g., Linux):
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version:
 - GPU models and configuration:
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Thanks for catching. Just to make sure, is this blocking anything? If not, I will do it later and add a unit test for it too.
It's not highpri, but this blocks replacing fairseq's LegacyDistributedDataParallel with the one in fairscale. Here is the fix we did in fairseq: https://github.com/pytorch/fairseq/commit/631023200a8e0c8531ccdb9e5d7f54411149efdf
Thanks, Myle! Will try to get this fix this week.

Resolving patch:
diff --git a/fairscale/nn/data_parallel/oss_ddp.py b/fairscale/nn/data_parallel/oss_ddp.py
--- a/fairscale/nn/data_parallel/oss_ddp.py
+++ b/fairscale/nn/data_parallel/oss_ddp.py
@@ -99,6 +99,15 @@ def __getstate__(self) -> Dict:
         attrs = copy.copy(self.__dict__)
         return attrs
 
+    def train(self, mode: bool = True) -> "OssDdp":
+        pre_mode = self.module.training
+        self.module.train(mode)
+        if self.module.training:
+            assert not self.need_reduction or pre_mode, "incorrect state transition"
+        else:
+            assert not self.need_reduction, "try to enter eval with grads unreduced"
+        return self
+
     @contextmanager
     def no_sync(self) -> Generator:
         """A context manager to disable gradient synchronization."""
@@ -108,10 +117,11 @@ def no_sync(self) -> Generator:
         self.accumulate_grads = old_accumulate_grads
 
     def forward(self, *inputs: Any, **kwargs: Any) -> Tensor:
-        if self.need_reduction:
-            raise RuntimeError("OssDdp requires explicit reduction, must call OssDdp.reduce")
-        if not self.accumulate_grads:
-            self.need_reduction = True
+        if self.module.training:
+            if self.need_reduction:
+                raise RuntimeError("OssDdp requires explicit reduction, must call OssDdp.reduce")
+            if not self.accumulate_grads:
+                self.need_reduction = True
         return self.module(*inputs, **kwargs)
 
     def reduce(self) -> None:
@@ -119,6 +129,7 @@ def reduce(self) -> None:
         This function must be called explicitly after backward to reduce
         gradients. There is no automatic hook like c10d.
         """
+        assert self.module.training, "Cannot call reduce in eval"
 
         def reduce_params(params: List[Parameter], params_rank: int) -> None:
             """ Helper to reduce a list of params that should fix in the buffer. """


Test patch:
diff --git a/tests/nn/data_parallel/test_oss_ddp.py b/tests/nn/data_parallel/test_oss_ddp.py
--- a/tests/nn/data_parallel/test_oss_ddp.py
+++ b/tests/nn/data_parallel/test_oss_ddp.py
@@ -54,3 +54,32 @@ def run_one_step(rank, world_size, backend, device, temp_file_name):
 def run_test(backend, device, world_size=2):
     temp_file_name = tempfile.mkstemp()[1]
     mp.spawn(run_one_step, args=(world_size, backend, device, temp_file_name), nprocs=world_size, join=True)
+
+
+def run_eval_mode(_unused):
+    """ Testing eval mode make sure this is no asserts. """
+    dist.init_process_group(
+        init_method=f"file://{tempfile.mkstemp()[1]}", backend=dist.Backend.GLOO, rank=0, world_size=1
+    )
+    model = Sequential(Linear(2, 3), Linear(3, 4))
+    optimizer = OSS(model.parameters(), lr=0.1, momentum=0.99)
+    ddp = OssDdp(model, optimizer, 1)
+
+    ddp.eval()
+    for _ in range(5):
+        input_tensor = torch.rand((64, 2))
+        output = ddp(input_tensor)
+
+    ddp.train()
+    try:
+        for _ in range(5):
+            input_tensor = torch.rand((64, 2))
+            output = ddp(input_tensor)
+    except RuntimeError:
+        pass
+    else:
+        assert False, "Multiple forward passes on training mode should not pass"
+
+
+def test_eval_mode():
+    mp.spawn(run_eval_mode, args=(), join=True)