Adapting Unsloth's implementation for classification tasks

[Gladiator07]

Hi @danielhanchen ,

I am experimenting with training causal models for classification tasks. So far HuggingFace's implementation projects the last hidden states to the final classification layer (instead of a normal language modelling head as in causal lm architecture). I am trying to see if we can stay in unsloth and still finetune for classification.

So far I've tried this notebook: https://github.com/timothelaborie/text_classification_scripts/blob/main/unsloth_classification.ipynb
and it works great. The way shown in the notebook effectively matches the HuggingFace's implementation of AutoModelForSequenceClassification for causal models. But when we want to save the model and actually do the inference, unsloth throws an error of size mis-match as we are trying to copy the weights of last lm_head which we modified in the notebook. There's a monkey-patch to avoid this error and get it working, but seems very ugly. Is there a cleaner way to save the model (lora adapters) and load it for inference with the modified lm_head. Also a native support for finetuning classification models would be great too if it's in the plans.

Again thanks for all the work you do!

[shimmyshimmer]

Interesting we'll take a closer look!

[SvelaT]

Are there any updates on this issue? I am also interested in alternatives to that patch.

[pySilver]

@Gladiator07 are you still using that approach?

[Gladiator07]

Hi @pySilver , nope I just switched to plain HuggingFace for that project

[pySilver]

@Gladiator07 do you actually observe better results with attached classification head VS using something like guided output?

[Gladiator07]

@pySilver In my case the results were slightly better with attached classification head. Also we save on a lot of training memory and inference time in case of classification head compared to a complete lm head output (with like 100k vocabulary prediction) compared to just n_classes in case of classification head.

[pySilver]

@Gladiator07 thx for your feedback! Just to clarify things, by saying "plain HF" you mean using something like Gemma2ForSequenceClassification from transformers? Do you have any example of such setup (fine tuning)?

[Ryan-gsq]

It works fine for me. I'll use the warp mode.model without the last MLP layer and add a classification header.

// model.py
from typing import Optional, Unpack, Any, Type

import torch

from config import cfg
import torch.nn as nn
from transformers.modeling_outputs import SequenceClassifierOutputWithPast

from transformers.cache_utils import Cache
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs


def ModelForClassification(base):
    """Directly modify the forward method of the base model"""

    raw_forward = base.forward
    # Add classification head
    base.classification_head = nn.Linear(
        in_features=base.config.hidden_size, out_features=cfg.model.num_labels, bias=False
    ).to(base.device)
    base.classification_head_loss_fn = nn.BCEWithLogitsLoss().to(base.device)

    # Define new forward function
    def new_forward(*args, **kwargs):
        labels = kwargs.pop("labels", None)
        kwargs["use_cache"] = False
        kwargs["output_hidden_states"] = False
        # Call the original forward method
        transformer_outputs = raw_forward(*args, **kwargs)

        hidden_states = transformer_outputs.last_hidden_state
        logits = base.classification_head(hidden_states[:, -1]).reshape(-1)

        loss = None
        if labels is not None:
            labels = labels.to(logits.dtype).reshape(-1)
            loss = base.classification_head_loss_fn(logits, labels)

        return SequenceClassifierOutputWithPast(
            loss=loss,
            logits=logits,
            hidden_states=hidden_states,
            attentions=transformer_outputs.attentions,
        )

    # Replace the forward method
    base.forward = new_forward

    # Return the modified base model
    return base

// train.py
from unsloth import FastModel


import evaluate
from datasets import Dataset
import torch
from config import cfg

from transformers import (
    get_cosine_schedule_with_warmup,
    AutoTokenizer,
    DataCollatorWithPadding,
)

import pandas as pd
from torch.utils.data import DataLoader

from tqdm.auto import tqdm

from omegaconf import OmegaConf
import math
import time
from accelerate import Accelerator
from accelerate.utils import set_seed
from torch import optim
import bitsandbytes as bnb
from model import ModelForClassification
from peft import LoraConfig, TaskType, get_peft_model


def train():
  model, tokenizer = FastModel.from_pretrained(
      model_name=cfg.model.backbone_path,
      max_seq_length=cfg.model.max_length,
      load_in_4bit=True,
      load_in_8bit=False,
      full_finetuning=False,
  )
  
  // attensiton layer
  model = model.model
  model = ModelForClassification(model)
  # lora
  peft_config = LoraConfig(
      r=cfg_dict["model"]["lora"]["r"],
      lora_alpha=cfg_dict["model"]["lora"]["lora_alpha"],
      lora_dropout=cfg_dict["model"]["lora"]["lora_dropout"],
      bias="none",
      task_type=TaskType.SEQ_CLS,
      target_modules=cfg_dict["model"]["lora"]["target_modules"],
      modules_to_save=cfg_dict["model"]["lora"]["modules_to_save"],
  )
  
  model = get_peft_model(model, peft_config)
  
  model.print_trainable_parameters()

