allow streaming data

configs/pretrain_all_xl_3.yaml

@@ -37,18 +37,19 @@ target_modules:
   - down_proj
 # data
 train_ds_names:
-  - fw_qa_3_mini
-  - ctx_qa
-  - pwc
-  - hotpot_qa
-  - squad
-  - drop
-  - narrativeqa
-  - quoref
-  - ropes
-  - synthetic_convqa
+  - fw_qa_3 # ~ 267M
+  - ctx_qa # 300k
+  - pwc # 240k
+  - hotpot_qa # 90k
+  - squad # 90k
+  - drop # 77k
+  - narrativeqa # 40k
+  - quoref # 11k
+  - ropes # 11k
+  - synthetic_convqa # 40k
 
 val_ds_names:
+  - fw_qa_3
   - fw_qa_xl
   - ctx_qa
   - pwc

intx_sft.py

@@ -3,6 +3,7 @@ import os
 import random
 import string
 import time
+from math import ceil
 from collections import defaultdict
 from copy import copy, deepcopy
 from functools import partial
@@ -15,6 +16,7 @@ import wandb
 import yaml
 
 from ctx_to_lora.data_utils import (
+    DS_LEN,
     convert_ctx_prompt_response_to_messages,
     get_preprocessing_fn,
     get_sft_prompt_formatting_fn,
@@ -282,7 +284,7 @@ def main():
         training_args.lr_scheduler_type == "cosine_with_min_lr"
         and training_args.lr_scheduler_kwargs is None
     ):
-        training_args.lr_scheduler_kwargs = {"min_lr": 1e-7}
+        training_args.lr_scheduler_kwargs = {"min_lr": 1e-8}
     args = {
         **vars(deepcopy(data_args)),
         **vars(deepcopy(ctx_args)),
@@ -380,6 +382,7 @@ def main():
     add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
     tokenizer_kwargs = {"max_length": ctx_args.max_base_len}
     ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len}  # not used for now
+
     _get_tokenized_dataset = partial(
         get_tokenized_dataset,
         tokenizer=tokenizer,
@@ -400,12 +403,16 @@ def main():
     ):
         if not ds_names:
             continue
+        streaming = (split == "train") and data_args.streaming
         tokenized_ds[split] = {
-            os.path.basename(ds_name): _get_tokenized_dataset(ds_name, split)
+            os.path.basename(ds_name): _get_tokenized_dataset(
+                ds_name, split, streaming=streaming
+            )
             for ds_name in ds_names
         }
 
     train_ds = tokenized_ds["train"]
+    logging.info(f"train_ds: {train_ds}")
     val_ds = dict()
     if "validation" in tokenized_ds:
         n_val_samples = data_args.max_val_samples_per_ds
@@ -419,35 +426,31 @@ def main():
             val_indices = np.random.permutation(len(ds))[:n_val_samples]
             val_ds[ds_name] = val_ds[ds_name].select(val_indices)
 
-    # train_ds = concatenate_datasets(list(train_ds.values()))
+    max_steps = ceil(
+        sum([DS_LEN[ds] for ds in train_ds])
+        * training_args.num_train_epochs
+        / training_args.per_device_train_batch_size
+        / training_args.gradient_accumulation_steps
+        / training_args.world_size
+    )
+    training_args.max_steps = max_steps
 
-    # total_len = sum(len(ds) for ds in train_ds.values())
-    train_ds_len = [len(ds) for ds in train_ds.values()]
-    total_len = sum(train_ds_len)
+    # interleaving streaming datasets
+    # simplify the probs for smaller datasets
+    # slightly upsample those datasets
+    probs = [
+        0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
+        for ds_name in train_ds
+    ]
     train_ds = interleave_datasets(
         list(train_ds.values()),
-        probabilities=[l / total_len for l in train_ds_len],
+        probabilities=probs,
+        stopping_strategy="all_exhausted",
         seed=training_args.seed,
     )
 
-    # val_train_indices = np.random.permutation(len(train_ds))[:500]
-    # val_ds["train"] = train_ds.select(val_train_indices)
-
-    test_ds = dict()
-    if "test" in tokenized_ds:
-        n_test_samples = data_args.max_test_samples_per_ds
-        for ds_name, ds in tokenized_ds["test"].items():
-            test_ds[ds_name] = ds
-            test_indices = np.random.permutation(len(ds))[:n_test_samples]
-            test_ds[ds_name] = test_ds[ds_name].select(test_indices)
-
     logger.info(f"train_ds: {train_ds}")
     logger.info(f"val_ds: {val_ds}")
