train.py

#!/usr/bin/env python3

import math
import copy
import torch
import os
import random
import time
import shutil
from torch.utils.data import SequentialSampler, Subset, RandomSampler
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, \
    PreTrainedTokenizerFast, HfArgumentParser, GPTQConfig, AutoConfig, TrainerCallback, BitsAndBytesConfig
from transformers.trainer_utils import EvalPrediction
from transformers.integrations.integration_utils import TensorBoardCallback
from datasets import load_dataset, Dataset
from dataclasses import dataclass, field
from typing import Dict, Optional, Sequence, Sized, Iterator


MULTI_GPU_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1"))
MULTI_GPU_RANK = int(os.environ.get("RANK", "0"))
IS_MULTI_GPU = os.environ.get("RANK") != None
IS_MASTER_PROCESS = MULTI_GPU_RANK == 0

@dataclass
class TrainingRunArguments:
    run_name: str = field(metadata={"help": "The folder to save the output model under"})
    base_model: str = field(metadata={"help": "The base model to load for fine-tuning"})
    train_dataset: str = field(metadata={"help": "The JSON file containing the training dataset"})
    test_dataset: str = field(default=None, metadata={"help": "The JSON file containing the evaluation dataset"})
    ctx_size: int = field(default=2048, metadata={"help": "The number of tokens to pad & truncate the input examples to"})
    bf16: bool = field(default=False, metadata={"help": "If set, the model will the loaded and trained in bf16 instead of fp16"})
    batch_size: int = field(default=8, metadata={"help": "The simulated 'batch size' that we will train on. will tweak gradient accumulations steps"})
    micro_batch_size: int = field(default=2, metadata={"help": "The actual batch size that will fit into VRAM on this machine"})
    epochs: int = field(default=1, metadata={"help": "The number of times to train the model on each example"})
    learning_rate: float = field(default=1e-5, metadata={"help": "The starting learning rate (speed at which the model trains)"})
    learning_rate_schedule: str = field(default="cosine", metadata={"help": "How fast the learning rate is reduced during training"})
    learning_rate_warmup: float = field(default=0.0, metadata={"help": "The starting learning rate (speed at which the model trains)"})
    weight_decay: float = field(default=0.1, metadata={"help": ""})
    gradient_clip: float = field(default=1.0, metadata={"help": ""})
    resume_from_checkpoint: str = field(default="", metadata={"help": "The name of the checkpoint to resume training from"})
    eval_steps: int = field(default=200, metadata={"help": "The number of steps in between evaluations of the model; set to -1 to evaluate every epoch"})
    save_steps: int = field(default=-1, metadata={"help": "The number of steps in between model checkpoints; set to -1 to save every epoch"})
    save_total_limit: int = field(default=1, metadata={"help": "The number of recent checkpoints of the model to save (not including the final model)"})
    logging_steps: int = field(default=5, metadata={"help": "Sets the number of steps in between log output for the training run"})
    group_by_length: bool = field(default=False, metadata={"help": "If enabled, the training data will be grouped by length to optimize use of padding"})
    pre_allocate_cuda_buffers: bool = field(default=True, metadata={"help": "If enabled, runs a forward and backward pass on the model before training to force pytorch to allocate the correct size CUDA buffers up front"})
    
    # Quantization
    load_in_8bit: bool = field(default=False, metadata={"help": "Set to load the base model in 8-bit mode using bitsandbytes"})
    load_in_4bit: bool = field(default=False, metadata={"help": "Set to load the base model in 4-bit mode using bitsandbytes"})
    load_as_gptq: bool = field(default=False, metadata={"help": "Set to load the base model as a GPTQ using AutoGPTQ"})
    
    # lora config
    use_lora: bool = field(default=False, metadata={"help": "If set, then the trained model will be a LoRA"})
    lora_rank: int = field(default=4, metadata={"help": "Rank which determines LoRA matrix size. Rank typically starts at 8 but can go up to 256. Higher ranks can store more information but increase the computational and memory cost of LoRA."})
    lora_alpha: int = field(default=32, metadata={"help": "Alpha a scaling factor for updates. Alpha directly impacts the adapters contribution and is often set to 1x or 2x the rank value."})
    lora_dropout: float = field(default=0.05)
    lora_modules: str = field(default=None, metadata={"help": "Target modules: LoRA can be applied to various model components, including attention mechanisms (Q, K, V matrices), output projections, feed-forward blocks, and linear output layers. While initially focused on attention mechanisms, extending LoRA to other components has shown benefits. However, adapting more modules increases the number of trainable parameters and memory needs."})
    lora_modules_to_save: str = field(default=None, metadata={"help": "Additional modules to save"})
    lora_merge: bool = field(default=False, metadata={"help": "If set, the Lora will be merged back into the base model an saved"})

    dpo: bool = field(default=False, metadata={"help": "If set, performs Direct Preference Optimization instead of Supervised Fine Tuning"})
    beta: float = field(default=0.1, metadata={"help": "The implicit reward value used during DPO training"})
    dpo_loss: str = field(default="sigmoid", metadata={"help": "The loss type to use during DPO training"})

    add_pad_token: bool = field(default=False, metadata={"help": "If set, a pad token will be added to the tokenizer's vocabulary"})
    add_chatml_tokens: bool = field(default=False, metadata={"help": "If set, tokens for the ChatML format will be added specifically"})
    add_chatml_prompt_template: bool = field(default=False, metadata={"help": "If set, the ChatML prompt template will be set as the model's Jinja2 template"})
    gradient_checkpointing: bool = field(default=False, metadata={"help": "Enables gradient checkpointing which saves quite a lot of VRAM"})

    sync_to_bucket: str = field(default=None, metadata={"help": "If set, checkpoints will be synced to the s3 bucket specified by this argument"})
    flops_baseline: str = field(default=None, metadata={"help": "The baseline flops for the GPUs used for the training run. Outputs MFU"})

    prefix_ids:str = field(default=None, metadata={"help": "Determine the prefix tokens that surround the response from the assistant for SFT if model can not correctly recognise response."})
    suffix_ids:str = field(default=None, metadata={"help": "Determine the suffix tokens that surround the response from the assistant for SFT if model can not correctly recognise response."})


