Mamba new data fixes2

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