Expose doc breakpoint in cfg

Signed-off-by: Davis Wertheimer <[email protected]>

Author: daviswer

fms_fsdp/config/training.py

@@ -29,6 +29,7 @@ class train_config:
     logical_shards: int = 1024
     num_workers: int = 1
     doc_cutoff: int = 1_000_000
+    doc_breakpoint: int = 65_536
 
     # fsdp policies
     sharding_strategy: str = "hsdp"

fms_fsdp/utils/dataloader_utils.py

@@ -1,4 +1,5 @@
 import torch
+from math import ceil
 
 from fms_fsdp.utils.dataset_utils import (
     ArrowHandler,
@@ -104,6 +105,7 @@ def get_data_loader(cfg, rank, world_size):
         bos_token=cfg.bos_token,
         strip_tokens=set(droplist),
         min_length=3,
+        max_consecutive_chunks=ceil(cfg.doc_breakpoint/1024),
         seed=cfg.seed,
     )
     # Add rescaling/resharding