...

[pySilver]

@Ryan-gsq thanks for the code! Can you please add few more lines / details?

1. How your training samples looks like? Is that json-formatted responses or just classes that are in output vocabulary?
2. How do you specify list of supported classification classes?
3. Does this header returns confidence too?

[pySilver]

@Ryan-gsq ping ;)

[Ryan-gsq]

@pySilver

import os

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

script_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(script_dir)
os.environ["WANDB_PROJECT"] = "qwen3-lora-cls"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import random
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import wandb


# patch torch repr
def custom_repr(self):
    return f"{{Tensor:{tuple(self.shape)}}} {{Dtype:{str(self.dtype)}}}  {{Device:{str(self.device)}}} {original_repr(self)}"


original_repr = torch.Tensor.__repr__
torch.Tensor.__repr__ = custom_repr

SEED = 3407

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True


# ====== REST OF SCRIPT ======
from unsloth import (
    FastLanguageModel,
    is_bfloat16_supported,
    UnslothTrainer,
    UnslothTrainingArguments,
)
from datasets import Dataset
from transformers import (
    DataCollatorWithPadding,
    TrainerCallback,
)
from transformers.modeling_outputs import SequenceClassifierOutputWithPast
from sklearn.metrics import roc_auc_score, recall_score
from accelerate import Accelerator
from peft import TaskType, LoraConfig, get_peft_model

###
def ModelForClassification(base, hidden_size):
    """Directly modify the forward method of the base model"""

    raw_forward = base.forward
    # Add classification head
    base.classification_head = nn.Linear(
        in_features=hidden_size, out_features=1, bias=False
    ).to(base.device)
    nn.init.xavier_uniform_(base.classification_head.weight)
    # Define new forward function
    def new_forward(*args, **kwargs):
        labels = kwargs.pop("labels", None)
        # Call the original forward method
        transformer_outputs = raw_forward(*args, **kwargs)

        hidden_states = transformer_outputs.last_hidden_state
        mask = kwargs['attention_mask']
        # For right padding, we need to get the last non-padding token
        last_token_indices = mask.sum(dim=1) - 1
        logits = base.classification_head(hidden_states[torch.arange(hidden_states.shape[0]), last_token_indices]).reshape(-1)

        loss = None
        if labels is not None:
            labels = labels.to(logits.dtype).reshape(-1)
            loss = F.binary_cross_entropy_with_logits(logits, labels)

        return SequenceClassifierOutputWithPast(
            loss=loss,
            logits=logits,
            hidden_states=hidden_states if kwargs.get('output_hidden_states') else None,
            attentions=transformer_outputs.attentions if kwargs.get('output_attentions') else None,
        )

    # Replace the forward method
    base.forward = new_forward

    # Return the modified base model
    return base

### patch unsloth, skip lm_head embed size fix
import unsloth_zoo.tokenizer_utils as tokenizer_utils

def fix_untrained_tokens_(*k, **kk):
    ...
tokenizer_utils.fix_untrained_tokens = fix_untrained_tokens_

# Fix OOM issue when resuming from checkpoint
from peft.peft_model import PeftModel
from functools import wraps
# Store the original load_adapter method
original_load_adapter = PeftModel.load_adapter

@wraps(original_load_adapter)
def patched_load_adapter(
    self,
    model_id,
    adapter_name,
    is_trainable=False,
    torch_device=None,
    autocast_adapter_dtype=True,
    ephemeral_gpu_offload=False,
    low_cpu_mem_usage=False,
    **kwargs,
):
    # If torch_device is None, use self.device instead of infer_device()
    if torch_device is None and hasattr(self, "device"):
        torch_device = str(self.device)

    # Call the original method with our modified parameters
    return original_load_adapter(
        self,
        model_id,
        adapter_name,
        is_trainable,
        torch_device,
        autocast_adapter_dtype,
        ephemeral_gpu_offload,
        low_cpu_mem_usage,
        **kwargs,
    )
# Apply the patch
PeftModel.load_adapter = patched_load_adapter


max_seq_length = 768
dtype = torch.bfloat16
load_in_4bit = True
MODEL_NAME = "unsloth/Qwen3-8B-Base-unsloth-bnb-4bit"


# 获取当前设备ID，用于分布式训练
accelerator = Accelerator(mixed_precision="bf16")
device_index = accelerator.process_index
device_map = {"": device_index}

print(f"Process index: {device_index}, Device map: {device_map}")

# Load model with flash attention
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_NAME,
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
    device_map=device_map,
    use_gradient_checkpointing="unsloth",
)

model = ModelForClassification(model.model, model.config.hidden_size)

peft_config = LoraConfig(
    r=32,
    lora_alpha=32,
    lora_dropout=0.1,
    bias="none",
    task_type=TaskType.SEQ_CLS,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
    ],
    modules_to_save=[
        "classification_head"
    ],
)

