Support LLMs?

[RogerRogerUSC]

Hi @f-dangel,

I am trying to use Backpack in LLM training, but it failed. So I would like to inquire whether it supports LLM training? Thank you!

[f-dangel]

Hi Zhe,

BackPACK's support for transformers is relatively poor at the moment.
What are you trying to achieve?
Happy to point to alternative solutions or provide information how to get BackPACK to do the job.

Best,
Felix

[RogerRogerUSC]

Thank you so much for your reply, @f-dangel . We are thinking about using this tool in LLMs. My collaborator @ZidongLiu will have other questions about this tool.

[ZidongLiu]

Hi @f-dangel ,

I am trying to get DiagHessian for opt-125m model, and here is small code snippet I am using.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset as huggingface_load_dataset

import backpack

device = "cuda"

torch_dtype = torch.float16
hf_model_name = "facebook/opt-125m"
tokenizer = AutoTokenizer.from_pretrained(hf_model_name, padding_side="left", truncate_side="left")
model = AutoModelForCausalLM.from_pretrained(hf_model_name, torch_dtype=torch_dtype).to(device)
loss_fn = torch.nn.CrossEntropyLoss()
backpack.extend(model)
backpack.extend(loss_fn)

max_length = 32
dataset = huggingface_load_dataset("glue", "sst2")
raw_train_dataset = dataset["train"]
# raw_test_dataset = dataset["validation"]


class SST2Template:
    verbalizer = {0: " bad", 1: " good"}

    def get_verbalizer_id(self, tokenizer):
        return {k: tokenizer.encode(v)[-1] for k, v in self.verbalizer.items()}

    def verbalize_for_pred(self, sample):
        text = sample["sentence"].strip()
        return f"{text} It was"

    def verbalize(self, sample):
        label = sample["label"]
        return f"{self.verbalize_for_pred(sample)}{self.verbalizer[label]}"


template = SST2Template()
encoded_train_texts = list(map(template.verbalize, raw_train_dataset))


class CustomLMDataset(torch.utils.data.DataLoader):
    def __init__(self, texts, tokenizer, max_length):
        self.texts = texts
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        input_text = self.texts[idx]
        input_ids = self.tokenizer.encode(input_text, add_special_tokens=True)
        # left_truncation
        if len(input_ids) > self.max_length:
            input_ids = input_ids[-self.max_length :]
        return torch.tensor(input_ids, dtype=torch.long)


def collate_fn(batch):
    # Pad sequences to the max length in the batch
    padded_batch = tokenizer.pad(
        {"input_ids": batch},
        padding=True,
        max_length=max_length,
        return_tensors="pt",
    )
    input_ids = padded_batch["input_ids"]
    attention_mask = padded_batch["attention_mask"]
    return (
        {
            "input_ids": input_ids[:, :(-1)],
            "attention_mask": attention_mask[:, :(-1)],
        },  # get entire sentence except last token
        input_ids[:, 1:],  # get_all except first token (for loss calculation)
    )


train_dataset = CustomLMDataset(encoded_train_texts, tokenizer, max_length=max_length)
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=True,
    collate_fn=collate_fn,
)

first_batch_input, first_batch_target = next(iter(train_loader))
first_batch_input["input_ids"] = first_batch_input["input_ids"].to(device)
first_batch_input["attention_mask"] = first_batch_input["attention_mask"].to(device)
first_batch_target = first_batch_target.to(device)


verbalizer_id_map = template.get_verbalizer_id(tokenizer)
n_candidate = len(verbalizer_id_map)
verbalizer_id_list = [verbalizer_id_map[i] for i in range(n_candidate)]


def last_token_loss_fn(batch_pred, targets):
    logits = batch_pred.logits
    last_token_batch_pred = logits[:, -1, verbalizer_id_list].view(-1, len(verbalizer_id_list))
    last_token_label = (targets[:, -1] == verbalizer_id_map[1]).to(int)

    return loss_fn(last_token_batch_pred, last_token_label)


with backpack.backpack(backpack.extensions.DiagHessian()):
    outputs = model(
        input_ids=first_batch_input["input_ids"], attention_mask=first_batch_input["attention_mask"]
    )

    loss = last_token_loss_fn(outputs, first_batch_target)
    loss.backward()
    print("last token loss finish")


def full_sentence_loss_fn(batch_pred, targets):
    return loss_fn(batch_pred.logits.view(-1, batch_pred.logits.size(-1)), targets.view(-1))


with backpack.backpack(backpack.extensions.DiagHessian()):
    outputs = model(
        input_ids=first_batch_input["input_ids"], attention_mask=first_batch_input["attention_mask"]
    )

    loss = full_sentence_loss_fn(outputs, first_batch_target)
    loss.backward()
    print("full sentence loss finish")

If I use the full_sentence_loss_fn, I got CUDA out of memory:

Want to see if it's too much data to store when calculating DiagHessian?

I want to swtich to use my own last_token_loss_fn to potentially reduce memory cost, I got this error:

Wondering if it's because of my implementation for the last_token_loss_fn misses anything backpack wants.

NOTE: both loss functions are able to help me train with any torch optimizer.

Versions:

python:  3.10.15
torch: 2.3.1
backpack-for-pytorch: 1.7.1

[f-dangel]

Hi,

there are two problems with your approach which I think eliminate using BackPACK to achieve what you want:

1. Computing the exact Hessian diagonal will be very expensive for models with high-dimensional outputs (i.e. any LLM). This is just a result from the algorithm we use to compute the Hessian diagonal.
2. Since you are interested in computing the Hessian diagonal, all layers in your model must be supported by BackPACK. This will not be the case, since BackPACK is lacking support for LLM-specific layers.

There is a potential solution: If you are fine computing an unbiased estimate of the Hessian diagonal, you can use Hessian-vector products to do that. For that, you don't need BackPACK:

1. You could use the PyHessian library for Hessian-vector products and implement the diagonal estimation method (https://arxiv.org/pdf/2006.00719, Figure 4) yourself.
2. You could use the curvlinops library (https://github.com/f-dangel/curvlinops) and use its HessianLinearOperator and diagonal estimation algorithm. The library has an example how to estimate diagonals and how to use LLMs. Disclaimer: I am a maintainer of this library.