add back num_workers and worldsize dataloader resumption checks (#1250)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Dataloader checkpoints now capture worker and world-size settings to
ensure compatible restores.
* On restore, mismatches in workers or world size trigger a warning and
safely restart the dataloader from scratch; compatible states restore
with a success message.

* **Tests**
* Added tests validating warning logs and restart behavior when worker
or world-size settings differ from the saved checkpoint, and successful
restore when they match.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Peter St. John <[email protected]>

Author: pstjohn

bionemo-recipes/recipes/esm2_native_te/checkpoint.py

@@ -515,6 +515,8 @@ def save_dataloader(
     dataloader_path = ckpt_path / f"dataloader_rank_{dist_config.rank}.pt"
 
     dataloader_state = dataloader.state_dict()
+    dataloader_state["num_workers"] = dataloader.num_workers
+    dataloader_state["num_ranks"] = dist_config.world_size
     torch.save(dataloader_state, dataloader_path)
     if dist_config.is_main_process():
         logger.info(f"Saved dataloader state to {dataloader_path}")
@@ -545,6 +547,18 @@ def load_dataloader(
         return dataloader
 
     dataloader_state = torch.load(dataloader_path)
+
+    if (
+        dataloader.num_workers != dataloader_state["num_workers"]
+        or dist_config.world_size != dataloader_state["num_ranks"]
+    ):
+        logger.warning(
+            f"Dataloader num_workers mismatch: {dataloader.num_workers} != {dataloader_state['num_workers']} or "
+            f"num_ranks mismatch: {dist_config.world_size} != {dataloader_state['num_ranks']}, "
+            "starting dataloader from scratch."
+        )
+        return dataloader
+
     dataloader.load_state_dict(dataloader_state)
     if dist_config.is_main_process():
         logger.info(f"Loaded dataloader state from {dataloader_path}")

bionemo-recipes/recipes/esm2_native_te/tests/test_dataset.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 import shutil
 from dataclasses import dataclass
@@ -711,3 +712,124 @@ def test_lazy_tokenization_returns_batch():
 
     batch = next(iter(dataloader))
     assert batch is not None
+
+
+def test_stateful_dataloader_load_fails_if_num_workers_mismatch(tmp_path, caplog):
+    dataloader_path = tmp_path / "dl_test_num_workers_mismatch"
+    shutil.rmtree(dataloader_path, ignore_errors=True)
+    os.makedirs(dataloader_path, exist_ok=True)
+    tokenizer_name = "facebook/esm2_t6_8M_UR50D"
+    load_dataset_kwargs = {
+        "path": "parquet",
+        "split": "train",
+        "data_files": "train.parquet",
+        "streaming": False,
+    }
+
+    rank0_dist_config = MockDistributedConfig(
+        rank=0,
+        local_rank=0,
+        world_size=1,
+    )
+
+    reference_dataloader, _ = create_dataloader(
+        distributed_config=rank0_dist_config,
+        tokenizer_name=tokenizer_name,
+        load_dataset_kwargs=load_dataset_kwargs,
+        micro_batch_size=4,
+        num_workers=1,
+        mlm_probability=0,
+    )
+
+    save_dataloader(
+        dataloader=reference_dataloader,
+        ckpt_path=dataloader_path,
+        dist_config=rank0_dist_config,
+    )
+
+    del reference_dataloader
+
+    reference_dataloader, _ = create_dataloader(
+        distributed_config=rank0_dist_config,
+        tokenizer_name=tokenizer_name,
+        load_dataset_kwargs=load_dataset_kwargs,
+        micro_batch_size=4,
+        num_workers=2,
+        mlm_probability=0,
+    )
+
+    with caplog.at_level(logging.WARNING):
+        load_dataloader(
+            dataloader=reference_dataloader,
+            ckpt_path=dataloader_path,
+            dist_config=rank0_dist_config,
+        )
+
+    assert (
+        "Dataloader num_workers mismatch: 2 != 1 or num_ranks mismatch: 1 != 1, starting dataloader from scratch."
+        in caplog.text
+    )
+
+
+def test_stateful_dataloader_load_fails_if_num_ranks_mismatch(tmp_path, caplog):
+    dataloader_path = tmp_path / "dl_test_num_workers_mismatch"
+    shutil.rmtree(dataloader_path, ignore_errors=True)
+    os.makedirs(dataloader_path, exist_ok=True)
+    tokenizer_name = "facebook/esm2_t6_8M_UR50D"
+    load_dataset_kwargs = {
+        "path": "parquet",
+        "split": "train",
+        "data_files": "train.parquet",
+        "streaming": False,
+    }
+
+    rank0_dist_config = MockDistributedConfig(
+        rank=0,
+        local_rank=0,
+        world_size=1,
+    )
+
+    reference_dataloader, _ = create_dataloader(
+        distributed_config=rank0_dist_config,
+        tokenizer_name=tokenizer_name,
+        load_dataset_kwargs=load_dataset_kwargs,
+        micro_batch_size=4,
+        num_workers=1,
+        mlm_probability=0,
+    )
+
+    save_dataloader(
+        dataloader=reference_dataloader,
+        ckpt_path=dataloader_path,
+        dist_config=rank0_dist_config,
+    )
+
+    del reference_dataloader
+    del rank0_dist_config
+
+    rank2_dist_config = MockDistributedConfig(
+        rank=0,
+        local_rank=0,
+        world_size=2,
+    )
+
+    reference_dataloader, _ = create_dataloader(
+        distributed_config=rank2_dist_config,
+        tokenizer_name=tokenizer_name,
+        load_dataset_kwargs=load_dataset_kwargs,
+        micro_batch_size=4,
+        num_workers=1,
+        mlm_probability=0,
+    )
+
+    with caplog.at_level(logging.WARNING):
+        load_dataloader(
+            dataloader=reference_dataloader,
+            ckpt_path=dataloader_path,
+            dist_config=rank2_dist_config,
+        )
+
+    assert (
+        "Dataloader num_workers mismatch: 1 != 1 or num_ranks mismatch: 2 != 1, starting dataloader from scratch."
+        in caplog.text
+    )