BUG: fix bugs in lazy loader.

Author: msafari18

barcodebert/datasets.py

@@ -11,6 +11,8 @@
 from torch.utils.data import Dataset, IterableDataset
 from torchtext.vocab import vocab as build_vocab_from_dict
 from transformers import AutoTokenizer
+import torch.distributed as dist
+from torch.utils.data import get_worker_info
 
 
 class KmerTokenizer(object):
@@ -151,14 +153,44 @@ def parse_row(self, row):
         tokens, att_mask = self.tokenizer(dna_seq, offset=offset)
         return tokens, torch.tensor(int(label), dtype=torch.int64), att_mask
 
+    def __len__(self):
+        # count lines once at startup (cheap) so val‑throughput logging doesn't crash
+        if not hasattr(self, "_n_samples"):
+            with open(self.file_path, "r") as f:
+                # subtract header if CSV has one; adjust accordingly
+                self._n_samples = sum(1 for _ in f) - 1
+        return self._n_samples
+
+
     def __iter__(self):
+        # Determine global rank & world size
+        if dist.is_available() and dist.is_initialized():
+            rank = dist.get_rank()
+            world_size = dist.get_world_size()
+        else:
+            rank, world_size = 0, 1
+
+        # If we're also using multiple DataLoader workers (num_workers > 1),
+        # further subdivide per-worker:
+        worker_info = get_worker_info()
+        if worker_info is not None:
+            # each worker in the same process gets a unique ID
+            worker_id = worker_info.id
+            total_workers = worker_info.num_workers
+            # flatten ranks+workers into a single shard index
+            rank = rank * total_workers + worker_id
+            world_size = world_size * total_workers
+
+        # Now stream the file, and only yield rows where idx % world_size == rank
         df_iter = pd.read_csv(
             self.file_path,
             sep="\t" if self.file_path.endswith(".tsv") else ",",
             chunksize=1,
             keep_default_na=False,
         )
-        for chunk in df_iter:
+        for idx, chunk in enumerate(df_iter):
+            if idx % world_size != rank:
+                continue
             yield self.parse_row(chunk.iloc[0])
 
 

barcodebert/pretraining.py

@@ -181,46 +181,48 @@ def print_pass(*args, **kwargs):
 
     eval_set = "Val"
 
-    # Dataloader --------------------------------------------------------------
-    dl_train_kwargs = {
-        "batch_size": config.batch_size_per_gpu,
-        "drop_last": True,
-        "sampler": None,
-        "shuffle": True,
-        "worker_init_fn": utils.worker_seed_fn,
-    }
-    dl_val_kwargs = {
-        "batch_size": config.batch_size_per_gpu,
-        "drop_last": False,
-        "sampler": None,
-        "shuffle": False,
-        "worker_init_fn": utils.worker_seed_fn,
-    }
-    if config.cpu_workers is None:
-        config.cpu_workers = utils.get_num_cpu_available()
-    if use_cuda:
-        cuda_kwargs = {"num_workers": config.cpu_workers, "pin_memory": True}
-        dl_train_kwargs.update(cuda_kwargs)
-        dl_val_kwargs.update(cuda_kwargs)
-
-    if config.distributed:
-        # The DistributedSampler breaks up the dataset across the GPUs
-        dl_train_kwargs["sampler"] = DistributedSampler(
-            dataset_train,
-            shuffle=True,
-            seed=config.seed if config.seed is not None else 0,
-            drop_last=False,
-        )
-        dl_train_kwargs["shuffle"] = None
-        dl_val_kwargs["sampler"] = DistributedSampler(
-            dataset_val,
-            shuffle=False,
-            drop_last=False,
-        )
-        dl_val_kwargs["shuffle"] = None
-
-    dataloader_train = torch.utils.data.DataLoader(dataset_train, **dl_train_kwargs)
-    dataloader_val = torch.utils.data.DataLoader(dataset_val, **dl_val_kwargs)
+    # Dataloaders -------------------------------------------------------------
+    if config.lazy_load:
+        # streaming IterableDataset → no sampler, no shuffle
+        stream_kwargs = {
+            "batch_size": config.batch_size_per_gpu,
+            "drop_last": True,
+            "num_workers": config.cpu_workers,
+            "pin_memory": use_cuda,
+            "worker_init_fn": utils.worker_seed_fn,
+        }
+        dataloader_train = torch.utils.data.DataLoader(dataset_train, **stream_kwargs)
+        dataloader_val = torch.utils.data.DataLoader(dataset_val, **stream_kwargs)
+    else:
+        # map‑style Dataset → use DistributedSampler in dist. mode
+        map_train_kwargs = {
+            "batch_size": config.batch_size_per_gpu,
+            "drop_last": True,
+            "shuffle": True,
+            "num_workers": config.cpu_workers,
+            "pin_memory": use_cuda,
+            "worker_init_fn": utils.worker_seed_fn,
+        }
+        map_val_kwargs = {
+            "batch_size": config.batch_size_per_gpu,
+            "drop_last": False,
+            "shuffle": False,
+            "num_workers": config.cpu_workers,
+            "pin_memory": use_cuda,
+            "worker_init_fn": utils.worker_seed_fn,
+        }
+        if config.distributed:
+            map_train_kwargs["shuffle"] = False
+            map_train_kwargs["sampler"] = DistributedSampler(
+                dataset_train, shuffle=True,
+                seed=(config.seed or 0),
+                drop_last=False,
+            )
+            map_val_kwargs["sampler"] = DistributedSampler(
+                dataset_val, shuffle=False, drop_last=False
+            )
+        dataloader_train = torch.utils.data.DataLoader(dataset_train, **map_train_kwargs)
+        dataloader_val = torch.utils.data.DataLoader(dataset_val, **map_val_kwargs)
 
     # MODEL ===================================================================
     base_pairs = "ACGT"