add test for cp_dataloader that doesn't require datacenter hardware

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

Author: pstjohn

bionemo-recipes/recipes/esm2_native_te/dataset.py

@@ -244,6 +244,11 @@ def create_cp_dataloader(
     Returns:
         A tuple of (dataloader, dataset_or_sampler).
     """
+    # Ensure pad_sequences_to_be_divisible_by is passed to create_thd_dataloader
+    if kwargs.get("pad_sequences_to_be_divisible_by", None) is None:
+        logger.info("pad_sequences_to_be_divisible_by is not provided, using cp_mesh.size() * 2")
+        kwargs["pad_sequences_to_be_divisible_by"] = cp_mesh.size() * 2
+
     train_dataloader, tokenized_dataset = create_thd_dataloader(*args, **kwargs)
 
     train_dataloader.collate_fn = DataCollatorForContextParallel(

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

@@ -16,12 +16,21 @@
 import logging
 import os
 import shutil
+import subprocess
 from dataclasses import dataclass
 
+import pytest
 import torch
+from torch.distributed.device_mesh import init_device_mesh
 
 from checkpoint import load_dataloader, save_dataloader
-from dataset import create_bshd_dataloader, create_thd_dataloader
+from dataset import DistributedConfig, create_bshd_dataloader, create_cp_dataloader, create_thd_dataloader
+
+
+requires_multi_gpu = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.device_count() < 2,
+    reason="Test requires at least 2 GPUs",
+)
 
 
 @dataclass
@@ -880,3 +889,78 @@ def test_token_packing_dataloader():
     batch = next(iter(dataloader))
     assert batch["input_ids"].shape[1] == 8 * 1024
     assert batch["labels"].shape[1] == 8 * 1024
+
+
+@requires_multi_gpu
+def test_cp_dataloader(recipe_path):
+    import os
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = str(recipe_path)
+
+    cmd = [
+        "torchrun",
+        "--nproc_per_node=2",
+        "tests/test_dataset.py",
+    ]
+
+    result = subprocess.run(
+        cmd,
+        check=False,
+        text=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        timeout=240,
+        cwd=str(recipe_path),
+        env=env,
+    )
+    if result.returncode != 0:
+        print(f"STDOUT:\n{result.stdout}")
+        print(f"STDERR:\n{result.stderr}")
+        pytest.fail(f"Command failed with exit code {result.returncode}")
+
+
+if __name__ == "__main__":
+    dist_config = DistributedConfig()
+    device = torch.device(f"cuda:{dist_config.local_rank}")
+    torch.distributed.init_process_group(backend="nccl", device_id=device)
+    torch.cuda.set_device(dist_config.local_rank)
+    device_mesh = init_device_mesh("cuda", mesh_shape=(1, 2), mesh_dim_names=("dp", "cp"))
+
+    dataloader, _ = create_cp_dataloader(
+        distributed_config=dist_config,
+        cp_mesh=device_mesh["cp"],
+        tokenizer_name="facebook/esm2_t6_8M_UR50D",
+        load_dataset_kwargs={
+            "path": "parquet",
+            "split": "train",
+            "data_files": "train.parquet",
+            "streaming": True,
+        },
+        token_micro_batch_size=8 * 1024,
+        num_workers=1,
+    )
+
+    batch = next(iter(dataloader))
+    # With CP size 2, each sequence is split into 2 * cp_world_size = 4 slices.
+    # Each rank gets 2 slices (beginning and end), so each rank gets approximately
+    # (8 * 1024) / 2 = 4096 tokens per rank
+    # Note: Sequences are padded to be divisible by pad_sequences_to_be_divisible_by
+    # (which defaults to cp_mesh.size() * 2 = 4 if not provided)
+    # The actual token count per rank can vary due to:
+    # 1. Sequence packing (variable-length sequences packed up to token_micro_batch_size)
+    # 2. Per-sequence padding to be divisible by pad_sequences_to_be_divisible_by
+    # 3. CP splitting logic that takes slices from beginning and end
+    expected_tokens_per_rank = (8 * 1024) // device_mesh["cp"].size()
+    actual_shape = batch["input_ids"].shape[1]
+    # Allow for variance due to sequence packing, padding, and CP splitting
+    # The actual shape should be close to expected_tokens_per_rank but can vary
+    # Allow up to 100 tokens of variance (both above and below) to account for
+    # sequence packing and padding effects
+    assert actual_shape >= expected_tokens_per_rank - 100, (
+        f"Expected at least {expected_tokens_per_rank - 100} tokens, got {actual_shape}"
+    )
+    assert actual_shape <= expected_tokens_per_rank + 100, (
+        f"Expected at most {expected_tokens_per_rank + 100} tokens, got {actual_shape}"
+    )
+    assert batch["labels"].shape[1] == actual_shape