# Werfault.exe errors appear during model fine-tuning

**URL:** <https://discuss.huggingface.co/t/werfault-exe-errors-appear-during-model-fine-tuning/94284>\
**Category:** Beginners\
**Created:** [June 26, 2024, 8:09pm UTC](https://discuss.huggingface.co/t/werfault-exe-errors-appear-during-model-fine-tuning/94284 "2024-06-26T20:09:58Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![gacekk](https://avatars.discourse-cdn.com/v4/letter/g/7993a0/32.png) [@gacekk](https://discuss.huggingface.co/u/gacekk)\
**Post date:** [June 26, 2024, 8:09pm UTC](https://discuss.huggingface.co/t/werfault-exe-errors-appear-during-model-fine-tuning/94284/1 "2024-06-26T20:09:58Z")

</div>

Hi

I am trying to fine tune a model using custom dataset, the script is runnig, but I keep getting windows errors. I am running below script and using these setting. Runnig this in windows 11, RTX3090, 32GB of RAM

> python train.py --run\_name home1bPolish --base model speakleash/Bielik-7B-v0.1 --bf16 --train\_dataset data/home\_assistant\_train.jsonl --test\_dataset data/home\_assistant\_test.jsonl --learning\_rate 2e-5 --batch\_size 32 --micro\_batch\_size 32 --gradient\_checkpointing --group\_by\_length --ctx\_size 2048 --save\_steps 100 --save\_total\_limit 10

> #!/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, AutoConfig, TrainerCallback  
from transformers.trainer\_utils import EvalPrediction

from datasets import load\_dataset, Dataset  
from dataclasses import dataclass, field  
from typing import Dict, Optional, Sequence, Sized, Iterator  
from huggingface\_hub import login

login(token=“hf\_DtCOvpcvhUqoaYgDCfaJIChWNKGKgKuTEb”)  
os.environ[‘TF\_ENABLE\_ONEDNN\_OPTS’] = ‘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)
lora_alpha: int = field(default=32)
lora_dropout: float = field(default=0.05)
lora_modules: str = field(default=None)
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"})

```

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”)

print(f"Loading model ‘{training\_run\_args.base\_model}’…")

model\_kwargs = {}  
if training\_run\_args.load\_in\_8bit:  
model\_kwargs[“load\_in\_8bit”] = True  
elif training\_run\_args.load\_in\_4bit:  
model\_kwargs[“load\_in\_4bit”] = 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)

```
    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(“speakleash/Bielik-7B-v0.1”, trust\_remote\_code=True,  
max\_memory=find\_max\_vram)

tokenizer = AutoTokenizer.from\_pretrained(“speakleash/Bielik-7B-v0.1”, 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  
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"./{training\_run\_args.run\_name}"

training\_kwargs = {}

if training\_run\_args.test\_dataset:  
training\_kwargs.update({  
“per\_device\_eval\_batch\_size”: training\_run\_args.batch\_size,  
“evaluation\_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 = None, suffix_ids = None):
    
    self.tokenizer = tokenizer
    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]

    if prefix_ids:
        self.prefix_ids = prefix_ids
    else:
        self.prefix_ids = self.tokenizer["input_ids"]

    if suffix_ids:
        self.suffix_ids = suffix_ids
    else:
        self.suffix_ids = self.tokenizer["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):
                print("warning! example had no assistant response in it!")
            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),
    )

```

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)))

```

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:  
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
tokenized_train_dataset = datasets["train"].map(tokenize_function, batched=True, num_proc=os.cpu_count()).remove_columns(columns_to_remove)
if training_run_args.test_dataset:
    tokenized_test_dataset = datasets["test"].map(tokenize_function, batched=True, num_proc=os.cpu_count()).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)
print(f"Train dataset has {int(tokens_in_train_set / 1000000)}M tokens. Longest Example: {longest_example} tokens")

# data_collator = DataCollatorForSupervisedFineTuning(tokenizer=tokenizer)
# fix for tinyllama not detecting split properly
data_collator = DataCollatorForSupervisedFineTuning(
    tokenizer=tokenizer,
    prefix_ids=[29966, 29989, 465, 22137, 29989, 29958, 13],
    suffix_ids=[2],
)

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:  
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())

```