<jupyter_start><jupyter_text>PyTorch AO (torchao) with int8_dynamic_activation_int8_weight  Imports<jupyter_code>import argparse
import os

import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader
from peft import (
    get_peft_config,
    get_peft_model,
    get_peft_model_state_dict,
    set_peft_model_state_dict,
    LoraConfig,
    PeftType,
    PrefixTuningConfig,
    PromptEncoderConfig,
)

import evaluate
from datasets import load_dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TorchAoConfig, get_linear_schedule_with_warmup, set_seed
from tqdm import tqdm<jupyter_output><empty_output><jupyter_text>Parameters<jupyter_code>batch_size = 16
model_name_or_path = "google/gemma-2-2b"
task = "mrpc"
device = "cuda"
num_epochs = 5
lr = 2e-5

lora_rank = 16
lora_alpha = 32
lora_dropout = 0.1<jupyter_output><empty_output><jupyter_text>Data<jupyter_code>if any(k in model_name_or_path for k in ("gpt", "opt", "bloom")):
    padding_side = "left"
else:
    padding_side = "right"

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side=padding_side)
if getattr(tokenizer, "pad_token_id") is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

datasets = load_dataset("glue", task)
metric = evaluate.load("glue", task)
def tokenize_function(examples):
    # max_length=None => use the model max length (it's actually the default)
    outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
    return outputs
tokenized_datasets = datasets.map(
    tokenize_function,
    batched=True,
    remove_columns=["idx", "sentence1", "sentence2"],
)

# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
# transformers library
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
def collate_fn(examples):
    return tokenizer.pad(examples, padding="longest", return_tensors="pt")
# Instantiate dataloaders.
train_dataloader = DataLoader(
    tokenized_datasets["train"],
    shuffle=True,
    collate_fn=collate_fn,
    batch_size=batch_size,
)
eval_dataloader = DataLoader(
    tokenized_datasets["validation"],
    shuffle=False,
    collate_fn=collate_fn,
    batch_size=batch_size,
)<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>quant_config = TorchAoConfig(quant_type="int8_dynamic_activation_int8_weight")
model = AutoModelForSequenceClassification.from_pretrained(
    model_name_or_path, return_dict=True, device_map=0, torch_dtype=torch.bfloat16, quantization_config=quant_config
)
peft_config = LoraConfig(
    task_type="SEQ_CLS",
    r=lora_rank,
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    target_modules=["q_proj", "v_proj"],
)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()<jupyter_output>trainable params: 3,199,488 || all params: 2,617,545,984 || trainable%: 0.1222<jupyter_text>Training<jupyter_code>optimizer = AdamW(params=model.parameters(), lr=lr)

# Instantiate scheduler
lr_scheduler = get_linear_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
    num_training_steps=(len(train_dataloader) * num_epochs),
)
model.config.use_cache = False
model.to(device)
%%time
for epoch in range(1, num_epochs + 1):
    model.train()
    train_losses = []
    for step, batch in enumerate(tqdm(train_dataloader)):
        batch.to(device)
        outputs = model(**batch)
        loss = outputs.loss
        if not torch.isfinite(loss):
            raise ValueError("non-finite loss encountered")

        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        train_losses.append(loss.item())

    model.eval()
    for step, batch in enumerate(tqdm(eval_dataloader)):
        batch.to(device)
        with torch.no_grad():
            outputs = model(**batch)
        predictions = outputs.logits.argmax(dim=-1)
        predictions, references = predictions, batch["labels"]
        metric.add_batch(
            predictions=predictions,
            references=references,
        )

    eval_metric = metric.compute()
    train_loss = sum(train_losses) / len(train_losses)
    print(f"epoch {epoch} | train loss {train_loss:.4f} |", eval_metric)
# memory: 4122MiB<jupyter_output><empty_output>