Repository: facebookresearch/fairscale

Issue:
[feat] OSS: Support nvidia's LARC
## 🚀 Feature
Make it possible to support LARC with OSS

## Motivation
LARC is a must have for large batch jobs, right now OSS will break on LARC because of the closure() being passed

## Pitch
Should be doable to gracefully handle optimizers with do not support closures in step()

## Alternatives
Not supporting LARC, reduces a lot of OSS interest

## Additional context

cc @mannatsingh @prigoyal @msbaines 


Discussion:
SGTM

But if you are using LARC to wrap Adam then you may want to consider using the [fused implementation](https://github.com/NVIDIA/apex/blob/master/apex/optimizers/fused_lamb.py) of [LAMB](https://arxiv.org/pdf/1904.00962.pdf) available in apex. It is implemented as a fused multi-tensor CUDA kernel so will run orders of magnitude faster than LARC, implemented in python, wrapping Adam.

If you are wrapping something other than Adam, you may obtain a noticeable speedup implementing a fused multi-tensor kernel. Something like this could be added to fairscale.
The need for now (classy/vissl) is mostly around SGD, using [this](https://github.com/NVIDIA/apex/blob/master/apex/parallel/LARC.py) as a wrap, which breaks with the `closure` keyword argument. I'll keep the rest of Apex in mind, we should probably converge indeed

Resolving patch:
diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -103,8 +103,11 @@ def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) ->
         # Sync oss param_groups attributes in case they've been updated by a scheduler.
         self._sync_param_groups()
 
-        # Run the optimizer step on this shard only
-        loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
+        # Run the optimizer step on this shard only:
+        if closure is not None:
+            loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
+        else:
+            loss = self.optim.step(**kwargs)
 
         # Sync all the states. Broadcast requests are issued async, we check completeness before moving on
         requests = []


Test patch:
diff --git a/tests/optim/test_oss.py b/tests/optim/test_oss.py
--- a/tests/optim/test_oss.py
+++ b/tests/optim/test_oss.py
@@ -103,13 +103,12 @@ def test_lr_scheduler():
         assert x == x2
 
 
-class SGDWithStepKWArg(torch.optim.SGD):
-    def step(self, closure=None, kwarg=[]):
-        super().step()
-        kwarg.append(5)
-
-
 def test_step_with_kwargs():
+    class SGDWithStepKWArg(torch.optim.SGD):
+        def step(self, closure=None, kwarg=[]):
+            super().step()
+            kwarg.append(5)
+
     kwarg = []
     x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
     o = optim.OSS([x], SGDWithStepKWArg, lr=0.1)
@@ -119,6 +118,18 @@ def test_step_with_kwargs():
     assert x == torch.tensor([0.9], device=DEVICE)
 
 
+def test_step_without_closure():
+    class SGDWithoutClosure(torch.optim.SGD):
+        def step(self):
+            return super().step()
+
+    x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
+    o = optim.OSS([x], SGDWithoutClosure, lr=0.1)
+    x.backward()
+    o.step()
+    assert x == torch.tensor([0.9], device=DEVICE)
+
+
 def test_local_state_dict():
     x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
     o = optim.OSS([x], lr=0.1)