Merge pull request #23 from tplr-ai/feat/validator_sync

Feat/validator sync

Author: distributedstatemachine

docker/compose.yml

@@ -35,7 +35,6 @@ services:
     image: containrrr/watchtower
     volumes:
       - /var/run/docker.sock:/var/run/docker.sock
-      - ${HOME}/.docker/config.json:/config.json:ro
     command: --interval 30 --cleanup --label-enable
     restart: unless-stopped
     environment:

docker/docker-compose-test.yml

@@ -1,9 +1,15 @@
+networks:
+  test:
+    driver: bridge
+
 services:
   miner1:
     build:
       context: ..
       dockerfile: docker/Dockerfile
     container_name: templar-miner-M111
+    networks:
+      - test
     volumes:
       - ~/.bittensor/wallets:/root/.bittensor/wallets
       - ../logs:/app/logs
@@ -27,22 +33,24 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              device_ids: [ '0', '1', '2' ]
+              device_ids: [ '3' ]
               capabilities: [ gpu ]
 
   miner2:
     build:
       context: ..
       dockerfile: docker/Dockerfile
     container_name: templar-miner-M222
+    networks:
+      - test
     volumes:
       - ~/.bittensor/wallets:/root/.bittensor/wallets
       - ../logs:/app/logs
     environment:
       NODE_TYPE: miner
       WALLET_NAME: Bistro
       WALLET_HOTKEY: M222
-      CUDA_DEVICE: cuda:1
+      CUDA_DEVICE: cuda:0
       NETWORK: test
       DEBUG: 'true'
       WANDB_API_KEY: ${WANDB_API_KEY}
@@ -58,36 +66,39 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              device_ids: [ '0', '1', '2' ]
+              device_ids: [ '1' ]
               capabilities: [ gpu ]
 
   validator:
     build:
       context: ..
       dockerfile: docker/Dockerfile
     container_name: templar-validator-V11
+    networks:
+      - test
     volumes:
       - ~/.bittensor/wallets:/root/.bittensor/wallets
       - ../logs:/app/logs
     environment:
       NODE_TYPE: validator
       WALLET_NAME: Bistro
       WALLET_HOTKEY: V11
-      CUDA_DEVICE: cuda:2
+      CUDA_DEVICE: cuda:0
       NETWORK: test
       DEBUG: 'true'
       WANDB_API_KEY: ${WANDB_API_KEY}
       NETUID: 268
-      HOST_CUDA_VERSION : 12.6
+      HOST_CUDA_VERSION: 12.6
       R2_ACCOUNT_ID: ${R2_ACCOUNT_ID}
       R2_READ_ACCESS_KEY_ID: ${R2_READ_ACCESS_KEY_ID}
       R2_READ_SECRET_ACCESS_KEY: ${R2_READ_SECRET_ACCESS_KEY}
-      R2_WRITE_ACCESS_KEY_ID : ${R2_WRITE_ACCESS_KEY_ID}
-      R2_WRITE_SECRET_ACCESS_KEY : ${R2_WRITE_SECRET_ACCESS_KEY}
+      R2_WRITE_ACCESS_KEY_ID: ${R2_WRITE_ACCESS_KEY_ID}
+      R2_WRITE_SECRET_ACCESS_KEY: ${R2_WRITE_SECRET_ACCESS_KEY}
+    restart: always
     deploy:
       resources:
         reservations:
           devices:
             - driver: nvidia
-              device_ids: [ '0', '1', '2' ]
+              device_ids: [ '2' ]
               capabilities: [ gpu ]

hparams.json

@@ -24,5 +24,8 @@
     "warmup_steps": 250,
     "alpha_f": 0.1,
     "t_max": 20000,
-    "validator_offset": 4
+    "validator_offset": 4,
+    "checkpoint_frequency": 50,
+    "topk_peers": 20,
+    "minimum_peers": 5 
 }

neurons/miner.py

@@ -109,7 +109,7 @@ def __init__(self):
         )
         cosine_scheduler = CosineAnnealingWarmRestarts(
             self.optimizer,
-            T_0=1000,
+            T_0=10000,
             T_mult=2,
             eta_min=self.hparams.learning_rate * 0.1,
         )
@@ -137,6 +137,11 @@ def __init__(self):
             hparams=self.hparams,
         )
 
+        self.bucket = self.comms.get_own_bucket()
+        self.comms.try_commit(self.wallet, self.bucket)
+        self.comms.fetch_commitments()
+
+
         # Init peers
         if not self.config.peers:
             self.peers = self.comms.peers
