Mamba new data fixes2 (#149)

* make mamba

* add quick debug

* add quick debug

* revert debug verbosity

* Learning rate scheduler changed (Constant)

* Add AutoHandler

* Add Auto cfg option for AutoHAndler

* Len gets called before open

* path/filepath typo fix

* Partitioning fix from mup-search

* Cosine 0.01 decay

* Warmup interval change

* Schedule change

* Constant schedule

* LR schedule change (cool down and constant lr)

* Update dataset_utils.py

Added a check for length of doc

* LR schedule change (Warmup + constant)

* Update dataset_utils.py

* Cosine schedule

* For constant lr 1.5e5

* Schedule change

* Schedule change

* Final singlefile checkpoint saves one folder up (#127)

* Final singlefile checkpoint saves one folder up

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

* save file under new pth subfolder

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

* Repath for easier consumption/conversion

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

---------

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

* Added cool down

* length of doc check

* splitstrip cols and pass to fhandler

* fhandler col_names support

* Warmup for annealing

* Debugging

* Debugging II

* Empty shard check

* Added constant lr schedule with warmup

* added print for lenght of doc

* added print for lenght of doc II

* Update dataset_utils.py

* Update dataset_utils.py

* Update dataset_utils.py

* Update dataset_utils.py

* Adding print for debug

* Revert "Pulled from data-fixes branch"

This reverts commit ac5194b, reversing
changes made to 1b50708.

reverting changes

* Revert all changes made after March 6 (before merge)

* Revert all changes made after March 6 (before merge)

* removed print

---------

Signed-off-by: Davis Wertheimer <[email protected]>
Co-authored-by: Linsong Chu <[email protected]>
Co-authored-by: divykum2 <[email protected]>
Co-authored-by: divya-kumari32 <[email protected]>

Author: 3 people

fms_fsdp/config/training.py

@@ -76,10 +76,3 @@ class train_config:
     stage2_prompt_length: int = 64
     stage2_batch_size: int = 96
     stage2_seq_length: int = 256
-
-    # FIM training
-    psm_rate: float = 0.0
-    spm_rate: float = 0.0
-    fim_pre: int = 1
-    fim_mid: int = 2
-    fim_suf: int = 3

fms_fsdp/utils/checkpointing_utils.py

@@ -324,7 +324,9 @@ def save_single_file(
     ):
         # Note: metadata kwargs cannot contain any of:
         # (step, model)
-        save_name = os.path.join(self.ckp_path, "step_" + str(step) + "_ckp.pth")
+        pth_path = os.path.join(self.ckp_path[:-12], "pth", "step_" + str(step))
+        os.makedirs(pth_path, exist_ok=True)
+        save_name = os.path.join(pth_path, "consolidated.00.pth")
         save_time = time.time()
         with FSDP.state_dict_type(
             model,

fms_fsdp/utils/dataloader_utils.py

@@ -6,7 +6,6 @@
     AutoHandler,
     BufferDataset,
     CheckpointDataset,
-    FIMDataset,
     ParquetHandler,
     PreloadBufferDataset,
     PreprocessDataset,
@@ -59,9 +58,9 @@ def __iter__(self):
     return torch.utils.data.DataLoader(data, batch_size=cfg.batch_size)
 
 
-def get_data_loader(cfg, rank, world_size):
+def get_data_loader(cfg, rank, world_size, postprocess=[causal_lm]):
     """
-    Pytorch dataloader for stateful, distributed, and rescalable language model training.
+    Pytorch dataloader for stateful, distributed, and rescalable causal language model (CLM) training.
     Assumes underlying data is sequences of integer values.
     ...
     Args
@@ -72,12 +71,11 @@ def get_data_loader(cfg, rank, world_size):
         Rank of current distributed worker. Used for handling dataset sharding logic.
     world_size : int
         Number of distributed workers. Used for handling dataset sharding logic.
+    postprocess : List[Callable]
+        Any task-specific postprocessing to apply before handing over data. Steps will apply in
+        the order provided by the user. For CLM training, use postprocess=[causal_lm].
     """
 
-    fim_training = cfg.psm_rate + cfg.spm_rate > 0
-    if fim_training:
-        assert cfg.bos_token is None, "No BOS in FIM training. Did you mean fim_pre?"
-
     datasets, weights, cols = parse_data_args(cfg.datasets, cfg.weights, cfg.col_name)
 
     # Base streaming dataset. Returns doc chunks in sequence.
@@ -94,7 +92,7 @@ def get_data_loader(cfg, rank, world_size):
             cfg.tokenizer_path, cols, cfg.doc_cutoff
         )
     else:
-        filehandler = _handler_map[cfg.file_type](cols)
+        filehandler = _handler_map[cfg.file_type, cols]
     # Base reader layer
     data = StreamingDocDataset(
         cfg.data_path,
@@ -124,34 +122,20 @@ def get_data_loader(cfg, rank, world_size):
         verbose=(rank == 0),
     )
     # Wrap above dataset in packing logic to form constant-length lines.
-    # Increment seq len to counteract CLM's one token removal.
     data = BufferDataset(
         data,
-        cfg.seq_length + 1,
+        cfg.seq_length if causal_lm not in postprocess else 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 = PreloadBufferDataset(data, 10000)
 
-    # Apply FIM transformation if needed
-    if fim_training:
-        data = FIMDataset(
-            data,
-            cfg.eos_token,
-            cfg.psm_rate,
-            cfg.spm_rate,
-            pre_token=cfg.fim_pre,
-            mid_token=cfg.fim_mid,
-            suf_token=cfg.fim_suf,
-        )
-
-    # Transform to tensors
+    # Apply desired postprocessing steps in sequence
     data = PreprocessDataset(data, torch.IntTensor)
-
-    # Apply CLM transformation
-    data = PreprocessDataset(data, causal_lm)
+    for p in postprocess:
+        data = PreprocessDataset(data, p)
 
     # Enable auto-saving
     data = CheckpointDataset(
@@ -181,4 +165,4 @@ def splitstrip(x):
     datas = splitstrip(datas)
     weights = [float(x) for x in splitstrip(weights)]
     cols = splitstrip(cols)
-    return datas, weights, cols
+    return datas, weights, cols