manager: final changes for init_sync

Author: d4l3k

src/manager.rs

@@ -901,6 +901,7 @@ mod tests {
                     step: 0,
                     world_size: 1,
                     shrink_only: false,
+                    data: String::new(),
                 },
                 QuorumMember {
                     replica_id: "replica_1".to_string(),
@@ -909,6 +910,7 @@ mod tests {
                     step: 1,
                     world_size: 1,
                     shrink_only: false,
+                    data: String::new(),
                 },
                 QuorumMember {
                     replica_id: "replica_2".to_string(),
@@ -917,6 +919,7 @@ mod tests {
                     step: 0,
                     world_size: 1,
                     shrink_only: false,
+                    data: String::new(),
                 },
                 QuorumMember {
                     replica_id: "replica_3".to_string(),
@@ -925,6 +928,7 @@ mod tests {
                     step: 1,
                     world_size: 1,
                     shrink_only: false,
+                    data: String::new(),
                 },
                 QuorumMember {
                     replica_id: "replica_4".to_string(),
@@ -933,6 +937,7 @@ mod tests {
                     step: 0,
                     world_size: 1,
                     shrink_only: false,
+                    data: String::new(),
                 },
             ],
             created: None,

torchft/manager.py

@@ -34,7 +34,7 @@
 from contextlib import nullcontext
 from datetime import timedelta
 from enum import Enum
-from typing import Callable, cast, Dict, List, Optional, TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, TypeVar, cast
 
 import torch
 from torch.distributed import ReduceOp, TCPStore
@@ -106,6 +106,7 @@ def __init__(
         hostname: str = socket.gethostname(),
         heartbeat_interval: timedelta = timedelta(milliseconds=100),
         checkpoint_transport: Optional[CheckpointTransport[Dict[str, T]]] = None,
+        init_sync: bool = True,
     ) -> None:
         """
         Args:
@@ -143,6 +144,9 @@ def __init__(
             hostname: if rank==0, the hostname to advertise to the lighthouse server
             checkpoint_transport: the checkpoint transport to use for
                 transfering checkpoints to recovering replicas, defaults to HTTPTransport
+            init_sync: whether to synchronize the model weights on step 0. If
+                all of the model weights are initialized identically via
+                ``torch.set_seed`` you should set this to False.
         """
         self._load_state_dict = load_state_dict
         self._user_state_dict = state_dict
@@ -152,6 +156,7 @@ def __init__(
         self._quorum_timeout = quorum_timeout
         self._connect_timeout = connect_timeout
         self._world_size_mode = world_size_mode
+        self._init_sync = init_sync
 
         store_addr = store_addr or os.environ["MASTER_ADDR"]
         store_port = store_port or int(os.environ["MASTER_PORT"])
@@ -455,7 +460,7 @@ def _async_quorum(
             checkpoint_metadata=self._checkpoint_transport.metadata(),
             shrink_only=shrink_only,
             timeout=quorum_timeout,
-            init_sync=self.init_sync,
+            init_sync=self._init_sync,
         )
 
         quorum_id = quorum.quorum_id

torchft/manager_test.py

@@ -40,6 +40,7 @@ def _create_manager(
         min_replica_size: int = 2,
         world_size_mode: WorldSizeMode = WorldSizeMode.DYNAMIC,
         timeout: timedelta = timedelta(seconds=10),
+        init_sync: bool = True,
     ) -> Manager:
         pg = create_autospec(ProcessGroup)
         pg.errored.return_value = None
@@ -67,6 +68,7 @@ def _create_manager(
                 use_async_quorum=use_async_quorum,
                 world_size_mode=world_size_mode,
                 timeout=timeout,
+                init_sync=init_sync,
             )
             self.manager = manager
         return manager
@@ -617,7 +619,12 @@ def test_quorum_happy_timeouts(self, client_mock: MagicMock) -> None:
 
     @patch("torchft.manager.ManagerClient", autospec=True)
     def test_quorum_skip_init(self, client_mock: MagicMock) -> None:
-        manager = self._create_manager(use_async_quorum=False)
+        manager = self._create_manager(
+            use_async_quorum=False,
+            init_sync=False,
+        )
+
+        self.assertFalse(manager._init_sync)
 
         quorum = QuorumResult()
         quorum.quorum_id = 123
@@ -633,16 +640,8 @@ def test_quorum_skip_init(self, client_mock: MagicMock) -> None:
         client_mock()._quorum.return_value = quorum
 
         manager.start_quorum()
-        self.assertEqual(
-            client_mock()._quorum.call_args.kwargs["init_sync"], True
-        )
+        self.assertEqual(client_mock()._quorum.call_args.kwargs["init_sync"], False)
 
-        manager.start_quorum(init_sync=True)
-        self.assertEqual(
-            client_mock()._quorum.call_args.kwargs["init_sync"], True
-        )
-        
-        manager.start_quorum(init_sync=False)
-        self.assertEqual(
-            client_mock()._quorum.call_args.kwargs["init_sync"], False
-        )
+        manager._init_sync = True
+        manager.start_quorum()
+        self.assertEqual(client_mock()._quorum.call_args.kwargs["init_sync"], True)