class UploadToS3Callback(TrainerCallback):
    def __init__(self, s3_bucket, s3_prefix, save_total_limit=None):
        import boto3
        self.s3_client = boto3.client('s3')
        self.s3_bucket = s3_bucket
        self.s3_prefix = s3_prefix
        self.save_total_limit = save_total_limit

    def on_save(self, args, state, control, **kwargs):
        output_dir = kwargs['output_dir']
        checkpoint = os.path.basename(output_dir)
        
        # Upload current checkpoint
        for root, dirs, files in os.walk(output_dir):
            for file in files:
                local_path = os.path.join(root, file)
                s3_path = os.path.join(self.s3_prefix, checkpoint, os.path.relpath(local_path, start=output_dir))
                self.s3_client.upload_file(local_path, self.s3_bucket, s3_path)
                print(f"Uploaded {local_path} to s3://{self.s3_bucket}/{s3_path}")

        # Manage checkpoints in S3
        if self.save_total_limit:
            s3_checkpoints = self.list_s3_checkpoints()
            if len(s3_checkpoints) > self.save_total_limit:
                sorted_checkpoints = sorted(s3_checkpoints)
                to_delete = sorted_checkpoints[:-self.save_total_limit]
                for checkpoint in to_delete:
                    self.delete_checkpoint_from_s3(checkpoint)

        # Clean local checkpoints, keeping only the most recent
        all_checkpoints = [os.path.join(args.output_dir, d) for d in os.listdir(args.output_dir) if os.path.isdir(os.path.join(args.output_dir, d))]
        if all_checkpoints:
            latest_checkpoint = max(all_checkpoints, key=os.path.getmtime)
            for checkpoint_dir in all_checkpoints:
                if checkpoint_dir != latest_checkpoint:
                    shutil.rmtree(checkpoint_dir)
                    print(f"Deleted local checkpoint {checkpoint_dir}")

    def list_s3_checkpoints(self):
        paginator = self.s3_client.get_paginator('list_objects_v2')
        page_iterator = paginator.paginate(Bucket=self.s3_bucket, Prefix=self.s3_prefix, Delimiter='/')
        return [prefix.get('Prefix').rstrip('/').split('/')[-1] for page in page_iterator for prefix in page.get('CommonPrefixes', [])]

    def delete_checkpoint_from_s3(self, checkpoint_name):
        resp = self.s3_client.list_objects_v2(Bucket=self.s3_bucket, Prefix=os.path.join(self.s3_prefix, checkpoint_name))
        for obj in resp.get('Contents', []):
            self.s3_client.delete_object(Bucket=self.s3_bucket, Key=obj['Key'])
            print(f"Deleted s3://{self.s3_bucket}/{obj['Key']}")

