Flatten 1_raw_dataset submodules into plain tracked files

FocalLoRA, Should-It-Be-Executed-Or-Processed, and topicattack were
nested git repos (with an inner FocalLoRA/data/FocalLoRA/.git as well).
Drop their .git history and track the contents directly in this repo
instead of as submodules/gitlinks.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Codes/1_raw_dataset/Should-It-Be-Executed-Or-Processed/fine-tuning/fsdp_lora_dafaults.yaml

@@ -0,0 +1,28 @@
+model_id: "meta-llama/Meta-Llama-3-8B-Instruct"
+dataset_path: "."
+max_seq_len: 3072 # 2048
+
+output_dir: "./checkpoints/llama-3-8b-sep-qlora"
+report_to: "wandb"
+learning_rate: 0.0002
+lr_scheduler_type: "constant"
+num_train_epochs: 3
+per_device_train_batch_size: 1
+per_device_eval_batch_size: 1
+gradient_accumulation_steps: 2
+optim: adamw_torch
+logging_steps: 10
+save_strategy: epoch
+evaluation_strategy: epoch
+max_grad_norm: 0.3
+warmup_ratio: 0.03
+bf16: true
+tf32: true
+gradient_checkpointing: true
+
+#fsdp: "full_shard auto_wrap offload"
+fsdp: "full_shard auto_wrap"
+fsdp_config:
+  backward_prefetch: "backward_pre"
+  forward_prefetch: "false"
+  use_orig_params: "false"

Codes/1_raw_dataset/Should-It-Be-Executed-Or-Processed/fine-tuning/prepare_dataset.py

@@ -0,0 +1,64 @@
+import os.path
+
+import fire
+import pandas as pd
+
+
+def main(data_path: str, model_type: str, out_dir: str) -> None:
+    """
+    Process the input data, generate text based on the specified model type, and save the 
+    resulting datasets into training and testing sets.
+
+    Parameters:
+    data_path (str): The path to the input JSON file containing the data.
+    model_type (str): The type of model to be used for generating text. This should be one 
+                      of the keys in the `map_funcs` dictionary.
+    out_dir (str): The directory where the output JSON files will be saved.
+
+    The function performs the following steps:
+    1. Reads the input JSON file into a pandas DataFrame.
+    2. Retains only the 'inputs', 'goal_safe', and 'goal_unsafe' columns.
+    3. Extracts 'prompt_1' and 'prompt_2' from the 'inputs' column and removes 'inputs'.
+    4. Applies a text generation function based on the specified `model_type`.
+    5. Splits the data into training (80%) and testing (20%) sets.
+    6. Prints two random samples from the training set.
+    7. Saves the training and testing sets as JSON files in the specified output directory.
+
+    Raises:
+    KeyError: If `model_type` is not one of the predefined keys in the `map_funcs` dictionary.
+    """
+    df = pd.read_json(data_path)
+    # remove columns except prompt_1 prompt_2 goal_safe goal_unsafe
+    df = df[['inputs', 'goal_safe', 'goal_unsafe']]
+    df['prompt_1'] = df['inputs'].apply(lambda x: x['prompt_1'])
+    df['prompt_2'] = df['inputs'].apply(lambda x: x['prompt_2'])
+    df = df.drop(columns=['inputs'])
+
+    map_funcs = {
+        "gemma-1.1": lambda x: x['prompt_1'] + x['goal_safe'] + "<end_of_turn>",
+        "Starling-LM": lambda x: x['prompt_1'] + x['goal_safe'] + "<|end_of_turn|>",
+        "Llama-3": lambda x: x['prompt_1'] + x['goal_safe'] + "<|eot_id|>",
+        "Llama-2": lambda x: x['prompt_1'] + x['goal_safe'] + "</s>",
+        "zephyr": lambda x: x['prompt_1'] + x['goal_safe'] + "</s>",
+        "Phi-3": lambda x: x['prompt_1'] + x['goal_safe'] + "<|end|>",
+    }
+
+    df['text'] = df.apply(map_funcs[model_type], axis=1)
+
+    df_train = df.sample(frac=0.8, random_state=42)
+    df_test = df.drop(df_train.index)
+
+    #print a few random samples
+    print(df_train['text'].sample(2).to_list())
+
+    os.makedirs(out_dir, exist_ok=True)
+    df_train.to_json(f'{out_dir}/train_dataset.json',
+                     orient='records',
+                     force_ascii=False)
+    df_test.to_json(f'{out_dir}/test_dataset.json',
+                    orient='records',
+                    force_ascii=False)
+
+
+if __name__ == '__main__':
+    fire.Fire(main)

Codes/1_raw_dataset/Should-It-Be-Executed-Or-Processed/fine-tuning/train_fsdp.py

