v2.1.18 (#662)

Author: joellidin

docs/miner.md

@@ -173,9 +173,9 @@ This guide will help you set up and run a miner for **τemplar**. We'll cover bo
    # Dataset R2 credentials - You may set up your own Shared Sharded Dataset, but must at minimum set these keys
    # See docs/shared_sharded_dataset.md for instructions
    export R2_DATASET_ACCOUNT_ID="8af7f92a8a0661cf7f1ac0420c932980"
-   export R2_DATASET_BUCKET_NAME="gemma-migration"
-   export R2_DATASET_READ_ACCESS_KEY_ID="a733fac6c32a549e0d48f9f7cf67d758"
-   export R2_DATASET_READ_SECRET_ACCESS_KEY="f50cab456587f015ad21c48c3e23c7ff0e6f1ad5a22c814c3a50d1a4b7c76bb9"
+   export R2_DATASET_BUCKET_NAME="mixed-dataset-migration"
+   export R2_DATASET_READ_ACCESS_KEY_ID="e70cd26850f697479bbb5fd9413713f4"
+   export R2_DATASET_READ_SECRET_ACCESS_KEY="11e3364d6ef70e44d671863fb6de32d474aa6220fa2c9c3df45c5e012ebfbda3"
    export DATASET_BINS_PATH="tokenized/"
 
 

docs/shared_sharded_dataset.md

@@ -19,7 +19,7 @@ The Shared Sharded dataset is based on the [mlfoundations/dclm-baseline-1.0-parq
 For the fastest training, our optimized version includes:
 
 - Pretokenized numpy arrays in .npy files
-- Array slicing provided via .bin files
+- Sample ID arrays provided via .npy files
 
 ## System Requirements
 
@@ -43,9 +43,9 @@ Append the following env keys:
 
 ```bash
 R2_DATASET_ACCOUNT_ID=8af7f92a8a0661cf7f1ac0420c932980
-R2_DATASET_BUCKET_NAME=gemma-migration
-R2_DATASET_READ_ACCESS_KEY_ID=a733fac6c32a549e0d48f9f7cf67d758
-R2_DATASET_READ_SECRET_ACCESS_KEY=f50cab456587f015ad21c48c3e23c7ff0e6f1ad5a22c814c3a50d1a4b7c76bb9
+R2_DATASET_BUCKET_NAME=mixed-dataset-migration
+R2_DATASET_READ_ACCESS_KEY_ID=e70cd26850f697479bbb5fd9413713f4
+R2_DATASET_READ_SECRET_ACCESS_KEY=11e3364d6ef70e44d671863fb6de32d474aa6220fa2c9c3df45c5e012ebfbda3
 DATASET_BINS_PATH="tokenized/"
 ```
 
@@ -92,16 +92,16 @@ Use the CloudFlare migration tool for the easiest setup. Here are the key-value
 
 - Bucket Information
   `Source bucket provider`: `S3-Compatible Storage`
-  `Bucket name`: `gemma-migration`
-  `S3-compatible endpoint URL`: `https://8af7f92a8a0661cf7f1ac0420c932980.r2.cloudflarestorage.com/gemma-migration`
+  `Bucket name`: `mixed-dataset-migration`
+  `S3-compatible endpoint URL`: `https://8af7f92a8a0661cf7f1ac0420c932980.r2.cloudflarestorage.com/mixed-dataset-migration`
 - Required Credentials
-  `Access Key ID`: `a733fac6c32a549e0d48f9f7cf67d758`
-  `Secret Access Key`: `f50cab456587f015ad21c48c3e23c7ff0e6f1ad5a22c814c3a50d1a4b7c76bb9`
+  `Access Key ID`: `e70cd26850f697479bbb5fd9413713f4`
+  `Secret Access Key`: `11e3364d6ef70e44d671863fb6de32d474aa6220fa2c9c3df45c5e012ebfbda3`
 
 #### Page 2
 
 - Select destination R2 bucket
-  `Bucket name`: `gemma-migration`
+  `Bucket name`: `mixed-dataset-migration`
   `Access Key ID`: your_write_id
   `Access Key`: your_secret_write_id
   `Overwrite files?`: `Yes, overwrite (recommended)`
@@ -122,8 +122,8 @@ curl https://rclone.org/install.sh | sudo bash
 # Configure source (read-only)
 rclone config create r2-source s3 \
   provider=Cloudflare \
-  access_key_id=a733fac6c32a549e0d48f9f7cf67d758 \
-  secret_access_key=f50cab456587f015ad21c48c3e23c7ff0e6f1ad5a22c814c3a50d1a4b7c76bb9 \
+  access_key_id=e70cd26850f697479bbb5fd9413713f4 \
+  secret_access_key=11e3364d6ef70e44d671863fb6de32d474aa6220fa2c9c3df45c5e012ebfbda3 \
   endpoint=https://8af7f92a8a0661cf7f1ac0420c932980.r2.cloudflarestorage.com \
   acl=private
 
@@ -139,7 +139,7 @@ rclone config create r2-dest s3 \
 ##### Copy all shards (Full Migration)
 ```bash
 # Copy entire tokenized directory (all shards and sample IDs)
-rclone copy r2-source:gemma-migration/tokenized/ r2-dest:<your-bucket-name>/tokenized/ \
+rclone copy r2-source:mixed-dataset-migration/tokenized/ r2-dest:<your-bucket-name>/tokenized/ \
   --transfers 32 \
   --checkers 16 \
   --progress
@@ -149,10 +149,10 @@ rclone copy r2-source:gemma-migration/tokenized/ r2-dest:<your-bucket-name>/toke
 If you want to test with just the first two shards:
 ```bash
 # Copy first two training shards and their sample IDs
-rclone copy r2-source:gemma-migration/tokenized/train_000000.npy r2-dest:<your-bucket-name>/tokenized/ --progress
-rclone copy r2-source:gemma-migration/tokenized/train_000001.npy r2-dest:<your-bucket-name>/tokenized/ --progress
-rclone copy r2-source:gemma-migration/tokenized/sample_ids_000000.bin r2-dest:<your-bucket-name>/tokenized/ --progress
-rclone copy r2-source:gemma-migration/tokenized/sample_ids_000001.bin r2-dest:<your-bucket-name>/tokenized/ --progress
+rclone copy r2-source:mixed-dataset-migration/tokenized/train_000000.npy r2-dest:<your-bucket-name>/tokenized/ --progress
+rclone copy r2-source:mixed-dataset-migration/tokenized/train_000001.npy r2-dest:<your-bucket-name>/tokenized/ --progress
+rclone copy r2-source:mixed-dataset-migration/tokenized/sample_ids_000000.npy r2-dest:<your-bucket-name>/tokenized/ --progress
+rclone copy r2-source:mixed-dataset-migration/tokenized/sample_ids_000001.npy r2-dest:<your-bucket-name>/tokenized/ --progress
 ```
 
 After migration, update your environment variables to point to your bucket:

hparams/hparams.json

@@ -11,6 +11,7 @@
     "blocks_per_window": 115,
     "windows_per_weights": 3,
     "outer_steps_per_shard": 455,
+    "shard_reset_outer_step": 4040,
     "momentum_decay": 0.95,
     "topk_compression": 64,
     "target_chunk": 64,
@@ -43,8 +44,8 @@
     "eval_lr_factor": 0.2,
     "openskill_beta": 7,
     "openskill_tau": 0.1,
-    "checkpoint_init_version": "2.1.15",
-    "checkpoint_init_window": 59637,
+    "checkpoint_init_version": "2.1.17",
+    "checkpoint_init_window": 60711,
     "num_evaluation_bins": 5,
     "quantization_bins": 4,
     "quantization_range": 6,
@@ -62,6 +63,8 @@
       "scheduler": {
         "warmup_steps": 1500,
         "warmup_inner_steps": 30,
+        "initial_warmup_inner_steps": 200,
+        "replay_rewind_inner_steps": 20000,
         "t_max": 140000,
         "eta_min_factor": 0.1,
         "flatten_start_step": 2740,
@@ -78,6 +81,8 @@
       "scheduler": {
         "warmup_steps": 1500,
         "warmup_inner_steps": 30,
+        "initial_warmup_inner_steps": 200,
+        "replay_rewind_inner_steps": 20000,
         "t_max": 140000,
         "eta_min_factor": 0.1,
         "flatten_start_step": null,

neurons/miner.py

@@ -347,6 +347,9 @@ def __init__(self):
             token_dtype=np.uint32,  # Match preprocessing script dtype
         )
         self.outer_steps_per_shard = getattr(self.hparams, "outer_steps_per_shard")
+        self.shard_reset_outer_step = getattr(
+            self.hparams, "shard_reset_outer_step", None
+        )
 
         tplr.logger.info("[Init] ✔ fully done – entering run()")
 
@@ -417,7 +420,11 @@ async def run(self):
 
         self.comms.start_commitment_fetcher()
 
-        current_shard = self.global_step // self.outer_steps_per_shard
+        current_shard_epoch, current_shard = tplr.sharded_dataset.compute_shard_state(
+            self.global_step,
+            self.outer_steps_per_shard,
+            self.shard_reset_outer_step,
+        )
         tplr.logger.info(
             f"Starting with global_step={self.global_step} (actual outer steps)"
         )
@@ -432,6 +439,7 @@ async def run(self):
         self.set_dataloader()
 
         # Track the current shard to avoid double-swapping at initialization
+        last_shard_epoch = current_shard_epoch
         last_shard = current_shard
 
         # Put a dummy gradient to mark this miner as active for validators
@@ -474,8 +482,24 @@ async def run(self):
             self.sampler.set_window_uid(self.uid, step_window)
 
             # Check if we need to swap dataset based on shard index change
-            current_shard_check = self.global_step // self.outer_steps_per_shard
-            if current_shard_check > last_shard:
+            shard_epoch_check, current_shard_check = (
+                tplr.sharded_dataset.compute_shard_state(
+                    self.global_step,
+                    self.outer_steps_per_shard,
+                    self.shard_reset_outer_step,
+                )
+            )
+            if shard_epoch_check != last_shard_epoch:
+                tplr.logger.info(
+                    f"Resetting shard schedule at outer_step {self.global_step} "
+                    f"to shard {current_shard_check}"
+                )
+                await self.dataset_manager.initialize_datasets(current_shard_check)
+                self.set_dataloader()
+                dist_helper.safe_barrier("sync_shard_switch", self.local_rank)
+                last_shard_epoch = shard_epoch_check
+                last_shard = current_shard_check
+            elif current_shard_check > last_shard:
                 tplr.logger.info(
                     f"Swapping dataset after {self.global_step} outer steps at window {step_window}"
                 )

neurons/validator.py

@@ -521,6 +521,9 @@ def __init__(self):
         self.param_change_alpha = 0.2
 
         self.outer_steps_per_shard = getattr(self.hparams, "outer_steps_per_shard")
+        self.shard_reset_outer_step = getattr(
+            self.hparams, "shard_reset_outer_step", None
+        )
         self.dataset_manager = tplr.sharded_dataset.ShardedDatasetManager(
             sequence_length=self.hparams.sequence_length,
             rank=self.local_rank,  # Use local_rank for proper file operations
@@ -1245,7 +1248,11 @@ async def run(self):
             aggregator_device="cpu",
         )
 
-        current_shard = self.global_step // self.outer_steps_per_shard
+        shard_epoch, current_shard = tplr.sharded_dataset.compute_shard_state(
+            self.global_step,
+            self.outer_steps_per_shard,
+            self.shard_reset_outer_step,
+        )
 
         # Initialize datasets (only rank 0 downloads, handled internally by dataset_manager)
         _ = await self.dataset_manager.initialize_datasets(current_shard)
@@ -1256,6 +1263,7 @@ async def run(self):
         self.set_dataloader(validator=True)
 
         # Track the current shard to avoid double-swapping at initialization
+        last_shard_epoch = shard_epoch
         last_shard = current_shard
 
         if self.is_master:
@@ -1287,8 +1295,24 @@ async def run(self):
             window_start = tplr.T()
 
             # Check if we need to swap dataset based on shard index change
-            current_shard_check = self.global_step // self.outer_steps_per_shard
-            if current_shard_check > last_shard:
+            shard_epoch_check, current_shard_check = (
+                tplr.sharded_dataset.compute_shard_state(
+                    self.global_step,
+                    self.outer_steps_per_shard,
+                    self.shard_reset_outer_step,
+                )
+            )
+            if shard_epoch_check != last_shard_epoch:
+                tplr.logger.info(
+                    f"Resetting shard schedule at outer_step {self.global_step} "
+                    f"to shard {current_shard_check}"
+                )
+                await self.dataset_manager.initialize_datasets(current_shard_check)
+                self.set_dataloader(validator=True)
+                dist_helper.safe_barrier("sync_shard_switch", self.local_rank)
+                last_shard_epoch = shard_epoch_check
+                last_shard = current_shard_check
+            elif current_shard_check > last_shard:
                 tplr.logger.info(
                     f"Swapping dataset after {self.global_step} outer steps at window {self.current_window}"
                 )

scripts/dataset_prep/02_consolidate_shards.py

@@ -108,10 +108,37 @@ async def run_preprocessing(
                 )
             )
 
-        tokens_view = np.memmap(tokens_file, dtype=token_dtype, mode="r")
-        tok_u32 = tokens_view.view(np.uint32)  # reinterpret for 4-byte hashing
+        # Load tokens - if .npy file, use np.load to respect embedded dtype
+        if tokens_file.endswith(".npy"):
+            tokens_view = np.load(tokens_file, mmap_mode="r", allow_pickle=False)
+            # Ensure it's uint32 (step 01 saves as uint32)
+            if tokens_view.dtype != np.uint32:
+                tqdm.write(
+                    f"Warning: Shard {i} has dtype {tokens_view.dtype}, expected uint32. "
+                    f"Converting (this may indicate a preprocessing mismatch)."
+                )
+                tok_u32 = tokens_view.astype(np.uint32)
+            else:
+                tok_u32 = tokens_view
+        else:
+            # Raw binary file - use specified dtype then reinterpret as uint32
+            tokens_view = np.memmap(tokens_file, dtype=token_dtype, mode="r")
+            tok_u32 = tokens_view.view(np.uint32)
+
+        # Only create sample IDs for complete sequences (discard partial sequence at end)
+        total_tokens = tok_u32.shape[0]
+        num_complete_samples = total_tokens // seq_len
+
+        # Warn if there's a partial sequence that will be discarded
+        remainder = total_tokens % seq_len
+        if remainder > 0:
+            tqdm.write(
+                f"Warning: Shard {i} has {remainder} tokens remaining after chunking "
+                f"(total: {total_tokens}, seq_len: {seq_len}). "
+                f"Creating {num_complete_samples} complete samples."
+            )
 
-        raw_idx = np.arange(0, tok_u32.shape[0] + 1, seq_len)
+        raw_idx = np.arange(0, num_complete_samples * seq_len + 1, seq_len)
         starts = raw_idx[:-1]
         ends = raw_idx[1:]
 
@@ -208,7 +235,6 @@ async def main() -> None:
     print(f"  • Shards path: {args.r2_prefix}")
     print(f"  • Sequence length: {args.seq_len}")
     print(f"  • Token dtype: {args.token_dtype}")
-    print(f"  • Skip Validation: {args.skip_validation}")
     print()
 
     success = await run_preprocessing(args, args.seq_len, token_dtype)

src/tplr/__init__.py

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

src/tplr/neurons.py

@@ -689,6 +689,38 @@ async def handle_checkpoint_catchup(
         from_bootstrap: Whether checkpoint was from bootstrap version
         aggregator_device: which device to load aggregation results to
     """
+    # Determine scheduler config and warmup settings
+    optimizer_cfg = getattr(instance.hparams, "optimizer", {})
+    opt_type = optimizer_cfg.get("type", "adamw").lower()
+    opt_cfg = optimizer_cfg.get(opt_type, {})
+    scheduler_cfg = opt_cfg.get("scheduler", {})
+
+    default_warmup_inner = scheduler_cfg.get(
+        "warmup_inner_steps", getattr(instance, "warmup_inner_steps", 0)
+    )
+    startup_warmup_inner = scheduler_cfg.get(
+        "initial_warmup_inner_steps", default_warmup_inner
+    )
+
+    # Set warmup length:
+    # - If resuming from bootstrap, use the longer startup warmup.
+    # - If resuming from a regular checkpoint, use the default.
+    # - If global_step is 0, leave as-is; the scheduler's own warmup covers this.
+    if ckpt_global_step == 0:
+        tplr.logger.info("Global step is 0; leaving warmup settings unchanged.")
+    elif from_bootstrap:
+        instance.warmup_inner_steps = startup_warmup_inner
+        tplr.logger.info(
+            f"Applying startup warmup_inner_steps={startup_warmup_inner} (bootstrap resume)"
+        )
+        instance.warmup_steps_taken = 0
+    else:
+        instance.warmup_inner_steps = default_warmup_inner
+        tplr.logger.info(
+            f"Applying resumed warmup_inner_steps={default_warmup_inner} (checkpoint resume)"
+        )
+        instance.warmup_steps_taken = 0
+
     # Decide catch-up windows and run catch-up on ALL ranks
     # When loading from bootstrap, we always need to catch up from start_window
     # to ensure we're using current version's gradients
@@ -726,7 +758,17 @@ async def handle_checkpoint_catchup(
     # Replay scheduler steps based on windows completed from checkpoint
     # ckpt_global_step tracks windows, scheduler needs inner_steps per window
     total_inner_steps = ckpt_global_step * instance.hparams.inner_steps
-    if total_inner_steps > 0:
+
+    # Apply configurable rewind before replaying scheduler to give slack on restarts
+    rewind_inner_steps = scheduler_cfg.get("replay_rewind_inner_steps", 0)
+    if rewind_inner_steps > 0:
+        total_inner_steps = max(total_inner_steps - rewind_inner_steps, 0)
+        tplr.logger.info(
+            f"Rewinding scheduler replay by {rewind_inner_steps} inner steps; "
+            f"{total_inner_steps} steps remain to replay"
+        )
+
+    if total_inner_steps > 0 and getattr(instance, "inner_scheduler", None) is not None:
         for _ in range(total_inner_steps):
             # Respect flatten window during replay
             if not instance.should_skip_scheduler_step():