[WIP]: Trainings

configs/deepspeed_config_2.json

@@ -0,0 +1,48 @@
+{
+	"bf16": {
+		"enabled": "auto"
+	},
+	"fp16": {
+		"enabled": "auto",
+		"loss_scale": 0,
+		"loss_scale_window": 1000,
+		"initial_scale_power": 16,
+		"hysteresis": 2,
+		"min_loss_scale": 1
+	},
+	"optimizer": {
+		"type": "AdamW",
+		"params": {
+			"lr": "auto",
+			"betas": "auto",
+			"eps": "auto",
+			"weight_decay": "auto"
+		}
+	},
+	"scheduler": {
+		"type": "WarmupLR",
+		"params": {
+			"warmup_min_lr": "auto",
+			"warmup_max_lr": "auto",
+			"warmup_num_steps": "auto"
+		}
+	},
+	"zero_optimization": {
+		"stage": 2,
+		"offload_optimizer": {
+			"device": "cpu",
+			"pin_memory": true
+		},
+		"allgather_partitions": true,
+		"allgather_bucket_size": 5e8,
+		"overlap_comm": true,
+		"reduce_scatter": true,
+		"reduce_bucket_size": 5e8,
+		"contiguous_gradients": true
+	},
+	"gradient_accumulation_steps": "auto",
+	"gradient_clipping": "auto",
+	"train_batch_size": "auto",
+	"train_micro_batch_size_per_gpu": "auto",
+	"steps_per_print": 1e5
+}

configs/granite_test.yaml

@@ -0,0 +1,27 @@
+model_name: "ibm-granite/granite-7b-base"
+dataset_path: "src/mtl4ad/resources/train/dataset"
+per_device_train_batch_size: 2
+per_device_eval_batch_size: 2
+gradient_checkpointing: true
+gradient_accumulation_steps: 1
+experiment_name: test_no_deepspeed
+checkpoint_dir: "models"
+# deepspeed: "configs/deepspeed_config.json" #"configs/deepspeed_config.json"
+dataset_percentage: 0.2
+n_val_sample: 100
+dataset_seed: 42
+shuffle: true
+enable_peft: true
+eval_steps: 4
+do_eval: true
+logging_steps: 5
+save_total_limit: 10
+save_steps: 5
+eval_strategy: "steps"
+num_train_epochs: 1
+# max_steps: 100
+report_to: "mlflow"
+max_seq_length: 512
+process_group_timeout_in_minutes: 50
+resume_from_checkpoint: "/home/ynanateukam/mtl4ad/models/test_no_deepspeed/checkpoint-15"
+# optim: "adamw_bnb_8bit"

configs/granite_train.yaml

@@ -0,0 +1,27 @@
+model_name: "ibm-granite/granite-7b-base"
+dataset_path: "src/mtl4ad/resources/train/dataset"
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+gradient_checkpointing: true
+gradient_accumulation_steps: 4
+experiment_name: full_train_no_deepspeed
+checkpoint_dir: "models"
+# deepspeed: "configs/deepspeed_config.json" #"configs/deepspeed_config.json"
+dataset_percentage: 100
+n_val_sample: 12000
+dataset_seed: 42
+shuffle: true
+enable_peft: true
+eval_steps: 10000
+logging_steps: 100
+save_total_limit: 6
+save_steps: 1000
+eval_strategy: "steps"
+num_train_epochs: 1
+# max_steps: 100
+report_to: "none"
+max_seq_length: 4096
+ddp_backend: "nccl"
+process_group_timeout_in_minutes: 50
+# resume_from_checkpoint: "/home/ynanateukam/mtl4ad/models/test_no_deepspeed/checkpoint-50"
+# optim: "adamw_bnb_8bit"

pyproject.toml

@@ -42,7 +42,7 @@ priority = "primary"
 python = "^3.10"
 absl-py = '^2.1.0'
 datasets = '^2.18.0'
-deepspeed = '0.13.1'  # to avoid serialization errors when saving checkpoints
+deepspeed = "0.14.4"  # to avoid serialization errors when saving checkpoints
 evaluate = '^0.4.1'
 loguru = '^0.7.2'
 mlflow = '^2.11.1'
@@ -53,12 +53,13 @@ rouge-score = '^0.1.2'
 tqdm = '^4.25.0'
 importlib-resources = "^6.3.1"
 pandas = "^2.2.1"
-transformers = "^4.39.3"
 sacrebleu = "^2.4.2"
 
 python-levenshtein = "0.25.1"
-trl = "^0.9.4"
 bitsandbytes = "^0.43.1"
+tokenizers = "^0.19.1"
+transformers = "^4.44.0"
+accelerate = "<0.30.0"
 [tool.poetry.dev-dependencies]
 jupyter = "^1.0.0"
 mypy = "^1.0.0"

src/mtl4ad/data_preprocessing.py

@@ -1,8 +1,9 @@
 """Dataset Utilities: Loading, Preprocessing, and Filtering"""
 
 import logging
+import multiprocessing as mp
 from pathlib import Path
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from datasets import Dataset, DatasetDict, concatenate_datasets, load_dataset
 from transformers import PreTrainedTokenizer
@@ -14,6 +15,34 @@
 )
 logger = logging.getLogger(__name__)
 
