Dataloader updates and streamlining

fms_fsdp/utils/dataloader_utils.py

@@ -1,12 +1,13 @@
 import torch
 
 from fms_fsdp.utils.dataset_utils import (
-    Buffer_Dataset,
-    Checkpoint_Dataset,
-    Preload_Buffer_Dataset,
-    Preprocess_Dataset,
-    Sampling_Dataset,
-    Scalable_Shard_Dataset,
+    BufferDataset,
+    CheckpointDataset,
+    PreloadBufferDataset,
+    PreprocessDataset,
+    SamplingDataset,
+    ScalableShardDataset,
+    StreamingDocDataset,
 )
 
 
@@ -68,35 +69,46 @@ def causal_lm(data_seq, prompt_len=0):
         int(x.strip()) for x in cfg.strip_tokens.split(",") if len(x.strip()) > 0
     ]
     droplist = droplist + [cfg.bos_token, cfg.eos_token, cfg.bol_token, cfg.eol_token]
-    data = Sampling_Dataset(
+    # Base reader layer
+    data = StreamingDocDataset(
         cfg.data_path,
-        Scalable_Shard_Dataset,
         rank,
         world_size,
         cfg.eos_token,
         bos_token=cfg.bos_token,
         strip_tokens=set(droplist),
         min_length=3,
+        seed=cfg.seed,
+    )
+    # Add rescaling/resharding
+    data = ScalableShardDataset(
+        data,
+        cfg.eos_token,
+        n_logical_shards=cfg.logical_shards,
+    )
+    # Add multi-dataset handling
+    data = SamplingDataset(
+        cfg.data_path,
+        data,
+        cfg.eos_token,
         datasets=datasets,
         weights=weights,
-        seed=cfg.seed,
         verbose=(rank == 0),
-        n_logical_shards=cfg.logical_shards,
     )
     # Wrap above dataset in packing logic to form constant-length lines.
-    data = Buffer_Dataset(
+    data = BufferDataset(
         data,
         cfg.seq_length + 1,
         bos_token=cfg.bol_token,
         eos_token=cfg.eol_token,
         pack_hard=True,
     )
     # Shuffle outputs in length 10k buffer. Consecutive lines appear 10k steps apart on average.
-    data = Preload_Buffer_Dataset(data, 10000)
+    data = PreloadBufferDataset(data, 10000)
     # Split line into input and target for the CLM task.
-    data = Preprocess_Dataset(data, causal_lm)
+    data = PreprocessDataset(data, causal_lm)
     # Enable auto-saving
-    data = Checkpoint_Dataset(
+    data = CheckpointDataset(
         data,
         cfg.ckpt_load_path,
         cfg.checkpoint_interval,