@@ -0,0 +1,200 @@
+from dataclasses import dataclass, field
+import os
+import random
+import torch
+from datasets import load_dataset
+from transformers import AutoTokenizer, TrainingArguments
+from trl.commands.cli_utils import TrlParser
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    set_seed,
+
+)
+
+from peft import LoraConfig
+
+from trl import SFTTrainer
+
+
+@dataclass
+class ScriptArguments:
+    """
+    A class to hold script arguments for model training.
+
+    Attributes:
+    dataset_path (str): Path to the dataset.
+    dataset_text_field (str): Dataset text field used for decoder-only training. Default is "text".
+    model_id (str): Model ID to use for SFT training.
+    max_seq_length (int): The maximum sequence length for SFT Trainer. Default is 512.
+    training_mode (str): Training mode: lora, qlora, or fft. Default is "lora".
+    attention_impl (str): Attention implementation: sdpa or flash_attention_2. Default is "sdpa".
+    lora_r (int): LoRA r parameter. Default is 16.
+    lora_alpha (int): LoRA alpha parameter. Default is 8.
+    lora_dropout (float): LoRA dropout parameter. Default is 0.05.
+    peft_target_modules (str): PEFT target modules. Default is "all-linear".
+    """
+    dataset_path: str = field(
+        default=None,
+        metadata={
+            "help": "Path to the dataset"
+        },
+    )
+    dataset_text_field: str = field(
+        default="text", metadata={"help": "Dataset text field used for decoder_only training"}
+    )
+    model_id: str = field(
+        default=None, metadata={"help": "Model ID to use for SFT training"}
+    )
+    max_seq_length: int = field(
+        default=512, metadata={"help": "The maximum sequence length for SFT Trainer"}
+    )
+    training_mode: str = field(
+        default="lora", metadata={"help": "Training mode: lora, qlora or fft"}
+    )
+    attention_impl: str = field(
+        default="sdpa", metadata={"help": "Attention implementation: sdpa or flash_attention_2"}
+    )
+    lora_r: int = field(
+        default=16, metadata={"help": "LoRA r parameter"}
+    )
+    lora_alpha: int = field(
+        default=8, metadata={"help": "LoRA alpha parameter"}
+    )
+    lora_dropout: float = field(
+        default=0.05, metadata={"help": "LoRA dropout parameter"}
+    )
+    peft_target_modules: str = field(
+        default="all-linear", metadata={"help": "PEFT target modules"}
+    )
+
+
+def training_function(script_args: ScriptArguments, training_args: TrainingArguments) -> None:
+    """
+    Train a model using the specified script arguments and training arguments.
+
+    Parameters:
+    script_args (ScriptArguments): The script arguments for model training.
+    training_args (TrainingArguments): The training arguments for the Trainer.
+
+    The function performs the following steps:
+    1. Load the training and testing datasets from JSON files.
+    2. Initialize the tokenizer using the specified model ID.
+    3. Print a few random samples from the training set.
+    4. Initialize the model with or without quantization based on the training mode.
+    5. Configure PEFT settings if using LoRA or QLoRA training mode.
+    6. Train the model using SFTTrainer and save the trained model.
+    """
+
+    train_dataset = load_dataset(
+        "json",
+        data_files=os.path.join(script_args.dataset_path, "train_dataset.json"),
+        split="train",
+    )
+    test_dataset = load_dataset(
+        "json",
+        data_files=os.path.join(script_args.dataset_path, "test_dataset.json"),
+        split="train",
+    )
+
+    ################
+    # Model & Tokenizer
+    ################
+
+    # Tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(script_args.model_id, use_fast=True)
+    tokenizer.pad_token = tokenizer.eos_token
+
+
+    # print random sample
+    with training_args.main_process_first(
+            desc="Log a few random samples from the processed training set"
+    ):
+        for index in random.sample(range(len(train_dataset)), 2):
+            print(train_dataset[index][script_args.dataset_text_field])
+
+    # Model
+    torch_dtype = torch.bfloat16
+    quant_storage_dtype = torch.bfloat16
+
+    if script_args.training_mode == "qlora":
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch_dtype,
+            bnb_4bit_quant_storage=quant_storage_dtype,
+        )
+    else:
+        quantization_config = None
+
+    model = AutoModelForCausalLM.from_pretrained(
+        script_args.model_id,
+        quantization_config=quantization_config,
+        attn_implementation=script_args.attention_impl,
+        torch_dtype=quant_storage_dtype,
+        use_cache=False if training_args.gradient_checkpointing else True,
+        trust_remote_code=True if 'microsoft/Phi-3' in script_args.model_id else False,
+    )
+
+    if training_args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+
+    ################
+    # PEFT
+    ################
+
+    if script_args.training_mode in ["lora", "qlora"]:
+        peft_config = LoraConfig(
+            lora_alpha=script_args.lora_alpha,
+            lora_dropout=script_args.lora_dropout,
+            r=script_args.lora_r,
+            bias="none",
+            target_modules=script_args.peft_target_modules,
+            task_type="CAUSAL_LM",
+            modules_to_save=["lm_head", "embed_tokens"]
+        )
+    else:
+        peft_config = None
+
+    trainer = SFTTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        dataset_text_field=script_args.dataset_text_field,
+        eval_dataset=test_dataset,
+        peft_config=peft_config,
+        max_seq_length=script_args.max_seq_length,
+        tokenizer=tokenizer,
+        packing=True,
+        dataset_kwargs={
+            "add_special_tokens": False,
+            "append_concat_token": False,
+        },
+    )
+    if trainer.accelerator.is_main_process and hasattr(trainer.model, "print_trainable_parameters"):
+        trainer.model.print_trainable_parameters()
+
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    trainer.train(resume_from_checkpoint=checkpoint)
+
+    if trainer.is_fsdp_enabled:
+        trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
+    trainer.save_model()
+
+
+if __name__ == "__main__":
+    parser = TrlParser((ScriptArguments, TrainingArguments))
+    script_args, training_args = parser.parse_args_and_config()
+
+    # set use reentrant to False
+    if training_args.gradient_checkpointing:
+        training_args.gradient_checkpointing_kwargs = {"use_reentrant": True}
+    # set seed
+    set_seed(training_args.seed)
+
+    # launch training
+    training_function(script_args, training_args)