Generalize DiLoCo to support Streaming #205
Conversation
facebook-github-bot
added
the
CLA Signed
tushar00jain
force-pushed
the
feature/streaming-diloco
branch
2 times, most recently
from
cb34ff1 to
35ca865
Compare
tushar00jain
force-pushed
the
feature/streaming-diloco
branch
from
35ca865 to
ab00d7d
Compare
Summary: - Add option to perform qunatized allreduce in torchft manager - Update user level API's for DiLoCo to also support Streaming DiLoCo -- it takes a list of modules now as input - Create a class `_StreamingDiLoCoFragment` used by DiLoCo to support streaming. Each fragment independently determines its schedule (when to send/sync). - Adding support for "alpha" and "tao" parameters from the paper are left as a TODO. Plan to add this in a separate PR. Test Plan: ``` $ pytest -vs torchft/local_sgd_integ_test.py $ pytest -vs torchft/local_sgd_test.py ```
tushar00jain
force-pushed
the
feature/streaming-diloco
branch
from
ab00d7d to
0f07f2d
Compare
| @@ -267,7 +291,9 @@ def shutdown(self, wait: bool = True) -> None: | |||
| self._manager.shutdown() | |||
| self._executor.shutdown(wait=wait) | |||
|
|
|||
| def allreduce(self, tensor: torch.Tensor) -> torch.futures.Future[torch.Tensor]: | |||
| def allreduce( | |||
| self, tensor: torch.Tensor, should_quantize: bool = False | |||
There was a problem hiding this comment.
I would also update some of the tests in manager_test.py to also include using the should_quantize flag. Can be done in a follow up PR
torchft/manager.py
Outdated
| except ImportError: | ||
| from torch import cuda | ||
|
|
||
| def allreduce_quantized( |
There was a problem hiding this comment.
nit: is this stub necessary? Can't we just have a constant like TRITON_AVAILABLE and then check that in the if statement in the implementation
There was a problem hiding this comment.
Less configuration options 🥲 let's say someone changes platforms, then they can start using triton automatically without having to modify the constant. Also avoids having us to configure CI properly for different platforms.
| def __init__( | ||
| self, | ||
| manager: Manager, | ||
| model_fragments: List[nn.Module], |
There was a problem hiding this comment.
nit: maybe from an API / UX perspective we can support nn.Module | List[nn.Module] with the specification that passing in a single nn.Module means whole model.
There was a problem hiding this comment.
Think it's better to have a smaller api surface and avoid having a special case?
| model_fragment, | ||
| math.floor((sync_every / len(model_fragments)) * (i + 1)), | ||
| inner_optimizer, | ||
| # TODO: Support different outer optimizers for each fragment |
There was a problem hiding this comment.
oh interesting, is that mentioned in the paper?
There was a problem hiding this comment.
I think they should be different otherwise things like momentum end up being the same for all fragments? Maybe it's not very important though
tushar00jain
force-pushed
the
feature/streaming-diloco
branch
from
0f07f2d to
600864f
Compare
Summary:
_StreamingDiLoCoFragmentused by DiLoCo to support streaming. Each fragment independently determines its schedule (when to send/sync).Test Plan: