Currently you can let SFTTrainer teach your models to learn to predict every token in your dataset, or you can let it train on “completions only”, using the DataCollatorForCompletionOnlyLM class.

I would like something in between, where certain tokens have a higher weight than others.

I thought it would be fairly trivial, but nope.

Here’s what I currently came up with (using Unsloth, so I can try this out on Google Collab):

import transformers
import torch.nn as nn
import torch
from datetime import datetime
from transformers import PreTrainedTokenizerBase
from typing import List, Dict, Any
from unsloth import is_bfloat16_supported
from trl import SFTTrainer

from transformers.utils import logging
logging.set_verbosity_info()
logger = logging.get_logger("transformers.modeling_utils")

class WeightedLossTrainer(SFTTrainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def compute_loss(self, model, inputs, return_outputs=False):
        logger.info("Compute loss starts")
        
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        weight_ranges = inputs.get("weight_ranges")

        batch_size, seq_len, num_classes = logits.shape

        loss_fct = nn.CrossEntropyLoss(reduction='none')
        total_weighted_loss = 0.0
        total_weights = 0.0

        logger.info(f"Doing {batch_size} batch sizes")

        for batch_idx in range(batch_size):
            # Collect weights and losses.
            batch_weighted_losses = []
            for start_idx, end_idx, weight in weight_ranges[batch_idx]:
                logit_chunk = logits[batch_idx, start_idx:end_idx + 1]
                label_chunk = labels[batch_idx, start_idx:end_idx + 1]

                loss = loss_fct(logit_chunk.view(-1, num_classes), label_chunk.view(-1))
                weighted_loss = loss * weight
                batch_weighted_losses.append(weighted_loss.sum())
                total_weights += weight * (end_idx - start_idx + 1)  # Total token count in this range
            
            # Sum the weighted losses for the batch.
            batch_weighted_loss_sum = torch.stack(batch_weighted_losses).sum()
            total_weighted_loss += batch_weighted_loss_sum.detach()

        # Compute the mean loss.
        mean_loss = total_weighted_loss / total_weights
        mean_loss = torch.tensor(mean_loss, dtype=torch.float32, device=logits.device, requires_grad=True)

        logger.info(f"Mean loss: {mean_loss}")

        return (mean_loss, outputs) if return_outputs else mean_loss



class WeightedDataCollator:
    def __init__(self, tokenizer: PreTrainedTokenizerBase):
        self.tokenizer = tokenizer

    def __call__(self, examples: List):
        all_input_ids = []
        all_attention_masks = []
        all_weight_ranges = []

        for entry in examples:
            example_input_ids = []
            example_attention_masks = []
            example_weight_ranges = []
            current_length = 0  # Initialize length counter

            for item in entry['pieces']:
                tokenized = self.tokenizer(item['text'], truncation=True, padding=False, return_tensors='pt')
                input_ids = tokenized.input_ids.squeeze()  # Get tensor, remove batch dimension
                attention_mask = tokenized.attention_mask.squeeze()  # Get tensor, remove batch dimension

                start_idx = current_length
                end_idx = start_idx + len(input_ids) - 1

                example_input_ids.append(input_ids)
                example_attention_masks.append(attention_mask)
                example_weight_ranges.append((start_idx, end_idx, item['weight']))

                current_length = end_idx + 1  # Update current length

            concatenated_input_ids = torch.cat(example_input_ids, dim=0) if example_input_ids else torch.tensor([], dtype=torch.long)
            concatenated_attention_masks = torch.cat(example_attention_masks, dim=0) if example_attention_masks else torch.tensor([], dtype=torch.long)

            pad_length = max_seq_length - len(concatenated_input_ids)  # Assuming max_length = 512 for padding if needed
            if pad_length > 0:
                concatenated_input_ids = torch.cat([concatenated_input_ids, torch.tensor([self.tokenizer.pad_token_id] * pad_length)])
                concatenated_attention_masks = torch.cat([concatenated_attention_masks, torch.tensor([0] * pad_length)])

            all_input_ids.append(concatenated_input_ids)
            all_attention_masks.append(concatenated_attention_masks)
            all_weight_ranges.append(example_weight_ranges)

            logger.info(f"All ranges: {all_weight_ranges}")

        return {
            "input_ids": torch.stack(all_input_ids),
            "attention_mask": torch.stack(all_attention_masks),
            "labels": torch.stack(all_input_ids).clone(),
            "weight_ranges": all_weight_ranges
        }

# Define data collator
data_collator = WeightedDataCollator(tokenizer=tokenizer)

# Prepare dataset for the data collator
#collated_data = data_collator(dataset)

training_args = transformers.TrainingArguments(
      per_device_train_batch_size = 2,
      gradient_accumulation_steps = 4,
      warmup_steps = 5,
      max_steps = 60,
      learning_rate = 2e-4,
      fp16 = not is_bfloat16_supported(),
      bf16 = is_bfloat16_supported(),
      logging_steps = 5,
      optim = "adamw_8bit",
      weight_decay = 0.01,
      lr_scheduler_type = "linear",
      seed = 3407,
      output_dir = "outputs",
      remove_unused_columns=False,
)

from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

trainer = WeightedLossTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    data_collator=data_collator,
    max_seq_length = max_seq_length,
    dataset_num_proc = 2,
    args = training_args,
    packing=False,
    dataset_text_field='text',
    dataset_kwargs={'skip_prepare_dataset': True}
)

trainer_stats = trainer.train()

Each entry in my dataset is an object that has a single property pieces.
pieces is an array, and it contains other objects. Each object inside it has a text and a weight property.

As soon as it starts to calculate the loss, it seems to take a long while (a few seconds) until it eventually just OOMs: ran out of CUDA memory.

So what exactly am I doing wrong, and how can I fix it?