Make doc slicing flaggable

Author: daviswer

fms_fsdp/config/training.py

@@ -30,6 +30,7 @@ class train_config:
     doc_breakpoint: int = 65_536
     filter_exp: int = 2
     target_doclen: int = 8192
+    slice_rate: float = 0.0
 
     # FIM training
     psm_rate: float = 0.0

fms_fsdp/utils/dataloader_utils.py

@@ -142,11 +142,12 @@ def get_data_loader(cfg, rank, world_size, dp_degree):
     # Shuffle outputs in length 10k buffer. Consecutive lines appear 10k steps apart on average.
     data = PreloadBufferDataset(data, 1000)
     # Slice and rearrange docs to force long-context retrieval
-    data = DocSliceDataset(
-        data,
-        cfg.eos_token,
-        slice_rate=.75,
-    )
+    if cfg.slice_rate > 0:
+        data = DocSliceDataset(
+            data,
+            cfg.eos_token,
+            slice_rate=cfg.slice_rate,
+        )
     # Apply FIM transformation if needed
     if fim_training:
         data = FIMDataset(