class MFUCallback(TrainerCallback):
    def __init__(self, peak_flops):
        self.total_iterations = 0
        self.start_time = time.time()
        self.flops_promised = peak_flops
        self.last_total_flos = 0

    def on_log(self, args, state, control, **kwargs):
        if state.global_step == 0:  # Avoid computation at the very beginning
            return
        
        current_time = time.time()
        elapsed_time = current_time - self.start_time

        # Calculate and log MFU
        new_flops = state.total_flos - self.last_total_flos
        kwargs['logs']['mfu'] = round(new_flops / elapsed_time / self.flops_promised, 4)

        self.start_time = current_time
        self.last_total_flos = state.total_flos
   

parser = HfArgumentParser([TrainingRunArguments])
training_run_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)

if sum([training_run_args.load_in_8bit, training_run_args.load_in_4bit, training_run_args.load_as_gptq]) > 1:
    raise Exception("Please select exactly one of 'load_in_8bit', 'load_in_4bit', or 'load_as_gptq")

if IS_MASTER_PROCESS:
    print(f"Loading model '{training_run_args.base_model}'...")

model_kwargs = {}
if training_run_args.load_in_8bit:
    model_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
elif training_run_args.load_in_4bit:
    model_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
elif training_run_args.load_as_gptq:
    model_kwargs["quantization_config"] = GPTQConfig(bits=4, disable_exllama=True)


if training_run_args.bf16:
    model_kwargs["torch_dtype"] = torch.bfloat16
else:
    model_kwargs["torch_dtype"] = torch.float16

# model_kwargs["resid_pdrop"] = 0.0
# model_kwargs["revision"] = "accfee56d8988cae60915486310362db5831b1bd"
model_kwargs["use_cache"] = False

def find_max_vram(min_buffer_mib=800):
    max_memory = {}
    for i in range(torch.cuda.device_count()):
        total_mem = (torch.cuda.get_device_properties(i).total_memory / (1024 * 1024))
        suggestion = round((total_mem - 1000) / 1000) * 1000
        suggestion = min(suggestion, total_mem - min_buffer_mib)

        if IS_MASTER_PROCESS:
            print(f"Model will target using {suggestion}MiB of VRAM on GPU {i}")
        max_memory[i] = f'{suggestion}MiB'

    return max_memory

if "LOCAL_RANK" not in os.environ:
    model_kwargs["device_map"] = "auto"

model = AutoModelForCausalLM.from_pretrained(
    training_run_args.base_model,
    trust_remote_code=True,
    max_memory=find_max_vram(),
    **model_kwargs
)
tokenizer = AutoTokenizer.from_pretrained(training_run_args.base_model, trust_remote_code=True)

if training_run_args.add_pad_token:
    tokenizer.add_special_tokens({'pad_token': '<|pad|>'})
    model.config.pad_token_id = tokenizer.pad_token_id

if training_run_args.add_chatml_tokens:
    tokenizer.add_special_tokens({
        'bos_token': '<|im_start|>',
        'eos_token': '<|im_end|>'
    })

    model.config.bos_token_id = tokenizer.bos_token_id
    model.config.eos_token_id = tokenizer.eos_token_id

if training_run_args.add_chatml_prompt_template:
    tokenizer.chat_template = (
        "{% for message in messages %}"
        "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
        "{% endfor %}"
        "{% if add_generation_prompt %}"
        "{{ '<|im_start|>assistant\n' }}"
        "{% endif %}"
    )

