transformers - 💡(How to fix) Fix MPS OOM error, finetuning T5Gemma2 with Seq2SeqTrainer [1 participants]

# %%
from huggingface_hub import login

login()

# %%
import os
os.environ["PYTORCH_MPS_LOW_WATERMARK_RATIO"] = "0.9"
os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "1.4" # OOM still happens with default value

# %%
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer
import evaluate
import torch


# %%
dataset = load_dataset("aaai25withanonymous/MathBridge")
dataset = dataset["train"].train_test_split(test_size=0.01, seed=488373954176)

# %%
MODEL_ID = "google/t5gemma-2-270m-270m"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, backend="torchvision")


# %%
model = AutoModelForSeq2SeqLM.from_pretrained(
    MODEL_ID, 
    device_map="auto",
)

# %%
collator = DataCollatorForSeq2Seq(tokenizer, model=model)
metric = evaluate.load("chrf")

# %%
import numpy as np

def postprocess_text(preds, labels):
    preds = [p.lower().strip() for p in preds]
    labels = [[l.lower().strip()] for l in labels]
    return preds, labels

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    if isinstance(preds, tuple):
        preds = preds[0]

    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)

    # Replace -100 padding inserted by the data collator before decoding
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

    result = metric.compute(predictions=decoded_preds, references=decoded_labels)

    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
    result["gen_len"] = np.mean(prediction_lens)
    result = {k: round(v, 4) for k, v in result.items()}
    return {"chrf": result["score"], "gen_len": result["gen_len"]}

# %%
prefix="Convert the LaTeX equation to spoken English: "
def preprocess_function(examples):
    inputs = [
        prefix + ctx_before + " " + eq + " " + ctx_after
        for ctx_before, eq, ctx_after in zip(
            examples["context_before"],
            examples["equation"],
            examples["context_after"]
        )
    ]
    # Dataset has some trailing punctuation noise, which is removed here.
    # There is also sometimes leading capitalisation noise, which is mitigated by
    # a case-insensitive metric at the training stage.
    targets = [s.lower().rstrip(". ") for s in examples["spoken_English"]]

    model_inputs = tokenizer(inputs)

    # Setup the tokenizer for targets
    labels = tokenizer(text_target=targets)

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

preprocess_function(dataset["train"][0:5])

# %%
dataset_tokenized = dataset.map(preprocess_function, batched=True)

# %%
batch_size = 2
model_name = MODEL_ID.split("/")[-1]
args = Seq2SeqTrainingArguments(
    f"MathBridge-{model_name}",
    num_train_epochs=1,
    warmup_steps=0.05,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    gradient_accumulation_steps=64,
    bf16=True,
    eval_strategy="steps",
    eval_steps=500,
    save_strategy="steps",
    save_steps=500,
    save_total_limit=3,
    predict_with_generate=True,
    push_to_hub=True,
    generation_max_length=145,   # tune to your target sequence lengths
)

# %%
trainer = Seq2SeqTrainer(
    model,
    args,
    train_dataset=dataset_tokenized["train"],
    eval_dataset=dataset_tokenized["test"],
    data_collator=collator,
    processing_class=tokenizer,
    compute_metrics=compute_metrics,
)

# %%
trainer.train()

# %%