Llama3 pre-context parallel dataloader changes (#1400)

Makes a number of updates in preparation for llama3 context-parallel
training. It's still not currently working, need to further update the
model to handle the `cu_seq_lens_q_padded` kwargs and would like to add
a single-GPU CP test that uses BSHD inputs to at least exercise this
code in CI.

This PR:
* Only materializes the dataloader on the cp_rank=0, and returns None on
other ranks.
* Uses the scatter operation in the dataloader to synchronize
`StopIteration` exceptions
* Adds tests for the CP dataloader on 1 and 2-gpu machines
* moves llama3 to use DLCM data as the sanity dataset, turns off some
genome collation options by default. This is larger than the dummy
sequences currently used in training, and will make sure we can fill a
few batches in CP testing. We may want to revert this eventually once
we're done bringing up llama3; since it does trigger the tokenizer
download during testing.
* removes `lazy tokenization` from llama3, this wont work. See
https://nvidia.slack.com/archives/C074Z808N05/p1767818883160949
* starts adding CP files for llama3

Closes BIO-11

---------

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

bionemo-recipes/models/esm2/src/esm/collator.py

@@ -20,7 +20,7 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, TypedDict
 
 import datasets
 import torch
@@ -334,7 +334,7 @@ class ContextParallelDataLoaderWrapper:
 
     def __init__(
         self,
-        dataloader: torch.utils.data.DataLoader,
+        dataloader: torch.utils.data.DataLoader | None,
         cp_mesh: torch.distributed.device_mesh.DeviceMesh,
     ):
         """A dataloader wrapper that distributes the data across the context parallelism group.
@@ -348,15 +348,28 @@ def __init__(
             cp_mesh: The context parallel mesh.
             cp_rank: The rank of the current context parallel process.
         """
-        self.dataloader = dataloader
+        if cp_mesh.get_local_rank() == 0:
+            assert dataloader is not None, "dataloader must be provided on rank 0"
+            self.dataloader = dataloader
+
+        else:
+            assert dataloader is None, "Dataloader on non-rank 0 will not be used"
+
         self.cp_rank = cp_mesh.get_local_rank()
         self.cp_group = cp_mesh.get_group()
         self.num_cp_ranks = cp_mesh.size()
         self._iterator = None
 
+        logger.debug(
+            "Created ContextParallelDataLoaderWrapper on global rank %s, cp rank %s",
+            torch.distributed.get_rank() if torch.distributed.is_initialized() else "<not initialized>",
+            self.cp_rank,
+        )
+
     def __iter__(self):
         """Make the dataloader iterable."""
-        self._iterator = iter(self.dataloader)  # < --- collator output.
+        if self.cp_rank == 0:
+            self._iterator = iter(self.dataloader)  # < --- collator output.
         return self
 
     def __next__(self):
@@ -385,24 +398,19 @@ def _send_data_to_cp_ranks(self):
             batch: The batch for the current CP rank.
 
         """
-        if self.cp_rank == 0:
-            # Get data once, then make copies for each rank.
-            if self._iterator is None:
-                self._iterator = iter(self.dataloader)
-            combined_batch = next(self._iterator)
+        try:
+            combined_batch = next(self._iterator) if self.cp_rank == 0 else None
+        except StopIteration as ex:
+            # If we encounter a StopIteration in the dataloader, we want to raise this error on all the CP ranks, so
+            # that the dataloader can be restarted.
+            combined_batch = [ex] * self.num_cp_ranks
 
-        else:
-            combined_batch = None
-
-        scatter_object_output_list = [None]
-        # Note: This does not provide an async_op handle. Thus its blocking.
-        torch.distributed.scatter_object_list(
-            scatter_object_output_list=scatter_object_output_list,
-            scatter_object_input_list=combined_batch,
-            group=self.cp_group,
-            group_src=0,
-        )
-        return scatter_object_output_list[0]
+        batch_on_this_rank = _scatter_batch_to_cp_ranks(combined_batch, self.cp_group)
+
+        if isinstance(batch_on_this_rank, StopIteration):
+            raise batch_on_this_rank
+
+        return batch_on_this_rank
 
 
 def _split_sample_by_num_tokens(sample: dict[str, Any], num_tokens: int) -> tuple[dict[str, Any], dict[str, Any]]:
@@ -670,3 +678,31 @@ def _get_group_local_rank(group: torch.distributed.ProcessGroup | None = None) -
         return torch.distributed.get_rank()
     global_rank = torch.distributed.get_rank()
     return torch.distributed.get_group_rank(group, global_rank)
+
+
+class BatchType(TypedDict):
+    """The fields in the batch dictionary for context parallel."""
+
+    input_ids: torch.Tensor
+    labels: torch.Tensor
+    cu_seq_lens_q: torch.Tensor
+    cu_seq_lens_k: torch.Tensor
+    cu_seq_lens_q_padded: torch.Tensor
+    cu_seq_lens_k_padded: torch.Tensor
+    max_length_q: int
+    max_length_k: int
+
+
+def _scatter_batch_to_cp_ranks(
+    batch: list[BatchType] | list[StopIteration], cp_group: torch.distributed.ProcessGroup | None = None
+) -> BatchType | StopIteration:
+    """Scatter a batch to all the CP ranks."""
+    scatter_object_output_list = [None]
+    # Note: This does not provide an async_op handle. Thus its blocking.
+    torch.distributed.scatter_object_list(
+        scatter_object_output_list=scatter_object_output_list,
+        scatter_object_input_list=batch,
+        group=cp_group,
+        group_src=0,
+    )
+    return scatter_object_output_list[0]

bionemo-recipes/models/esm2/tests/test_collator_context_parallel.py

@@ -15,7 +15,6 @@
 
 import copy
 import unittest
-from itertools import pairwise
 from typing import Dict, Iterator, List
 from unittest import mock
 
@@ -229,57 +228,6 @@ def size(self) -> int:
         return self._size
 
 
-def _fake_get_batch(
-    cu_seqlens_padded,
-    input_ids_padded,
-    labels_padded,
-    cp_size,
-    qvk_format,
-    cp_rank,
-):
-    total_slices = 2 * cp_size
-    seq_tokens = input_ids_padded.view(-1)
-    seq_labels = labels_padded.view(-1)
-    shard_tokens: List[torch.Tensor] = []
-    shard_labels: List[torch.Tensor] = []
-
-    for start, end in pairwise(cu_seqlens_padded):
-        start_idx = int(start)
-        end_idx = int(end)
-        slice_size = (end_idx - start_idx) // total_slices
-
-        first_start = start_idx + (cp_rank * slice_size)
-        first_end = first_start + slice_size
-        second_start = start_idx + ((total_slices - cp_rank - 1) * slice_size)
-        second_end = second_start + slice_size
-
-        shard_tokens.append(torch.cat([seq_tokens[first_start:first_end], seq_tokens[second_start:second_end]]))
-        shard_labels.append(torch.cat([seq_labels[first_start:first_end], seq_labels[second_start:second_end]]))
-
-    return (
-        torch.cat(shard_tokens).unsqueeze(0),
-        torch.cat(shard_labels).unsqueeze(0),
-    )
-
-
-def _make_cp_shards(base_batch: Dict[str, torch.Tensor], cp_size: int):
-    combined_batch = []
-    for cp_rank in range(cp_size):
-        input_ids_sharded, labels_sharded = _fake_get_batch(
-            cu_seqlens_padded=base_batch["cu_seq_lens_q_padded"],
-            input_ids_padded=base_batch["input_ids"],
-            labels_padded=base_batch["labels"],
-            cp_size=cp_size,
-            qvk_format="thd",
-            cp_rank=cp_rank,
-        )
-        batch_shard = dict(base_batch)
-        batch_shard["input_ids"] = input_ids_sharded
-        batch_shard["labels"] = labels_sharded
-        combined_batch.append(batch_shard)
-    return combined_batch
-
-
 def test_pad_thd_sequences_for_cp():
     pid = 1  # The pad token id.
     label_pad = -100  # The label pad id.
@@ -410,7 +358,7 @@ def run_roundtrip(base_batch):
         cp_mesh_rank0 = _DummyDeviceMesh(size=cp_size, rank=0)
         cp_mesh_rank1 = _DummyDeviceMesh(size=cp_size, rank=1)
         loader_rank0 = ContextParallelDataLoaderWrapper(_DummyLoader(combined_batch), cp_mesh_rank0)
-        loader_rank1 = ContextParallelDataLoaderWrapper(_DummyLoader(combined_batch), cp_mesh_rank1)
+        loader_rank1 = ContextParallelDataLoaderWrapper(None, cp_mesh_rank1)
 
         scatter_payload: Dict[str, List[Dict[str, torch.Tensor]]] = {}
         current_rank = {"value": None}
@@ -499,7 +447,7 @@ def run_roundtrip(base_batch):
         cp_mesh_rank0 = _DummyDeviceMesh(size=cp_size, rank=0)
         cp_mesh_rank1 = _DummyDeviceMesh(size=cp_size, rank=1)
         loader_rank0 = ContextParallelDataLoaderWrapper(_DummyLoader(combined_batch), cp_mesh_rank0)
-        loader_rank1 = ContextParallelDataLoaderWrapper(_DummyLoader(combined_batch), cp_mesh_rank1)
+        loader_rank1 = ContextParallelDataLoaderWrapper(None, cp_mesh_rank1)
 
         scatter_payload: Dict[str, List[Dict[str, torch.Tensor]]] = {}
         current_rank = {"value": None}

bionemo-recipes/recipes/esm2_native_te/collator.py

@@ -20,7 +20,7 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, TypedDict
 
 import datasets
 import torch
@@ -334,7 +334,7 @@ class ContextParallelDataLoaderWrapper:
 
     def __init__(
         self,
-        dataloader: torch.utils.data.DataLoader,
+        dataloader: torch.utils.data.DataLoader | None,
         cp_mesh: torch.distributed.device_mesh.DeviceMesh,
     ):
         """A dataloader wrapper that distributes the data across the context parallelism group.
@@ -348,15 +348,28 @@ def __init__(
             cp_mesh: The context parallel mesh.
             cp_rank: The rank of the current context parallel process.
         """
-        self.dataloader = dataloader
+        if cp_mesh.get_local_rank() == 0:
+            assert dataloader is not None, "dataloader must be provided on rank 0"
+            self.dataloader = dataloader
+
+        else:
+            assert dataloader is None, "Dataloader on non-rank 0 will not be used"
+
         self.cp_rank = cp_mesh.get_local_rank()
         self.cp_group = cp_mesh.get_group()
         self.num_cp_ranks = cp_mesh.size()
         self._iterator = None
 
+        logger.debug(
+            "Created ContextParallelDataLoaderWrapper on global rank %s, cp rank %s",
+            torch.distributed.get_rank() if torch.distributed.is_initialized() else "<not initialized>",
+            self.cp_rank,
+        )
+
     def __iter__(self):
         """Make the dataloader iterable."""
-        self._iterator = iter(self.dataloader)  # < --- collator output.
+        if self.cp_rank == 0:
+            self._iterator = iter(self.dataloader)  # < --- collator output.
         return self
 
     def __next__(self):
@@ -385,24 +398,19 @@ def _send_data_to_cp_ranks(self):
             batch: The batch for the current CP rank.
 
         """
-        if self.cp_rank == 0:
-            # Get data once, then make copies for each rank.
-            if self._iterator is None:
-                self._iterator = iter(self.dataloader)
-            combined_batch = next(self._iterator)
+        try:
+            combined_batch = next(self._iterator) if self.cp_rank == 0 else None
+        except StopIteration as ex:
+            # If we encounter a StopIteration in the dataloader, we want to raise this error on all the CP ranks, so
+            # that the dataloader can be restarted.
+            combined_batch = [ex] * self.num_cp_ranks
 
-        else:
-            combined_batch = None
-
-        scatter_object_output_list = [None]
-        # Note: This does not provide an async_op handle. Thus its blocking.
-        torch.distributed.scatter_object_list(
-            scatter_object_output_list=scatter_object_output_list,
-            scatter_object_input_list=combined_batch,
-            group=self.cp_group,
-            group_src=0,
-        )
-        return scatter_object_output_list[0]
+        batch_on_this_rank = _scatter_batch_to_cp_ranks(combined_batch, self.cp_group)
+
+        if isinstance(batch_on_this_rank, StopIteration):
+            raise batch_on_this_rank
+
+        return batch_on_this_rank
 
 
 def _split_sample_by_num_tokens(sample: dict[str, Any], num_tokens: int) -> tuple[dict[str, Any], dict[str, Any]]:
@@ -670,3 +678,31 @@ def _get_group_local_rank(group: torch.distributed.ProcessGroup | None = None) -
         return torch.distributed.get_rank()
     global_rank = torch.distributed.get_rank()
     return torch.distributed.get_group_rank(group, global_rank)
+
+
+class BatchType(TypedDict):
+    """The fields in the batch dictionary for context parallel."""
+
+    input_ids: torch.Tensor
+    labels: torch.Tensor
+    cu_seq_lens_q: torch.Tensor
+    cu_seq_lens_k: torch.Tensor
+    cu_seq_lens_q_padded: torch.Tensor
+    cu_seq_lens_k_padded: torch.Tensor
+    max_length_q: int
+    max_length_k: int
+
+
+def _scatter_batch_to_cp_ranks(
+    batch: list[BatchType] | list[StopIteration], cp_group: torch.distributed.ProcessGroup | None = None
+) -> BatchType | StopIteration:
+    """Scatter a batch to all the CP ranks."""
+    scatter_object_output_list = [None]
+    # Note: This does not provide an async_op handle. Thus its blocking.
+    torch.distributed.scatter_object_list(
+        scatter_object_output_list=scatter_object_output_list,
+        scatter_object_input_list=batch,
+        group=cp_group,
+        group_src=0,
+    )
+    return scatter_object_output_list[0]

bionemo-recipes/recipes/esm2_native_te/dataset.py

@@ -252,11 +252,16 @@ def create_cp_dataloader(
         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)
+    if cp_mesh.get_local_rank() == 0:
+        train_dataloader, tokenized_dataset = create_thd_dataloader(*args, **kwargs)
 
-    train_dataloader.collate_fn = DataCollatorForContextParallel(
-        collator=train_dataloader.collate_fn,
-        cp_world_size=cp_mesh.size(),
-    )
+        train_dataloader.collate_fn = DataCollatorForContextParallel(
+            collator=train_dataloader.collate_fn,
+            cp_world_size=cp_mesh.size(),
+        )
+
+    else:
+        train_dataloader = None
+        tokenized_dataset = None
 
     return ContextParallelDataLoaderWrapper(train_dataloader, cp_mesh), tokenized_dataset

bionemo-recipes/recipes/esm2_native_te/train_ddp_cp.py

@@ -197,7 +197,8 @@ def main(args: DictConfig) -> float | None:
 
         # Dataloader exhausted, incrementing epoch
         epoch += 1
-        dataset_or_sampler.set_epoch(epoch)
+        if dataset_or_sampler is not None:  # The dataset only exists on rank 0
+            dataset_or_sampler.set_epoch(epoch)
 
     # Save final model to a .safetensors file.
     if args.checkpoint.save_final_model and ckpt_path:

bionemo-recipes/recipes/esm2_native_te/train_fsdp2_cp.py

@@ -214,7 +214,8 @@ def main(args: DictConfig) -> float | None:
 
         # Dataloader exhausted, incrementing epoch
         epoch += 1
-        dataset_or_sampler.set_epoch(epoch)
+        if dataset_or_sampler is not None:  # The dataset only exists on rank 0
+            dataset_or_sampler.set_epoch(epoch)
 
     # Save final model to a .safetensors file.
     if args.checkpoint.save_final_model and ckpt_path:

bionemo-recipes/recipes/llama3_native_te/checkpoint.py

@@ -328,13 +328,15 @@ def save_checkpoint_fsdp2(
         )
         logger.info(f"Saved FSDP2 dataloader to {ckpt_path}")
 
-    # If we're using asynchronous checkpointing, make sure we only have one checkpoint future at a time.
-    if async_save and "fsdp2" in _ckpt_futures and _ckpt_futures["fsdp2"] is not None:
-        _ckpt_futures["fsdp2"].result()
-
     state_dict = {"app": AppState(model=model, optimizer=optimizer, scheduler=scheduler, step=step, epoch=epoch)}
-    ckpt_save_func = dcp_async_save if async_save else dcp_save
-    _ckpt_futures["fsdp2"] = ckpt_save_func(state_dict, checkpoint_id=checkpoint_path, process_group=process_group)
+    if async_save:
+        # If we're using asynchronous checkpointing, make sure we only have one checkpoint future at a time.
+        if "fsdp2" in _ckpt_futures and _ckpt_futures["fsdp2"] is not None:
+            _ckpt_futures["fsdp2"].result()
+
+        _ckpt_futures["fsdp2"] = dcp_async_save(state_dict, checkpoint_id=checkpoint_path, process_group=process_group)
+    else:
+        dcp_save(state_dict, checkpoint_id=checkpoint_path, process_group=process_group)
 
     if max_checkpoints is not None and dist_config.is_main_process():
         prune_checkpoints(ckpt_path, max_checkpoints)