embeddings_len = math.ceil(len(tokenizer) / 32) * 32
if model.get_input_embeddings().num_embeddings < embeddings_len:
    model.resize_token_embeddings(embeddings_len)
else:
    model.tie_weights()

# model.tie_weights()

original_model = model
peft_config = None
if training_run_args.use_lora:
    from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
    if IS_MASTER_PROCESS:
        print("Creating LoRA for model...")
    target_modules = training_run_args.lora_modules.split(",") if training_run_args.lora_modules else None
    modules_to_save = training_run_args.lora_modules_to_save.split(",") if training_run_args.lora_modules_to_save else None
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        inference_mode=False,
        r=training_run_args.lora_rank,
        lora_alpha=training_run_args.lora_alpha,
        lora_dropout=training_run_args.lora_dropout,
        target_modules=target_modules,
        modules_to_save=modules_to_save,
    )
    if training_run_args.load_in_8bit or training_run_args.load_in_4bit or training_run_args.load_as_gptq:
        model = prepare_model_for_kbit_training(
            model, use_gradient_checkpointing=training_run_args.gradient_checkpointing
        )
    model = get_peft_model(model, peft_config)
    model.enable_input_require_grads()

    model.print_trainable_parameters()
    

base_dir = "loras" if training_run_args.use_lora else "models"
model_dir = f"./{base_dir}/{training_run_args.run_name}"

training_kwargs = {}

if training_run_args.test_dataset:
    training_kwargs.update({
        "per_device_eval_batch_size": training_run_args.micro_batch_size,
        "eval_strategy": ("steps" if training_run_args.eval_steps != -1 else "epoch"),
        "eval_steps": (training_run_args.eval_steps if training_run_args.eval_steps != -1 else None),
        "bf16_full_eval": training_run_args.bf16,
    })

training_args = TrainingArguments(
    per_device_train_batch_size=training_run_args.micro_batch_size,
    gradient_accumulation_steps=training_run_args.batch_size//training_run_args.micro_batch_size,
    gradient_checkpointing=training_run_args.gradient_checkpointing,
    weight_decay=training_run_args.weight_decay,
    max_grad_norm=training_run_args.gradient_clip,
    save_strategy=("steps" if training_run_args.save_steps != -1 else "epoch"),
    save_steps=(training_run_args.save_steps if training_run_args.save_steps != -1 else None),
    save_safetensors=True,
    logging_steps=training_run_args.logging_steps, 
    output_dir=model_dir,
    num_train_epochs=training_run_args.epochs,
    save_total_limit=training_run_args.save_total_limit,
    report_to='none',
    learning_rate=training_run_args.learning_rate,
    lr_scheduler_type=training_run_args.learning_rate_schedule,
    warmup_ratio=training_run_args.learning_rate_warmup,
    log_level="info",
    bf16=training_run_args.bf16,
    group_by_length=training_run_args.group_by_length,
    # include_num_input_tokens_seen=True,
    **training_kwargs,
)

