support none outer lr

samsja · samsja · commit 532af191af85 · 2024-08-22T09:41:02.000Z
diff --git a/open_diloco/hivemind_diloco.py b/open_diloco/hivemind_diloco.py
@@ -406,7 +406,7 @@ def __init__(
         )
         self.diloco_grad_averager = self._make_gradient_averager(compression=grad_compression)
 
-        self.outer_scheduler = outer_scheduler(self.state_averager.optimizer)
+        self.outer_scheduler = outer_scheduler(self.state_averager.optimizer) if outer_scheduler else None
 
     def _check_kwargs(self, kwargs) -> None:
         """DiLoCo Optimizer only support a subset of Hivemind Optimizer kwargs.

Original file line number	Diff line number	Diff line change
`@@ -406,7 +406,7 @@ def __init__(`
`406`	`406`	`)`
`407`	`407`	`self.diloco_grad_averager = self._make_gradient_averager(compression=grad_compression)`
`408`	`408`
`409`		`- self.outer_scheduler = outer_scheduler(self.state_averager.optimizer)`
	`409`	`+ self.outer_scheduler = outer_scheduler(self.state_averager.optimizer) if outer_scheduler else None`
`410`	`410`
`411`	`411`	`def _check_kwargs(self, kwargs) -> None:`
`412`	`412`	`"""DiLoCo Optimizer only support a subset of Hivemind Optimizer kwargs.`