+def load_dataset_from_folder(folder: Path) -> Dict[str, Optional[Dataset]]:
+    """Load train and validation datasets from a single folder.
+
+    Args:
+        folder: The folder to load the datasets from.
+
+    Returns:
+        A dictionary containing 'train' and 'validation' datasets.
+    """
+    datasets = {"train": None, "validation": None}
+
+    train_file = folder / "train.parquet"
+    valid_file = folder / "validation.parquet"
+
+    if train_file.exists():
+        train_dataset = load_dataset(
+            "parquet", data_files=str(train_file), split="train", num_proc=mp.cpu_count(), cache_dir="/projects/0/prjs1108/.cache"
+        )
+        datasets["train"] = train_dataset
+
+    if valid_file.exists():
+        valid_dataset = load_dataset(
+            "parquet", data_files=str(valid_file), split="train", num_proc=mp.cpu_count(), cache_dir="/projects/0/prjs1108/.cache"
+        )
+        datasets["validation"] = valid_dataset
+
+
+    return datasets
 
 def load_dataset_from_folders(base_path: Union[str, Path]) -> DatasetDict:
     """Load and concatenate datasets from multiple folders.
@@ -24,31 +53,23 @@ def load_dataset_from_folders(base_path: Union[str, Path]) -> DatasetDict:
     Returns:
         The concatenated dataset with 'train' and 'validation' splits.
     """
-    all_train_datasets = []
-    all_valid_datasets = []
-
     base_path = Path(base_path) if isinstance(base_path, str) else base_path
-    for folder in base_path.iterdir():
-        if folder.is_dir():
-            train_file = folder / "train.parquet"
-            valid_file = folder / "validation.parquet"
-
-            if train_file.exists():
-                train_dataset = load_dataset(
-                    "parquet", data_files=str(train_file), split="train", num_proc=32
-                )
-                all_train_datasets.append(train_dataset)
-
-            if valid_file.exists():
-                valid_dataset = load_dataset(
-                    "parquet", data_files=str(valid_file), split="train", num_proc=32
-                )
-                all_valid_datasets.append(valid_dataset)
+
+    folders = [folder for folder in base_path.iterdir() if folder.is_dir()]
+
+    with mp.Pool(mp.cpu_count()) as pool:
+        results = pool.map(load_dataset_from_folder, folders)
+
+    all_train_datasets = [res["train"] for res in results if res["train"] is not None]
+    all_valid_datasets = [res["validation"] for res in results if res["validation"] is not None]
 
     train_dataset = concatenate_datasets(all_train_datasets) if all_train_datasets else None
     valid_dataset = concatenate_datasets(all_valid_datasets) if all_valid_datasets else None
 
     dataset = DatasetDict({"train": train_dataset, "validation": valid_dataset})
+    for entry in dataset:
+        dataset[entry] = dataset[entry].rename_column("target", "labels")
+        dataset[entry] = dataset[entry].rename_column("source", "text")
 
     return dataset
 
@@ -80,21 +101,22 @@ def process(examples: dict) -> dict:
         """
         formatted_prompts = [
             f"{source} \n\n {target}"
-            for source, target in zip(examples["source"], examples["target"])
+            for source, target in zip(examples["text"], examples["labels"])
         ]
         model_inputs = tokenizer(
             formatted_prompts,
             truncation=True,
             max_length=max_length,
             return_tensors=None,
+            padding="max_length"
         )
         model_inputs["labels"] = model_inputs["input_ids"].copy()
         return model_inputs
 
-    return dataset.map(process, batched=True)
+    return dataset.map(process, batched=True, num_proc=mp.cpu_count())
 
 
-def filter_dataset(dataset: DatasetDict, dataset_percentage: Optional[float]) -> DatasetDict:
+def filter_dataset(dataset: DatasetDict, dataset_percentage: Optional[float], n_val_sample: Optional[int]) -> DatasetDict:
     """
     Filters the dataset based on the given percentage.
 
@@ -109,39 +131,26 @@ def filter_dataset(dataset: DatasetDict, dataset_percentage: Optional[float]) ->
         raise ValueError("dataset_percentage cannot be None")
 
     train_new_size = round(len(dataset["train"]) * dataset_percentage / 100)
-    val_new_size = round(len(dataset["validation"]) * dataset_percentage / 100)
     dataset["train"] = dataset["train"].select(range(train_new_size))
-    dataset["validation"] = dataset["validation"].select(range(val_new_size))
+    if not n_val_sample:
+        val_new_size = round(len(dataset["validation"]) * dataset_percentage / 100)
+        dataset["validation"] = dataset["validation"].select(range(val_new_size))
+    else:
+        dataset["validation"] = dataset["validation"].shuffle(seed=42)
+        dataset["validation"] = dataset["validation"].take(n_val_sample)
     return dataset
 
 
-def load_dataset_and_preprocess(config: Dict[str, Any], tokenizer) -> Dataset:
-    """
-    Loads and preprocesses the dataset according to the given configuration.
+def generate_formatted_prompts(examples) -> List[str]:
+    """Generate formatted prompts from examples.
 
     Args:
-        config: The configuration object containing dataset and preprocessing parameters.
-        tokenizer: The tokenizer to be used for preprocessing the dataset.
+        examples: A dictionary with 'source' and 'target' lists.
 
     Returns:
-        The preprocessed dataset.
+        A list of formatted prompts.
     """
-    logger.info("Loading dataset...")
-    dataset = load_dataset_from_folders(config["dataset_path"])
-    for element in dataset:
-        logger.info(f"Original size of {element}: {len(dataset[element])}")
-
-    if "dataset_percentage" in config:
-        logger.info("Sampling dataset...")
-        dataset = filter_dataset(dataset, config.get("dataset_percentage"))
-        for element in dataset:
-            logger.info(f"Filtered size of {element}: {len(dataset[element])}")
-
-    dataset = preprocess(dataset, tokenizer, config.get("model_max_length", None))
-
-    if "shuffle" in config:
-        logger.info("Shuffling dataset...")
-        dataset = dataset.shuffle(seed=config.get("dataset_seed", 42))
-        logger.info("Done shuffling")
-
-    return dataset
+    return [
+        f"{source} \n\n {target}"
+        for source, target in zip(examples["text"], examples["labels"])
+    ]