class DataCollatorForSupervisedFineTuning(object):
    """Collate examples for supervised fine-tuning."""

    tokenizer: AutoTokenizer
    prompt_split: str
    response_prefix: str
    response_suffix: str
    prefix_ids: list[int]
    suffix_ids: list[int]

    def __init__(self, *, tokenizer: AutoTokenizer, prefix_ids: Optional[list[int]] = None, suffix_ids: Optional[list[int]] = None):
        
        self.tokenizer = tokenizer
        if not prefix_ids and not suffix_ids:
            assistant_prompt = tokenizer.apply_chat_template(
                conversation=[{"role": "assistant", "content":  r"%%%%%%%%%%%%%%%%"}], 
                tokenize=False).split( r"%%%%%%%%%%%%%%%%")
            
            self.response_prefix = assistant_prompt[0]
            self.response_suffix = assistant_prompt[1]

            # check for inserted system prompt and remove it
            if tokenizer.eos_token in self.response_prefix:
                self.response_prefix = self.response_prefix.split(tokenizer.eos_token)[1].lstrip()

            # some chat templates ALWAYS add the bos token
            if tokenizer.bos_token in self.response_prefix:
                self.response_prefix = self.response_prefix.replace(tokenizer.bos_token, "")

        if prefix_ids:
            self.prefix_ids = prefix_ids
        else:
            self.prefix_ids = self.tokenizer(self.response_prefix, add_special_tokens=False)["input_ids"]

        if suffix_ids:
            self.suffix_ids = suffix_ids
        else:
            self.suffix_ids = self.tokenizer(self.response_suffix, add_special_tokens=False)["input_ids"]

    def _find_mask_ranges(self, input_ids):
        """
        Returns a mask that blocks out everything but the response from the assistant
        The mask does NOT include the response_prefix but DOES include the response_suffix.
        The resulting behavior is the model uses the prefix as a prompt and the suffix as the end of text token
        """
        ranges = []
        i = 0

        while i < len(input_ids):
            try:
                # Find the start index of the prefix
                start_idx = input_ids.index(self.prefix_ids[0], i)
            except ValueError:
                break

            # Check if the entire prefix is present
            if input_ids[start_idx:start_idx + len(self.prefix_ids)] == self.prefix_ids:
                end_prefix_idx = start_idx + len(self.prefix_ids)
                start_response_idx = end_prefix_idx + 1

                # Find the start index of the suffix
                try:
                    # Find the start index of the suffix
                    suffix_start_idx = input_ids.index(self.suffix_ids[0], end_prefix_idx)
                except ValueError:
                    ranges.append((start_response_idx, len(input_ids)))
                    break

                # Check if the entire suffix is present
                if input_ids[suffix_start_idx:suffix_start_idx + len(self.suffix_ids)] == self.suffix_ids:
                    ranges.append((start_response_idx, suffix_start_idx))
                    i = suffix_start_idx + len(self.suffix_ids)
                else:
                    i = suffix_start_idx + 1
            else:
                i = start_idx + 1

        inverse_ranges = []
        current = 0

        for start, end in sorted(ranges):
            if start > current:
                inverse_ranges.append((current, start - 1))
            current = max(current, end + 1)
        
        if current < len(input_ids):
            inverse_ranges.append((current, len(input_ids) - 1))

        return inverse_ranges
    
    def _pad(self, examples, pad_value):
        longest = max([len(ex) for ex in examples])
        result = []
        for example in examples:
            cur_len = len(example)
            result.append(example + [pad_value] * (longest - cur_len))

        return result

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids = [instance["input_ids"] for instance in instances]
        labels = copy.deepcopy(input_ids)

        for label in labels:
            mask_ranges = self._find_mask_ranges(label)
            for start, end in mask_ranges:
                if end - start == len(label) - 1:
                    print("warning! example had no assistant response in it!")
                    print(input_ids)
                label[start:end] = [-100] * (end - start)

        input_ids = torch.LongTensor(self._pad(input_ids, self.tokenizer.pad_token_id or self.tokenizer.eos_token_id))
        labels = torch.LongTensor(self._pad(labels, -100))

        return dict(
            input_ids=input_ids,
            labels=labels,
            attention_mask=input_ids.ne(self.tokenizer.pad_token_id or self.tokenizer.eos_token_id),
        )

if IS_MASTER_PROCESS:
    print("Loading dataset...")
data_files = { "train": training_run_args.train_dataset }
if training_run_args.test_dataset:
    data_files["test"] = training_run_args.test_dataset
datasets = load_dataset("json", data_files=data_files)

def tokenize_raw_example(batch):
    return tokenizer(
        text=batch["text"],
        max_length=training_run_args.ctx_size,
        truncation=True,
        add_special_tokens=False,
    )

def tokenize_sharegpt_example(batch):
    # TODO: figure out how to properly batch this
    result = []
    for example in batch["conversations"]:
        conversation = [ { "role": x["from"], "content": x["value"] }  for x in example ]
        result.append(
            tokenizer.apply_chat_template(
                conversation=conversation,
                max_length=training_run_args.ctx_size,
                truncation=True,
            )
        )

    return {"input_ids": result}