-    logger.info(f"test_ds: {test_ds}")
-
-    # TODO: change to a faster collator? e.g.,
-    # https://huggingface.co/blog/packing-with-FA2
-    # data_collator = DataCollatorForSeq2Seq(tokenizer, model, pad_to_multiple_of=8)
     flattener = DataCollatorWithFlattening()
 
     def train_packed_collator(inp_list):
@@ -549,7 +552,6 @@ def main():
         training_args,
         train_ds,
         val_ds,
-        test_ds,
         train_collator,
         # partial(generation_collator, tokenizer=tokenizer),
         compute_metrics=partial(
@@ -571,6 +573,8 @@ if __name__ == "__main__":
     os.environ["WANDB_WATCH"] = ""  # "all"
     os.environ["WANDB_CONSOLE"] = "off"
     os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    os.environ["KMP_AFFINITY"] = "disabled"  # fixing iterable dataset stuck
+    os.environ["OMP_NUM_THREADS"] = "16"
     if os.getenv("DEBUG", False):
         disable_caching()
     # randomly sleep to avoid run_name collision

src/ctx_to_lora/configs.py

@@ -338,6 +338,10 @@ class DataArguments:
         default=None,
         metadata={"help": "Training dataset names."},
     )
+    streaming: bool = field(
+        default=False,
+        metadata={"help": "Whether to use streaming dataset for training."},
+    )
     val_ds_names: Optional[list[str]] = field(
         default=None,
         metadata={"help": "Validation dataset names."},

src/ctx_to_lora/data_utils.py

@@ -1,4 +1,5 @@
 import logging
+from os import path
 import numpy as np
 from glob import glob
 from typing import Any, Callable, Iterator, Optional
@@ -14,6 +15,24 @@ logger = logging.getLogger()
 FW_QA_PATHS = [
     f"data/raw_datasets/fw_qa/{i:05d}.parquet" for i in [0, 1, 6, 7, 8, 10, 22, 30, 35]
 ]
+
+# approximate length of the datasets
+# needed for streaming datasets
+DS_LEN = {
+    "fw_qa_3_mini": 100_000,
+    "fw_qa_3": 270_000_000,
+    "fw_qa_xl": 27_000_000,
+    "ctx_qa": 300_000,
+    "pwc": 240_000,
+    "hotpot_qa": 100_000,
+    "squad": 90_000,
+    "drop": 77_000,
+    "narrativeqa": 40_000,
+    "quoref": 11_000,
+    "ropes": 11_000,
+    "synthetic_convqa": 40_000,
+}
+
 DS_KWARGS = {
     "hotpot_qa": dict(
         train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"),
@@ -101,10 +120,22 @@ DS_KWARGS = {
     "fw_qa_3_mini": dict(
         train=dict(
             path="parquet",
-            data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"),
+            data_files=glob("data/raw_datasets/fw_qa_3/000_00001.parquet"),
             split="train[:100000]",
         ),
     ),
+    "fw_qa_3": dict(
+        train=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"),
+            split="train",
+        ),
+        validation=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_3/*val.parquet"),
+            split="train",
+        ),
+    ),
     "ctx_qa": dict(
         train=dict(
             path="parquet",
@@ -225,8 +256,24 @@ def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]:
             f"No dataset kwargs found for '{ds_name}' with split '{split}'.\n"
             f"Using default kwargs: {kwargs}"
         )
-        return kwargs
-    return DS_KWARGS[ds_name][split]
+        # return kwargs
+    else:
+        # return DS_KWARGS[ds_name][split]
+        kwargs = DS_KWARGS[ds_name][split]
+
+    # custom logic for slicing iterable datasets
+    take, skip = None, None
+    kw_split = kwargs["split"]
+    if ("[" in kw_split) and kw_split.endswith("]"):
+        kwargs["split"], slice = kw_split.split("[")
+        slice = slice.strip("]")
+        skip = slice.split(":")[0]
+        if skip:
+            kwargs["skip"] = int(skip)
+        take = slice.split(":")[1]
+        if take:
+            kwargs["take"] = int(take)
+    return kwargs
 
 
 def validate_columns(tokenized_ds):
@@ -354,12 +401,24 @@ def get_tokenized_dataset(
     add_negative_prompt: bool,
     use_kl_loss: bool,
     set_format: Optional[str] = None,
+    streaming: bool = False,
 ) -> dict[str, Any]:
 
     logger.debug(f"Loading dataset {ds_name} with split {split}...")
     need_ctx_ids = not add_ctx_to_chat
     try:
-        ds = load_dataset(**get_ds_kwargs(ds_name, split), trust_remote_code=True)
+        ds_kwargs = get_ds_kwargs(ds_name, split)
+        skip = ds_kwargs.pop("skip", None)
+        take = ds_kwargs.pop("take", None)
+        ds = load_dataset(
+            **get_ds_kwargs(ds_name, split),
+            trust_remote_code=True,
+            streaming=streaming,
+        )
+        if skip is not None:
+            ds = ds.skip(skip)
+        if take is not None:
+            ds = ds.take(take)
     except ValueError as e:
         logger.info(
             f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..."
@@ -368,10 +427,10 @@ def get_tokenized_dataset(
     cols_to_remove = [
         col for col in ds.column_names if col not in ["context", "prompt", "response"]
     ]
-    ds = ds.map(get_preprocessing_fn(ds_name), num_proc=16)
+    ds = ds.map(get_preprocessing_fn(ds_name))
     ds = ds.remove_columns(cols_to_remove)
-    ds = ds.filter(filter_none, batched=True, num_proc=16)
-    ds = ds.filter(filter_long_samples, batched=True, num_proc=16)
+    ds = ds.filter(filter_none, batched=True)
+    ds = ds.filter(filter_long_samples, batched=True)
     if split == "train":
         if add_negative_prompt:
             ds = ds.map(add_negative_prompt_fn, batched=True)
@@ -407,14 +466,11 @@ def construct_and_tokenize_ctx_qa(
     ds = ds.map(
         convert_ctx_prompt_response_to_messages,
         fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
-        num_proc=16,
     )
     # add "chat" field
-    ds = ds.map(
-        get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), num_proc=16
-    )
+    ds = ds.map(get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer))
     # tokenize the chat + mask the assistant inputs
-    ds = ds.filter(filter_long_chat, batched=True, num_proc=16)
+    ds = ds.filter(filter_long_chat, batched=True)
 
     # add "input_ids", "attention_mask", "labels"
     tokenized_ds = ds.map(
@@ -424,7 +480,6 @@ def construct_and_tokenize_ctx_qa(
             "mask_assistant_inputs": True,
             "tokenizer_kwargs": tokenizer_kwargs,
         },
-        num_proc=16,
     )
 
     # for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
@@ -447,25 +502,25 @@ def construct_and_tokenize_ctx_qa(
                 "tokenizer_kwargs": tokenizer_kwargs,
                 "for_kl_loss": True,
             },
-            num_proc=16,
         )
 
     if need_ctx_ids:
-        # TODO: can we batch this?
-        # TODO: can we cache this?
         # tokenize the ctx_text to get ctx_ids and ctx_attn_mask
         tokenized_ds = tokenized_ds.map(
             tokenize_ctx_text,
             fn_kwargs={"tokenizer": ctx_tokenizer},
             batched=True,
-            num_proc=16,
         )
 
     tokenized_ds = tokenized_ds.remove_columns(
         ["messages", "chat", "context", "prompt", "response"]
     )
-    tokenized_ds.set_format(type=set_format)
-    validate_columns(tokenized_ds)
+    if set_format:
+        tokenized_ds.set_format(type=set_format)
+    # if isinstance(tokenized_ds, IterableDataset):
+    #     # the columns are unknown when using streaming dataset
+    #     tokenized_ds = tokenized_ds._resolve_features()
+    # validate_columns(tokenized_ds)
     return tokenized_ds
 
 

src/ctx_to_lora/modeling_utils.py

@@ -1212,6 +1212,9 @@ class ModulatedPretrainedModel(nn.Module):
                     A = A[0:1]
                     B = B[:, 0:1]
                 biases = [A, B.T]
+                # TODO: scale the biases by the output size?
+                # e.g., rank=16 gives bigger (2x?) gradient magnitudes at initialization
+                # compared to rank=8
 
                 # bias-hyperinit
                 # init weights to zeros and bias to the base weights

src/ctx_to_lora/training_utils.py

@@ -99,7 +99,7 @@ def train_model(
     training_args,
     train_dataset=None,
     val_dataset=None,
-    test_dataset=None,
+    # test_dataset=None,
     train_collator=None,
     # generation_collator=None,
     compute_metrics=None,