model = get_peft_model(model, peft_config=peft_config)

print("*" * 128)
print(model.print_trainable_parameters())
print("*" * 128)

df = pd.read_parquet("./data/data_combine_filter_ids_train.parquet")
df["attention_mask"] = df["input_ids"].apply(lambda x: [1] * len(x))

# 创建数据集时设置缓存文件
train_end_index = int(len(df) * 0.95)
train_df = df[:train_end_index]
train_dataset = Dataset.from_pandas(train_df, preserve_index=False)

eval_df = df[train_end_index:]
eval0_df = eval_df[eval_df["label"] == False][:2048]
eval1_df = eval_df[eval_df["label"] == True].sample(len(eval0_df), ignore_index=True)
eval_df = pd.concat([eval0_df, eval1_df], ignore_index=True)
eval_df = eval_df.sort_values(
    by="input_ids",
    key=lambda x: x.apply(lambda x: len(x)),
    ignore_index=True,
    ascending=False,
)
eval_dataset = Dataset.from_pandas(eval_df, preserve_index=False)

cls_weights = list(train_df["label"].value_counts())
cls_weights = torch.tensor(data=cls_weights, device=model.device) / cls_weights[0]
print('cls weight :' , cls_weights)

# 打印数据集长度
print("*" * 128)
print(
    f"训练集长度: {len(train_dataset)}, 评估集长度: {len(eval_dataset)}, 训练集占比: {len(train_dataset) / len(df):.2%}, 评估集占比: {len(eval_dataset) / len(df):.2%}"
)
print(f"训练 {train_df['label'].value_counts()}")
print(f"评估 {eval_df['label'].value_counts()}")
print("*" * 128)

# 释放DataFrame内存
del df
del train_df
del eval0_df
del eval1_df
del eval_df

def compute_metrics(eval_preds):
    valid_predictions = 1 / (1 + np.exp(-eval_preds.predictions))
    valid_references = eval_preds.label_ids

    cls_recall = recall_score(valid_references, valid_predictions >= 0.6)
    cls_auc = roc_auc_score(valid_references, valid_predictions)

    return {
        "cls_recall": cls_recall,
        "cls_auc": cls_auc,
    }


def compute_loss(outputs, labels: torch.Tensor, num_items_in_batch: torch.Tensor):
    logits = outputs.logits
    weight = torch.where(labels == 1.0, cls_weights[1], cls_weights[0])
    loss = F.binary_cross_entropy_with_logits(logits, labels, weight=weight)
    return loss


class WandbCallback(TrainerCallback):
    def on_train_begin(self, args, state, control, model, **kwargs):
        if accelerator.is_main_process:
            print(f"{__file__} is saved {wandb.save(__file__)}")


trainer = UnslothTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=DataCollatorWithPadding(
        tokenizer=tokenizer,
        max_length=max_seq_length,
        pad_to_multiple_of=8,
    ),
    callbacks=[WandbCallback()],
    compute_metrics=compute_metrics,
    compute_loss_func=compute_loss,
    args=UnslothTrainingArguments(
        run_name="bce_loss",
        remove_unused_columns=False,

        num_train_epochs=2,
        max_grad_norm=1,
        per_device_train_batch_size=8,
        gradient_accumulation_steps=1,
        per_device_eval_batch_size=8,
        eval_accumulation_steps=1,
        gradient_checkpointing=False,
        
        warmup_ratio=0.1,
        learning_rate=1e-4,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        weight_decay=0.01,
        logging_steps=1,
        optim="adamw_8bit",
        lr_scheduler_type="cosine",
        seed=SEED,
        output_dir="bce_outputs",
        report_to="wandb",
        eval_steps=500,
        metric_for_best_model="cls_auc",
        eval_strategy="steps",
        save_strategy="steps",
        save_steps=500,
        save_total_limit=5,
        local_rank=accelerator.process_index,
        dataloader_num_workers=accelerator.num_processes * 4,
        torch_compile=True,
    ),
)

trainer.train()

input shape like

wandb:

[Ryan-gsq]

I found something very strange. In the case of an 8-card 4090 setup, with per_device_train_batch_size=4 and gradient_accumulation_steps=2, if I use gradient_checkpointing=True, the training time is almost halved. In the first 300 steps, the loss with gradient_checkpointing set to True is only about the value of gradient_checkpointing set to False at around 150 steps.

The reason is still unclear.

[pluesclues]

Btw we have official AutoModelForSequenceClassification support, you use the below code to initialize it and use num_labels to initialize the number of classes you want. I will ask to push the reward modeling notebook from this PR soon. #2710

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_name, 
    max_seq_length=4096,
    dtype=None,
    load_in_4bit=False, #Can load true if planning to use unsloth qlora 
    #token=hf_auth_token,
    num_labels=num_labels, # This ensures a classification head is added
)

Similar to: #372