Werfault.exe errors appear during model fine-tuning
HiI 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 RAMpython 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 python3import mathimport copyimport torchimport osimport randomimport timeimport shutilfrom torch.utils.data import SequentialSampler, Subset, RandomSamplerfrom transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer,PreTrainedTokenizerFast, HfArgumentParser, AutoConfig, TrainerCallbackfrom transformers.trainer_utils import EvalPredictionfrom datasets import load_dataset, Datasetfrom dataclasses import dataclass, fieldfrom typing import Dict, Optional, Sequence, Sized, Iteratorfrom huggingface_hub import loginlogin(token=“REDACTED_TOKEN”)os.environ[‘TF_ENABLE_ONEDNN_OPTS’] = ‘0’@dataclassclass 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):definit(self, s3_bucket, s3_prefix, save_total_limit=None):import boto3self.s3_client = boto3.client(‘s3’)self.s3_bucket = s3_bucketself.s3_prefix = s3_prefixself.save_total_limit = save_total_limitdef 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):definit(self, peak_flops):self.total_iterations = 0self.start_time = time.time()self.flops_promised = peak_flopsself.last_total_flos = 0def 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_flosparser = 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”] = Trueelif training_run_args.load_in_4bit:model_kwargs[“load_in_4bit”] = Trueif training_run_args.bf16:model_kwargs[“torch_dtype”] = torch.bfloat16else:model_kwargs[“torch_dtype”] = torch.float16model_kwargs[“resid_pdrop”] = 0.0model_kwargs[“revision”] = “accfee56d8988cae60915486310362db5831b1bd”model_kwargs[“use_cache”] = Falsedef 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) * 1000suggestion = 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_memoryif “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_idif 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_idif 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) * 32if model.get_input_embeddings().num_embeddings < embeddings_len:model.resize_token_embeddings(embeddings_len)else:model.tie_weights()model.tie_weights()original_model = modelpeft_config = Noneif training_run_args.use_lora:from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_trainingprint(“Creating LoRA for model…”)target_modules = training_run_args.lora_modules.split(“,”) if training_run_args.lora_modules else Nonemodules_to_save = training_run_args.lora_modules_to_save.split(“,”) if training_run_args.lora_modules_to_save else Nonepeft_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_datasetdatasets = 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 thisresult =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 thisresult =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”“”definit(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_pctself.evaluate_full_dataset = Falseself.learning_rate_overshoot = learning_rate_overshootdef 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 resultif 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 DPOTrainermax_prompt_length = 0train_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 tryprint("Something bad happened! Try and save it?")
import code, traceback
traceback.print_exc()
code.interact(local=locals())