def template_dpo_example(batch):
    # TODO: figure out how to properly batch this
    result = []
    for example in zip(batch["system"], batch["question"]):
        conversation = [ 
            { "role": "system", "content": example[0] },
            { "role": "user", "content": example[1] },
        ]
        result.append(
            tokenizer.apply_chat_template(
                conversation=conversation,
                max_length=training_run_args.ctx_size,
                truncation=True,
                tokenize=False,
                add_generation_prompt=True
            )
        )

    return {"prompt": result}

training_callbacks = []
if training_run_args.sync_to_bucket:
    training_callbacks.append(UploadToS3Callback(
        s3_bucket=training_run_args.sync_to_bucket,
        s3_prefix=training_run_args.run_name,
        save_total_limit=training_run_args.save_total_limit
    ))

if training_run_args.flops_baseline:
    # A100 GPU bfloat16 peak flops is 312 TFLOPS (312e12)
    # 4090 GPU bfloat16 peak flops is 165.2 TFLOPS (1652e11)
    # 3090 GPU bfloat16 peak flops is 71 TFLOPS (71e12)

    training_callbacks.append(MFUCallback(peak_flops=float(training_run_args.flops_baseline)))


# log to tensorboard (but after MFU)
training_callbacks.append(TensorBoardCallback())

class CustomSFTTrainer(Trainer):
    """Implement different training tweaks"""
    def __init__(self, random_eval_sample_pct=0.1, learning_rate_overshoot=1.15, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.random_eval_sample_pct = random_eval_sample_pct
        self.evaluate_full_dataset = False
        self.learning_rate_overshoot = learning_rate_overshoot

    def evaluate_all(self):
        self.evaluate_full_dataset = True
        super().evaluate()
        self.evaluate_full_dataset = False

    # Randomly sample the eval dataset
    def _get_eval_sampler(self, eval_dataset):
        if self.evaluate_full_dataset:
            return SequentialSampler(eval_dataset)
        else:
            num_samples = int(self.random_eval_sample_pct * len(eval_dataset))
            random_indices = random.sample(range(len(eval_dataset)), num_samples)
            subset_eval_dataset = Subset(eval_dataset, random_indices)
            return SequentialSampler(subset_eval_dataset)
        
    def _get_train_sampler(self):
        if self.args.group_by_length:
            return super()._get_train_sampler()
        
        return RandomSampler(self.train_dataset, generator=torch.Generator(device='cpu'))
    
    def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optimizer = None):
        """
        Saw this in the chinchilla paper. It says not to go over 25% overshoot
        Should speed up training by skipping the final fine tuning part that doesn't affect accuracy much
        """
        return super().create_scheduler(int(num_training_steps * self.learning_rate_overshoot), optimizer=optimizer)
    
    def floating_point_ops(self, inputs):
        config = self.model.config
        examples_length = len(inputs["input_ids"][0])
        batch_size = len(inputs["input_ids"])

        # mfu is approximated using thoughtput and param count
        # the number of paramters is approximately the number of multiply-accumulates (MAC) in the network
        # each MAC has 2 FLOPs - we multiply by 2 ie 2 * n_param
        # there are 3 passes of a NN (fwd, bwd, delta) - we multiply by 3 ie 2 * 3 * n_param
        # this gets us FLOPs / token
        flops_per_token = 2 * sum(p.numel() for p in self.model.parameters())
        flops_per_seq = flops_per_token * examples_length

        # there are 2 FLOPS per mac; there is A=Q*K^T and out=A*V ops (ie mult by 2)
        attn_flops_per_seq = config.num_hidden_layers * 2 * 2 * (config.hidden_size * (examples_length**2))

        # there are 2 ops in bwd pass and 1 in fwd pass so we mult by 3
        result = (3 * flops_per_seq + 3 * attn_flops_per_seq) * batch_size
        return result