@@ -175,23 +180,49 @@ async def run(self):
         validator_uid, stake = self.comms.get_highest_stake_validator()
         if stake > 0:
             try:
-                state_dict = await self.comms.get(
-                    uid=str(validator_uid),
-                    window=self.current_window,
-                    key='checkpoint',
-                    timeout=240,
-                    local=False,
-                    stale_retention=10
-                )
-                if state_dict is not None:
-                    self.model.load_state_dict(state_dict)
-                    tplr.logger.info(f"Loaded checkpoint from validator {validator_uid} at window {self.current_window}")
+                # Calculate the most recent window that should have a checkpoint
+                expected_checkpoint_window = (self.current_window // self.hparams.checkpoint_frequency) * self.hparams.checkpoint_frequency
+
+                # Try last few windows in case of missed checkpoints
+                for window in range(expected_checkpoint_window, max(0, expected_checkpoint_window - 3 * self.hparams.checkpoint_frequency), -self.hparams.checkpoint_frequency):
+                    result = await self.comms.get(
+                        uid=str(validator_uid),
+                        window=window,
+                        key='checkpoint',
+                        timeout=240,
+                        local=False,
+                        stale_retention=10
+                    )
+                    if result is None:
+                        tplr.logger.debug(f"No checkpoint found for window {window}")
+                        continue
+
+                    checkpoint_data, global_step = result
+                    try:
+                        # Load state dicts from dictionary
+                        self.model.load_state_dict(checkpoint_data['model_state_dict'])
+                        self.optimizer.load_state_dict(checkpoint_data['optimizer_state_dict'])
+                        self.scheduler.load_state_dict(checkpoint_data['scheduler_state_dict'])
+                        self.momentum = checkpoint_data['momentum']
+                        self.global_step = checkpoint_data['global_step']
+                        
+                        # Update optimizer and scheduler steps to match
+                        self.optimizer._step_count = self.global_step
+                        self.scheduler.last_epoch = self.global_step
+                        
+                        tplr.logger.info(f"Loaded checkpoint from validator {validator_uid} at window {window}, global_step={self.global_step}")
+                        break  # Successfully loaded checkpoint, exit loop
+                    except KeyError as e:
+                        tplr.logger.error(f"Invalid checkpoint format: missing key {e}")
+                    except Exception as e:
+                        tplr.logger.error(f"Failed to load checkpoint: {e}")
                 else:
-                    tplr.logger.info("No checkpoint found, starting from scratch")
+                    tplr.logger.info("No valid checkpoints found in recent windows")
             except Exception as e:
                 tplr.logger.warning(f"Failed to load checkpoint: {e}")
         else:
             tplr.logger.info("No active validators found, starting from scratch")
+            self.global_step = 0
 
         # Start background block listener
         self.loop = asyncio.get_running_loop()
@@ -200,6 +231,7 @@ async def run(self):
             args=(self.loop,),
             daemon=True,
         ).start()
+        self.comms.start_commitment_fetcher()
 
         while True:
             step_window = self.current_window
@@ -317,20 +349,37 @@ async def run(self):
                 xshapes[n] = xshape
                 totalks[n] = totalk
 
-            # All-gather share state from peers
+            # Gather gradients from peers
             tplr.logger.info(f"Start gather: {self.peers}")
             gather_result = await self.comms.gather(
                 state_dict=gradient,
                 my_uid=self.uid,
                 uids=self.peers,
                 window=step_window,
                 key='gradient',
-                timeout=5,
+                timeout=20,
                 device=self.config.device,
                 local=False,
-                stale_retention=10
+                stale_retention=10,
+                global_step=self.global_step,
             )
-            
+
+            if gather_result is None:
+                tplr.logger.error("Failed to gather gradients from peers. Waiting for next window.")
+                # Wait for next window
+                while self.current_window == step_window:
+                    await asyncio.sleep(0.1)
+                continue  # Proceed to the next window
+
+            # Update self.global_step based on the maximum global_step received
+            max_global_step = max(gather_result.global_steps + [self.global_step])
+            if max_global_step > self.global_step:
+                tplr.logger.info(f"Updating global_step from {self.global_step} to {max_global_step}")
+                self.global_step = max_global_step
+                # Update optimizer and scheduler steps
+                self.optimizer._step_count = self.global_step
+                self.scheduler.last_epoch = self.global_step
+
             # Decompress state and apply to grad.
             for n, p in self.model.named_parameters():
                 idxs_key = n + 'idxs'