drop dispatching dataloader (#245)

mayank31398 · web-flow · commit 6cc864b53903 · 2025-06-29T19:24:46.000-04:00
Signed-off-by: Mayank Mishra &lt;mayank31398@gmail.com&gt;
diff --git a/configs/finetuning-example.yml b/configs/finetuning-example.yml
@@ -74,5 +74,3 @@ distributed_args:
   # use ZeRO-3 for model sharding, saves most memory but needs more communication. this is fine since we are doing training on 2 GPUs and they are connected via NVLink
   stage: 3
   torch_compile: true
-  # this will load dataset only on the first GPU and send part of the data to the other GPUs, not recommended unless the datasets are immensely large
-  dispatching_dataloader: false
diff --git a/lm_engine/arguments.py b/lm_engine/arguments.py
@@ -263,8 +263,6 @@ class DistributedArgs(BaseArgs):
     communication_dtype: str | None = None
     # whether to use torch.compile
     torch_compile: bool = False
-    # whether to use a dispatching dataloader
-    dispatching_dataloader: bool = False
     # tensor parallel world size
     tensor_parallel_world_size: int = 1
     # whether to use sequence parallel
diff --git a/lm_engine/data/__init__.py b/lm_engine/data/__init__.py
@@ -5,15 +5,12 @@
 import logging
 from functools import partial
 
-import torch
-import torch.distributed
-
 from ..arguments import DatasetArgs, InferenceArgs, TrainingArgs
 from ..enums import DatasetSplit, Mode
 from ..tokenizers import TOKENIZER_TYPE
 from ..utils import ProcessGroupManager, log_rank_0, run_rank_n
 from .base import BaseDataset, BlendedDatasets
-from .dataloader import DispatchingDataLoader, ResumableDataLoader
+from .dataloader import ResumableDataLoader
 from .debug import DebugDataset
 from .huggingface import HuggingFaceDataset
 from .ibm import get_ibm_dataloaders
@@ -107,111 +104,6 @@ def get_finetuning_dataloader(
     if ProcessGroupManager.get_tensor_parallel_rank() != 0:
         return
 
-    if args.distributed_args.dispatching_dataloader:
-        assert (
-            ProcessGroupManager.get_tensor_parallel_world_size() == 1
-        ), "tensor parallel doesn't support dispatching dataloader"
-
-        dataloader = _get_dispatching_dataloader(args, split=split, mode=mode, tokenizer=tokenizer)
-    else:
-        dataloader = _get_non_dispatching_dataloader(args, split=split, mode=mode, tokenizer=tokenizer)
-
-    return dataloader
-
-
-def get_pretraining_dataloaders(
-    args: TrainingArgs, tokenizer: TOKENIZER_TYPE, consumed_samples: int
-) -> tuple[ResumableDataLoader, list[ResumableDataLoader], list[ResumableDataLoader]]:
-    if args.datasets[0].class_name == "MegatronDataset":
-        dataloaders = get_megatron_gpt_dataloaders(args, tokenizer, consumed_samples=consumed_samples)
-    elif args.datasets[0].class_name == "IBMDataset":
-        dataloaders = get_ibm_dataloaders(args, tokenizer)
-
-    return dataloaders
-
-
-def _get_dispatching_dataloader(
-    args: TrainingArgs | InferenceArgs, split: DatasetSplit, mode: Mode, tokenizer: TOKENIZER_TYPE
-) -> ResumableDataLoader:
-    micro_batch_size = args.training_parameters.micro_batch_size
-
-    num_ranks_per_node = torch.cuda.device_count()
-    node_rank = ProcessGroupManager.get_global_rank() // num_ranks_per_node
-    num_nodes = ProcessGroupManager.get_world_size() // num_ranks_per_node
-
-    def _get_source_broadcast_mapping() -> dict:
-        result = {}
-        for i in range(num_nodes):
-            source = i * num_ranks_per_node
-            ranks = list(range(source, source + num_ranks_per_node))
-            result[source] = torch.distributed.new_group(ranks)
-        return result
-
-    source_broadcast_mapping = _get_source_broadcast_mapping()
-
-    # check if node's first rank
-    if ProcessGroupManager.get_global_rank() == node_rank * num_ranks_per_node:
-        datasets_list, data_sampling_ratios = get_datasets_list(
-            dataset_args_list=args.datasets, split=split, mode=Mode.training, tokenizer=tokenizer
-        )
-
-        if len(datasets_list) == 0:
-            return None
-
-        blended_dataset = BlendedDatasets(datasets=datasets_list, split=split)
-        data_sampling_ratios = [1] if len(datasets_list) == 1 else data_sampling_ratios
-
-        # each node is given a data sampler
-        # TODO modify this when we add model parallelism
-
-        # sampler routes to the dispatching parent worker
-        sampler = BlendedDistributedSampler(
-            dataset=blended_dataset,
-            data_sampling_ratios=data_sampling_ratios,
-            num_replicas=num_nodes,
-            rank=node_rank,
-            ignore_sampling_proportion_for_validation=args.training_parameters.ignore_sampling_proportion_for_validation,
-            shuffle=split == DatasetSplit.train,
-            seed=args.random_args.seed,
-            drop_last=False,
-        )
-    else:
-        blended_dataset = None
-        data_sampling_ratios = None
-        sampler = None
-
-    # dataloader does local dispatching and thus needs source_rank and broadcast_ranks
-    dataloader = DispatchingDataLoader(
-        blended_dataset,
-        batch_size=micro_batch_size,
-        sampler=sampler,
-        collate_fn=partial(
-            collate_fn,
-            mode=mode,
-            loss_mask=args.training_parameters.loss_mask,
-            eos_token_id=tokenizer.eos_token_id,
-            use_padding_free_transformer=args.model_args.use_padding_free_transformer,
-            pad_to_multiple_of=ProcessGroupManager.get_tensor_parallel_world_size(),
-        ),
-        source_broadcast_mapping=source_broadcast_mapping,
-        broadcast_world_size=num_ranks_per_node,
-    )
-
-    _log_dataset(
-        blended_dataset=blended_dataset,
-        sampler=sampler,
-        split=split,
-        num_training_steps=args.training_parameters.num_training_steps,
-        gradient_accumulation_steps=args.training_parameters.gradient_accumulation_steps,
-        micro_batch_size=args.training_parameters.micro_batch_size,
-    )
-
-    return dataloader
-
-
-def _get_non_dispatching_dataloader(
-    args: TrainingArgs | InferenceArgs, split: DatasetSplit, mode: Mode, tokenizer: TOKENIZER_TYPE
-) -> ResumableDataLoader:
     micro_batch_size = args.training_parameters.micro_batch_size
 
     datasets_list, data_sampling_ratios = get_datasets_list(
@@ -262,6 +154,17 @@ def _get_non_dispatching_dataloader(
     return dataloader
 
 
+def get_pretraining_dataloaders(
+    args: TrainingArgs, tokenizer: TOKENIZER_TYPE, consumed_samples: int
+) -> tuple[ResumableDataLoader, list[ResumableDataLoader], list[ResumableDataLoader]]:
+    if args.datasets[0].class_name == "MegatronDataset":
+        dataloaders = get_megatron_gpt_dataloaders(args, tokenizer, consumed_samples=consumed_samples)
+    elif args.datasets[0].class_name == "IBMDataset":
+        dataloaders = get_ibm_dataloaders(args, tokenizer)
+
+    return dataloaders
+
+
 @run_rank_n
 def _log_dataset(
     blended_dataset: BlendedDatasets,
diff --git a/lm_engine/data/dataloader.py b/lm_engine/data/dataloader.py
@@ -4,15 +4,7 @@
 
 from __future__ import annotations
 
-from typing import Callable, Iterable, Iterator
-
-import torch
-import torch.distributed
-from torch.distributed import ProcessGroup
-from torch.utils.data import DataLoader, Dataset, Sampler
-
-from ..communication import Communication
-from ..utils import ProcessGroupManager
+from torch.utils.data import DataLoader
 
 
 class ResumableDataLoader(DataLoader):
@@ -22,106 +14,3 @@ def state_dict(self) -> dict:
     def load_state_dict(self, state_dict: dict) -> None:
         self.dataset.load_state_dict(state_dict.get("dataset"))
         self.sampler.load_state_dict(state_dict.get("sampler"))
-
-
-class DispatchingDataLoader(ResumableDataLoader):
-    def __init__(
-        self,
-        dataset: Dataset,
-        batch_size: int | None = 1,
-        sampler: Sampler | Iterable | None = None,
-        batch_sampler: Sampler[list] | Iterable[list] | None = None,
-        num_workers: int = 0,
-        collate_fn: Callable | None = None,
-        pin_memory: bool = False,
-        drop_last: bool = False,
-        source_broadcast_mapping: dict[int, ProcessGroup] | None = None,
-        broadcast_world_size: int | None = None,
-        static_shape_per_rank: tuple[int, int] | None = None,
-        keys: list[str] = ["input_ids", "attention_mask", "labels"],
-    ) -> DispatchingDataLoader:
-        self.broadcast_world_size = broadcast_world_size
-
-        self.is_source, self.source_rank, self.local_rank_in_broadcast_group, self.broadcast_group = (
-            get_source_and_broadcast_group(source_broadcast_mapping)
-        )
-
-        super().__init__(
-            dataset=dataset,
-            batch_size=batch_size * self.broadcast_world_size if batch_sampler is None else 1,
-            sampler=sampler,
-            batch_sampler=batch_sampler,
-            num_workers=num_workers,
-            collate_fn=collate_fn,
-            pin_memory=pin_memory,
-            drop_last=drop_last,
-        )
-
-        _length = torch.tensor(
-            [super().__len__() if self.is_source else 0], dtype=torch.long, device=torch.cuda.current_device()
-        )
-        torch.distributed.broadcast(_length, src=self.source_rank, group=self.broadcast_group)
-        self._length = _length.item()
-
-        self.global_static_shape = None
-        if static_shape_per_rank is not None:
-            self.global_static_shape = (static_shape_per_rank[0] * self.broadcast_world_size, static_shape_per_rank[1])
-
-        self.keys = keys
-
-    def __iter__(self) -> Iterator[dict]:
-        iterator = super().__iter__() if self.is_source else range(self._length)
-
-        for batch in iterator:
-            # if using dynamic shapes at every batch or when batch buffer is None during static batch, we need to get shape
-            # send/recv tensor shapes
-            if self.global_static_shape is None:
-                batch_shape = batch[self.keys[0]].shape if self.is_source else None
-                batch_shape = Communication.broadcast_object(
-                    batch_shape, src=self.source_rank, group=self.broadcast_group
-                )
-            else:
-                batch_shape = self.global_static_shape
-
-            if self.is_source:
-                for key in self.keys:
-                    batch[key] = batch[key].to(torch.cuda.current_device())
-            else:
-                batch = {
-                    key: torch.empty(batch_shape, dtype=torch.long, device=torch.cuda.current_device())
-                    for key in self.keys
-                }
-
-            for key in self.keys:
-                # send/recv batch
-                torch.distributed.broadcast(batch[key], src=self.source_rank, group=self.broadcast_group)
-
-                # slice batch
-                local_batch_size = batch[key].shape[0] // self.broadcast_world_size
-                batch[key] = batch[key][
-                    self.local_rank_in_broadcast_group
-                    * local_batch_size : (self.local_rank_in_broadcast_group + 1)
-                    * local_batch_size
-                ]
-
-            yield batch
-
-    def __len__(self) -> int:
-        return self._length
-
-
-def get_source_and_broadcast_group(
-    source_broadcast_mapping: dict[int, ProcessGroup],
-) -> tuple[bool, int, int, ProcessGroup]:
-    global_rank = ProcessGroupManager.get_global_rank()
-
-    for source_rank, broadcast_group in source_broadcast_mapping.items():
-        ranks = torch.distributed.get_process_group_ranks(broadcast_group)
-
-        if global_rank in ranks:
-            is_source = global_rank == source_rank
-            local_rank_in_broadcast_group = ranks.index(global_rank)
-
-            return is_source, source_rank, local_rank_in_broadcast_group, broadcast_group
-
-    assert False, "code shouldn't reach here"
diff --git a/lm_engine/data/megatron/__init__.py b/lm_engine/data/megatron/__init__.py