if not training_run_args.dpo:
    if IS_MASTER_PROCESS:
        print("Tokenizing datasets...")

    if "text" in datasets["train"].column_names:
        tokenize_function = tokenize_raw_example
        columns_to_remove = ["text"]
    elif "conversations" in datasets["train"].column_names:
        tokenize_function = tokenize_sharegpt_example
        columns_to_remove = ["conversations"]
    else:
        raise Exception("Unknown dataset input format (not raw corpus or sharegpt)")

    tokenized_test_dataset = None
    num_proc = os.cpu_count() // MULTI_GPU_WORLD_SIZE
    tokenized_train_dataset = datasets["train"].map(tokenize_function, batched=True, num_proc=num_proc).remove_columns(columns_to_remove)
    if training_run_args.test_dataset:
        tokenized_test_dataset = datasets["test"].map(tokenize_function, batched=True, num_proc=num_proc).remove_columns(columns_to_remove)

    example_lengths = [ len(example) for example in tokenized_train_dataset["input_ids"] ]
    tokens_in_train_set, longest_example = sum(example_lengths), max(example_lengths)
    if IS_MASTER_PROCESS:
        print(f"Train dataset has {int(tokens_in_train_set / 1000000)}M tokens. Longest Example: {longest_example} tokens")
    
    provided_prefix_ids = None
    provided_suffix_ids = None
    try:
        if training_run_args.prefix_ids:
            provided_prefix_ids = [ int(x) for x in training_run_args.prefix_ids.split(",") ]
        if training_run_args.suffix_ids:
            provided_suffix_ids = [ int(x) for x in training_run_args.suffix_ids.split(",") ]
    except ValueError as ex:
        print(f"Error parsing prefix_ids or suffix_ids: '{ex}'")
        exit(-1)
    
    data_collator = DataCollatorForSupervisedFineTuning(
        tokenizer=tokenizer,
        prefix_ids=provided_prefix_ids,
        suffix_ids=provided_suffix_ids,
    )

    trainer = CustomSFTTrainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_train_dataset,
        eval_dataset=tokenized_test_dataset,
        data_collator=data_collator,
        callbacks=training_callbacks,
    )
else:
    from trl import DPOTrainer
    max_prompt_length = 0

    train_dataset = datasets["train"].map(lambda x: { "prompt_len": len(x["system"]) })

    test_dataset = None
    if training_run_args.test_dataset:
        test_dataset = datasets["test"]

    max_prompt_length = max(train_dataset["prompt_len"])

    print("Templating DPO Examples...")
    templated_test_dataset = None
    templated_train_dataset = train_dataset.map(template_dpo_example, batched=True).remove_columns(["system", "question"])
    if training_run_args.test_dataset:
        templated_test_dataset = datasets["test"].map(template_dpo_example, batched=True).remove_columns(["system", "question"])

    # tokenizer.model_input_names = [ "chosen_input_ids" ]

    # group_by_length doesn't work here
    # templated_train_dataset = templated_train_dataset.sort("prompt_len", reverse=True)

    training_args.length_column_name = "prompt_len"
    model.enable_input_require_grads()

    trainer = DPOTrainer(
        model,
        ref_model=None,
        # ref_model=original_model,
        peft_config=peft_config,
        args=training_args,
        beta=training_run_args.beta,
        loss_type=training_run_args.dpo_loss,
        train_dataset=templated_train_dataset,
        eval_dataset=templated_test_dataset,
        tokenizer=tokenizer,
        max_length=training_run_args.ctx_size,
        max_prompt_length=max_prompt_length,
        truncation_mode="keep_start",
        callbacks=training_callbacks,
    )

try:
    checkpoint = training_run_args.resume_from_checkpoint
    if checkpoint:
        trainer.train(checkpoint)
    else:
        trainer.train()

    if training_run_args.test_dataset:
        trainer.evaluate_all()

    if trainer.is_fsdp_enabled:
        trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")

    if training_run_args.use_lora and training_run_args.lora_merge:
        trainer.save_model() # save lora

        merged_model = model.merge_and_unload(progressbar=True)
        merged_model_dir = f"./models/{training_run_args.run_name}"
        merged_model.save_pretrained(merged_model_dir, safe_serialization=True, max_shard_size="2GB")
        
        tokenizer.save_pretrained(merged_model_dir)
    else:
        trainer.save_model()
        tokenizer.save_pretrained(model_dir)

except Exception as ex:
    if trainer.is_fsdp_enabled:
        raise ex # this doesn't play nice with FSDP so don't even try
    
    print("Something bad happened! Try and save it?")
    import code, traceback
    traceback.print_exc()
    code.interact(local=locals())