feat/new anneal shard

docs/shared_sharded_dataset.md

@@ -146,11 +146,11 @@ rclone copy r2-source:mixed-dataset-migration/anneal/ r2-dest:<your-bucket-name>
 ```
 
 ##### Copy specific shards (Partial Migration for Testing)
-If you want to test with just the first shard:
+If you want to test with the current anneal shard:
 ```bash
-# Copy first training shard and its sample IDs
-rclone copy r2-source:mixed-dataset-migration/anneal/anneal_000000.npy r2-dest:<your-bucket-name>/anneal/ --progress
-rclone copy r2-source:mixed-dataset-migration/anneal/sample_ids_anneal_000000.npy r2-dest:<your-bucket-name>/anneal/ --progress
+# Copy anneal shard 2 and its sample IDs
+rclone copy r2-source:mixed-dataset-migration/anneal/anneal_000002.npy r2-dest:<your-bucket-name>/anneal/ --progress
+rclone copy r2-source:mixed-dataset-migration/anneal/sample_ids_anneal_000002.npy r2-dest:<your-bucket-name>/anneal/ --progress
 ```
 
 After migration, update your environment variables to point to your bucket:

hparams/hparams.json

@@ -5,7 +5,7 @@
       "start_global_step": 6100,
       "warmup_inner_steps": 200,
       "decay_outer_steps": 550,
-      "peak_lr_factor": 0.3,
+      "peak_lr_factor": 0.25,
       "eta_min_factor": 0.0,
       "file_prefix": "anneal"
     },

neurons/miner.py

@@ -434,9 +434,9 @@ async def run(self):
             self.outer_steps_per_shard,
             self.shard_reset_outer_step,
         )
-        # In anneal mode, always use shard 0
+        # In anneal mode, always use shard 2
         if self.dataset_manager.anneal_mode:
-            current_shard = 0
+            current_shard = 2
             current_shard_epoch = 0
         tplr.logger.info(
             f"Starting with global_step={self.global_step} (actual outer steps)"

neurons/validator.py

@@ -1263,9 +1263,9 @@ async def run(self):
             self.outer_steps_per_shard,
             self.shard_reset_outer_step,
         )
-        # In anneal mode, always use shard 0
+        # In anneal mode, always use shard 2
         if self.dataset_manager.anneal_mode:
-            current_shard = 0
+            current_shard = 2
             shard_epoch = 0
 
         # Initialize datasets (only rank 0 downloads, handled internally by dataset_manager)

src/tplr/__init__.py

@@ -20,7 +20,7 @@
 # mypy: ignore-errors
 # type: ignore
 
-__version__ = "2.1.23"
+__version__ = "2.1.24"
 
 # Import package.
 from .chain import *

src/tplr/sharded_dataset.py

@@ -396,7 +396,7 @@ async def initialize_datasets(self, current_shard_index: int) -> None:
 
         self.active_dataset = await self.create_dataset(current_shard_index)
 
-        # In anneal mode, don't prepare next shard (we stay on shard 0)
+        # In anneal mode, don't prepare next shard (we stay on one shard)
         if self.anneal_mode:
             self.upcoming_dataset = asyncio.create_task(asyncio.sleep(0))
             return