faster processing (num_proc + only main preprocess ds)

intx_sft.py

@@ -55,6 +55,7 @@ from transformers import (
     HfArgumentParser,
     set_seed,
 )
+from accelerate import Accelerator, DeepSpeedPlugin
 from peft import PeftModel
 from transformers.utils import is_liger_kernel_available
 from transformers.data import DataCollatorWithFlattening
@@ -252,6 +253,8 @@ def main():
         ]
     )
 
+    accelerator = Accelerator(deepspeed_plugins=DeepSpeedPlugin("ds_config.json"))
+    device = accelerator.device
     set_seed(training_args.seed)
     checkpoint_dir = training_args.resume_from_checkpoint
     if checkpoint_dir and not os.path.isdir(checkpoint_dir):
@@ -309,6 +312,7 @@ def main():
             train=True,
             requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
             peft_config=get_lora_config(model_name, **vars(lora_args)),
+            device=device,
         )
         ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path
         if ctx_name is not None:
@@ -376,6 +380,8 @@ def main():
     logger.debug(model)
     log_num_train_params(model)
 
+    # model = accelerator.prepare(model)
+
     ############ Dataset setup
 
     logger.info("Loading dataset...")
@@ -405,12 +411,10 @@ 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, streaming=streaming
-            )
-            for ds_name in ds_names
-        }
+        for ds_name in ds_names:
+            with accelerator.main_process_first():
+                ds = _get_tokenized_dataset(ds_name, split, streaming=streaming)
+            tokenized_ds[split] = {os.path.basename(ds_name): ds}
 
     train_ds = tokenized_ds["train"]
     logging.info(f"train_ds: {train_ds}")
@@ -575,7 +579,7 @@ if __name__ == "__main__":
     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"
+    os.environ["OMP_NUM_THREADS"] = "23"
     # os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = "137438953472"  # 128 TB
     if os.getenv("DEBUG", False):
         disable_caching()

scripts/fw_qa_3/medium_gemma_per_rank_fac2_per_layer_4_gpus.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+#SBATCH --job-name=ctxlora_medium
+#SBATCH --partition=a3
+#SBATCH --nodes=1
+#SBATCH --gpus=4
+#SBATCH --output=outputs/%x-%j.out
+#SBATCH --error=outputs/%x-%j.out
+
+# module load
+# module load cuda/12.1
+# module load cudnn/8.9.7
+# module load nccl/cuda-12.1/2.18.3
+# module load hpcx/2.20
+
+# export OMP_NUM_THREADS=24
+# export TRITON_CACHE_DIR=/tmp/.triton/
+. ~/miniconda3/etc/profile.d/conda.sh
+conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora
+# eval "$@"
+
+accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \
+--gpu_ids all --main_process_port 29564 intx_sft.py configs/pretrain_all_xl_3_medium.yaml \
+--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=32 \
+--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \
+--target_modules=down_proj \
+--num_self_attends_per_block=8 --num_latent_factor=2 \
+--lora_r=8 \
+--eval_steps=5000 --save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 \
+--neftune_noise_alpha=5 --use_light_weight_lora=False \
+--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \
+--add_repeat_prompt=False \
+--use_sequence_packing=True --per_rank_gen=True \
+--per_layer_processing=True

src/ctx_to_lora/data_utils.py

@@ -410,6 +410,7 @@ def _load_and_process_dataset(
     add_repeat_prompt: bool,
     streaming: bool,
     ds_path: str,
+    num_proc: int,
 ):
     try:
         ds_kwargs = get_ds_kwargs(ds_name, split)
@@ -447,27 +448,27 @@ def _load_and_process_dataset(
     ds = ds.map(
         get_preprocessing_fn(ds_name),
         remove_columns=cols_to_remove,
-        num_proc=16,
+        num_proc=num_proc,
     )
     # 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, num_proc=num_proc)
+    ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc)
     if split == "train":
         if add_negative_prompt:
             ds = ds.map(
                 add_negative_prompt_fn,
                 batched=True,
                 batch_size=100_000,
-                num_proc=16,
+                num_proc=num_proc,
             )
         if add_repeat_prompt and "context_numbers" not in ds_name:
             ds = ds.map(
                 add_repeat_prompt_fn,
                 batched=True,
                 batch_size=100_000,
-                num_proc=16,
+                num_proc=num_proc,
             )
-    ds.save_to_disk(ds_path, num_proc=16)
+    ds.save_to_disk(ds_path, num_proc=num_proc)
     return ds
 
 
@@ -496,6 +497,7 @@ def get_tokenized_dataset(
         add_repeat_prompt=add_repeat_prompt,
         streaming=streaming,
     )
+    num_proc = None if streaming and split == "train" else 8
 
     ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest()
     ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
@@ -508,6 +510,7 @@ def get_tokenized_dataset(
         _load_and_process_dataset(
             **load_and_process_kwargs,
             ds_path=ds_path,
+            num_proc=num_proc,
         )
     ds = datasets.load_from_disk(ds_path)
 
@@ -521,6 +524,7 @@ def get_tokenized_dataset(
         need_ctx_ids,
         ds,
         set_format,
+        num_proc,
     )
     return tokenized_ds
 
@@ -535,6 +539,7 @@ def construct_and_tokenize_ctx_qa(
     need_ctx_ids,
     ds,
     set_format=None,
+    num_proc=None,
 ):
     kwargs = dict(
         tokenizer=repr(tokenizer),
@@ -561,21 +566,21 @@ 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,
+        num_proc=num_proc,
     )
     # add "chat" field
     ds = ds.map(
         get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
         batched=True,
         batch_size=100_000,
-        num_proc=16,
+        num_proc=num_proc,
     )
     # tokenize the chat + mask the assistant inputs
     ds = ds.filter(
         filter_long_chat,
         batched=True,
         batch_size=100_000,
-        num_proc=16,
+        num_proc=num_proc,
     )
 
     # add "input_ids", "attention_mask", "labels"
@@ -586,7 +591,7 @@ def construct_and_tokenize_ctx_qa(
             "mask_assistant_inputs": True,
             "tokenizer_kwargs": tokenizer_kwargs,
         },
-        num_proc=16,
+        num_proc=num_proc,
     )
 
     # for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
@@ -618,7 +623,7 @@ def construct_and_tokenize_ctx_qa(
             fn_kwargs={"tokenizer": ctx_tokenizer},
             batched=True,
             batch_size=100_000,
-            num_proc=16,
+            num_proc=num_proc,
         )
 
     tokenized_ds = tokenized_ds.remove_columns(
@@ -630,7 +635,7 @@ def construct_and_tokenize_ctx_qa(
     #     # the columns are unknown when using streaming dataset
     #     tokenized_ds = tokenized_ds._resolve_features()
     # validate_columns(tokenized_ds)
-    tokenized_ds.save_to_disk(ds_path, num_proc=16)
+    tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
     del tokenized_ds
     return datasets.load_from_disk(ds_path)