Understanding Hugging Face Trainer for Custom Model Training

Hugging Face's Trainer is an incredibly convenient tool, but

the Trainer source code exceeds 5,000 lines and is simply too large to revisit casually, so I'm writing this up to organize my own understanding.

---

Hugging Face provides several very handy tools.

1. AutoModel Class

Model architectures are bundled inside the Transformers library.

When you call AutoModel.from_pretrained(repo_id), it reads model.safetensors and config.json and maps the parameters to the appropriate model architecture.

model = ModernBertModel.from_pretrained(save_dir)

If a matching architecture exists, it loads and maps the corresponding weights.

The problem arises when you want to freely define a model architecture that the library doesn't provide.

Implementing a model from scratch is the cleanest approach, but the code quickly becomes unwieldy.

That's why the idea is to use AutoModel and reuse only the shared components — the Transformer's embeddings and attention — that are common across models.
(In practice, most Hugging Face models design their task-specific heads exactly this way.)

Let's walk through an example.
For ModernBERT, the ModernBertModel class returns the final hidden states.

Suppose you tokenize the input and produce inputs with a total length of 5.

Calling ModernBertModel's forward pass returns a tensor of shape [batch, length, hidden_size].
Assuming a batch size of 1:

A tensor of shape [1, 5, 768] is returned.

This is the output of the large Transformer model, and the work of tuning the model so that this output is correct is exactly what Pre-Training is.

Language models share the same pipeline for producing this output — though fine-tuning can make minor adjustments — and what you do with that output determines the downstream task.

This is precisely what we want to do.

outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            sliding_window_mask=sliding_window_mask,
            position_ids=position_ids,
            indices=indices,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
            batch_size=batch_size,
            seq_len=seq_len,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        last_hidden_state = outputs[0]

        if self.config.classifier_pooling == "cls":
            last_hidden_state = last_hidden_state[:, 0]
        elif self.config.classifier_pooling == "mean":
            last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum(
                dim=1, keepdim=True
            )

        pooled_output = self.head(last_hidden_state)
        pooled_output = self.drop(pooled_output)
        logits = self.classifier(pooled_output)

Here's one example.

This is the ModernBertForSequenceClassification class, which
receives the output of ModernBertModel internally and runs an additional forward pass.
(In the code above, self.model calls ModernBertModel's forward function.)

This model first applies pooling to the hidden state output we just saw,

then passes it through a head layer (a simple dense network) and a dropout layer to produce a 768-dimensional sentence embedding.

The input sentence is first embedded this way, and the resulting embedding is then passed through a classification layer.

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(logits, labels)

        if not return_dict:
            output = (logits,)
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

Finally, the forward function returns both the loss and the custom output (logits).
If the forward does not return a loss, the Trainer falls back to its internal compute_loss_func — but wiring things up that way is quite inconvenient, so it's much cleaner to handle the loss directly inside the model's forward.

The key points to remember are simple:

1. The shared components of the Transformers library become the backbone of your custom module.
2. Take that backbone's output and rewrite a forward function that computes the final loss.
3. The new class wrapping this logic should extend PreTrainedModel.

class ModernBERTSimCSE(PreTrainedModel):
    def __init__(self, modernbert, config: AutoConfig):
        super().__init__(config)
        self.modernbert = modernbert
        self.pooler = nn.Sequential(
            nn.Linear(config.hidden_size, 768, bias=True),
            nn.Tanh()
        )
        self.loss_fn = Loss()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        # Load config
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        config.architectures = ["ModernBertModel"]
        # Load base model
        base_model = ModernBertModel.from_pretrained(pretrained_model_name_or_path, config=config, **kwargs)
        # Initialize ModernBERTSimCSE
        return cls(modernbert=base_model, config=config)

    def forward(self, 
                anchor_input_ids=None, 
                anchor_attention_mask=None, 
                positive_input_ids=None, 
                positive_attention_mask=None, 
                negative_input_ids=None, 
                negative_attention_mask=None, 
                sentence_1_input_ids=None, 
                sentence_1_attention_mask=None, 
                sentence_2_input_ids=None, 
                sentence_2_attention_mask=None, 
                labels=None):
        try:
            if anchor_input_ids is not None:
                # Handle NLI (train) inputs
                anchor_outputs = self.modernbert(
                    input_ids=anchor_input_ids,
                    attention_mask=anchor_attention_mask,
                    return_dict=True,
                )
                positive_outputs = self.modernbert(
                    input_ids=positive_input_ids,
                    attention_mask=positive_attention_mask,
                    return_dict=True,
                )
                negative_outputs = self.modernbert(
                    input_ids=negative_input_ids,
                    attention_mask=negative_attention_mask,
                    return_dict=True,
                )

                # Apply pooling
                anchor_pooled = self.mean_pooling(anchor_outputs.last_hidden_state, anchor_attention_mask)
                anchor_pooled = self.pooler(anchor_pooled)
                positive_pooled = self.mean_pooling(positive_outputs.last_hidden_state, positive_attention_mask)
                positive_pooled = self.pooler(positive_pooled)
                negative_pooled = self.mean_pooling(negative_outputs.last_hidden_state, negative_attention_mask)
                negative_pooled = self.pooler(negative_pooled)

                # Compute loss
                loss = self.loss_fn.compute_loss(anchor_pooled, positive_pooled, negative_pooled)

                return SequenceClassifierOutput(
                    loss=loss,
                    logits=None,
                    hidden_states=(anchor_pooled, positive_pooled, negative_pooled),
                    attentions=None 
                )

            elif sentence_1_input_ids is not None and sentence_2_input_ids is not None:
                # STS (evaluation) processing
                sentence_1_outputs = self.modernbert(
                    input_ids=sentence_1_input_ids,
                    attention_mask=sentence_1_attention_mask,
                    return_dict=True,
                )
                sentence_2_outputs = self.modernbert(
                    input_ids=sentence_2_input_ids,
                    attention_mask=sentence_2_attention_mask,
                    return_dict=True,
                )

                # Perform pooling
                sentence_1_pooled = self.mean_pooling(sentence_1_outputs.last_hidden_state, sentence_1_attention_mask)
                sentence_1_pooled = self.pooler(sentence_1_pooled)
                sentence_2_pooled = self.mean_pooling(sentence_2_outputs.last_hidden_state, sentence_2_attention_mask)
                sentence_2_pooled = self.pooler(sentence_2_pooled)

                if labels is not None:
                    # Compute loss
                    cosine_similarity = nn.CosineSimilarity(dim=-1)
                    scores = cosine_similarity(sentence_1_pooled, sentence_2_pooled)
                    mse_loss = nn.MSELoss()

                    loss = mse_loss(scores, labels)

                    return SequenceClassifierOutput(
                        loss=loss,  # Loss value
                        logits=(sentence_1_pooled, sentence_2_pooled),
                        attentions=None 
                    )

                return None, (sentence_1_pooled, sentence_2_pooled)

            else:
                raise ValueError("Invalid input configuration for forward method.")
        except Exception as e:
            print(e)
            
     def mean_pooling(self, last_hidden_state: Tensor, attention_mask: Tensor) -> Tensor:
        """
        Mean pooling with attention mask weighting
        """
        weighted_sum = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1)
        mask_sum = attention_mask.sum(dim=1).unsqueeze(-1)
        return weighted_sum / mask_sum.clamp(min=1e-9)

This class is implemented to follow the methodology described in SimCSE and NLI-STS training.

As noted earlier, define a custom wrapper class around your model.

Any layers added on top that require gradient tracking must be defined in the wrapper class's __init__.

The output can be a tuple, but the Trainer expects a dict-like output by default — ModelOutput classes exist precisely to provide that format.

Several subclasses of ModelOutput are available; pick whichever one returns the values you need.

The forward method here is designed to behave differently depending on which input variables are provided.

One useful tip: if labels is passed as an input, the Trainer recognizes it and handles it automatically — so you don't need to include it in the return value of forward.

With minor adjustments to fit this pattern, you can build a new model that incorporates whatever task-specific logic you need.

2. Data Loader

Now, how do we feed data in?

We'll assume we're using the Dataset class from the datasets library as our base.

The requirements for a valid input dataset are:

1. It must be a custom Dataset class that wraps a datasets Dataset.
2. It must be able to return data in dict form.
3. It must expose a __getitem__(self, idx) method that retrieves a sample as that dict.
4. It must expose a __len__(self) method that returns the total number of samples.

Here's an example:

class NLIDataLoader(Dataset):
    def __init__(self, 
                 args: DataArguments, 
                 tokenizer: PreTrainedTokenizer):
        """
        Initialize the NLI data loader.

        Args:
        """
        try:
            login(token=args.hf_data_token)
        except:
            print("Fail to login hgf")
        
        self.dataset = datasets.load_dataset(path=args.train_data, data_dir=args.train_data_dir, split=args.train_data_split)
        self.tokenizer = tokenizer
        self.max_seq_length = args.max_len
        
        if args.data_filtering:
            self.dataset = self.dataset.filter(self.is_valid_input)
        
        self.dataset = self.dataset.map(self.preprocess, batched=True, num_proc=os.cpu_count())
        
    def preprocess(self, examples):
        """
        Preprocess the NLI dataset by tokenizing anchor, positive, and negative samples.

        Args:
            examples (dict): A dictionary containing raw dataset examples.

        Returns:
            dict: Tokenized inputs for anchor, positive, and negative samples.
        """
        anchor = self.tokenizer(
            examples["anchor"],
            max_length=self.max_seq_length,
            padding="max_length",
            truncation=True,
        )
        positive = self.tokenizer(
            examples["positive"],
            max_length=self.max_seq_length,
            padding="max_length",
            truncation=True,
        )
        negative = self.tokenizer(
            examples["negative"],
            max_length=self.max_seq_length,
            padding="max_length",
            truncation=True,
        )
        return {
            "anchor_input_ids": anchor["input_ids"],
            "anchor_attention_mask": anchor["attention_mask"],
            "positive_input_ids": positive["input_ids"],
            "positive_attention_mask": positive["attention_mask"],
            "negative_input_ids": negative["input_ids"],
            "negative_attention_mask": negative["attention_mask"],
        }
        
    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        return self.dataset[idx]

That's all the structure you need.

These field names aren't what a standard Transformer model expects as inputs — but since we customized the forward method earlier, the model handles them just fine.

That said, there's one more piece to prepare: the collator.

The collator's job is to take samples from the Dataset and batch them into the form the model receives during training.

Because the field names are non-standard, as mentioned above, we need to customize the collator as well.

In the code below, there's no real need to inherit from DataCollatorWithPadding — since we're defining a new __call__ method, none of the original behavior carries over anyway.

The collator returns different batch shapes depending on the input format:

class SimCSEDataCollator(DataCollatorWithPadding):
    def __call__(self, features):
        """
        Custom collator to handle both NLI and STS inputs.
        """
        # Check the type of input data
        
        if "anchor_input_ids" in features[0]:
            # NLI data processing
            batch = {
                "anchor_input_ids": [f["anchor_input_ids"] for f in features],
                "anchor_attention_mask": [f["anchor_attention_mask"] for f in features],
                "positive_input_ids": [f["positive_input_ids"] for f in features],
                "positive_attention_mask": [f["positive_attention_mask"] for f in features],
                "negative_input_ids": [f["negative_input_ids"] for f in features],
                "negative_attention_mask": [f["negative_attention_mask"] for f in features],
            }
        elif "sentence_1_input_ids" in features[0]:
            # STS data processing
            batch = {
                "sentence_1_input_ids": [f["sentence_1_input_ids"] for f in features],
                "sentence_1_attention_mask": [f["sentence_1_attention_mask"] for f in features],
                "sentence_2_input_ids": [f["sentence_2_input_ids"] for f in features],
                "sentence_2_attention_mask": [f["sentence_2_attention_mask"] for f in features],
                "labels": [f["labels"] for f in features],
            }
        else:
            raise ValueError("Features do not match NLI or STS format.")

        # Convert lists to tensors
                
        batch = {key: torch.tensor(value, dtype=torch.long) if "labels" not in key else torch.tensor(value, dtype=torch.float) for key, value in batch.items()}
        
        return batch

3. Trainer

The last piece is the Trainer.

trainer = Trainer(
    model=model,
    args=training_args,
    processing_class=tokenizer,
    optimizers=(optimizer,scheduler),
    train_dataset=nli_loader,
    eval_dataset=sts_loader,
    compute_metrics=model.compute_metrics,
    data_collator=SimCSEDataCollator(tokenizer=tokenizer)
)

Each of the following Trainer arguments must be defined correctly:

• model
• processing_class (tokenizer)
• optimizer
• dataset
• compute_metrics
• data_collator

For model, pass the custom class that wraps PreTrainedModel as defined earlier.

This class, as noted earlier, must have a custom forward() function defined.

processing_class is not strictly required, but it is needed when calling push_to_hub to push the tokenizer alongside the model.

optimizer defaults to the AdamW optimizer.

In general, any object that inherits from the Optimizer base class defined in torch.optim.optimizer can be passed here.

For dataset, any format matching the description above is acceptable — specifically, anything that supports dataset[idx] and returns a dict-type item.

For collator, as described earlier, the role is to take the dataset items and transform (and optionally process) them into batches.

The last argument to cover is compute_metrics.

It takes a callback function as input. That function receives prediction data and returns a dict of metrics.

The flow works as follows:

1. The model's forward function is called.
2. The return value is processed into an EvalPrediction object.
3. That EvalPrediction object is passed as input to the compute_metrics function provided earlier.
4. A dict of metrics is returned.

class EvalPrediction:
    """
    Evaluation output (always contains labels), to be used to compute metrics.

    Parameters:
        predictions (`np.ndarray`): Predictions of the model.
        label_ids (`np.ndarray`): Targets to be matched.
        inputs (`np.ndarray`, *optional*): Input data passed to the model.
        losses (`np.ndarray`, *optional*): Loss values computed during evaluation.
    """

    def __init__(
        self,
        predictions: Union[np.ndarray, Tuple[np.ndarray]],
        label_ids: Union[np.ndarray, Tuple[np.ndarray]],
        inputs: Optional[Union[np.ndarray, Tuple[np.ndarray]]] = None,
        losses: Optional[Union[np.ndarray, Tuple[np.ndarray]]] = None,
    ):
        self.predictions = predictions
        self.label_ids = label_ids
        self.inputs = inputs
        self.losses = losses
        self.elements = (self.predictions, self.label_ids)
        if self.inputs is not None:
            self.elements += (self.inputs,)
        if self.losses is not None:
            self.elements += (self.losses,)

    def __iter__(self):
        return iter(self.elements)

    def __getitem__(self, idx):
        if idx < 0 or idx >= len(self.elements):
            raise IndexError("tuple index out of range")
        return self.elements[idx]

    def compute_metrics(self, eval_pred):
        predictions, labels = eval_pred

        # Split the two sentence embeddings
        sentence_1_embeddings, sentence_2_embeddings = predictions[0], predictions[1]

        # Convert to NumPy
        embeddings1 = sentence_1_embeddings
        embeddings2 = sentence_2_embeddings
        labels = labels.flatten()

        # Compute distances and similarities
        cosine_scores = 1 - paired_cosine_distances(embeddings1, embeddings2)
        manhattan_distances = -paired_manhattan_distances(embeddings1, embeddings2)
        euclidean_distances = -paired_euclidean_distances(embeddings1, embeddings2)
        dot_products = [np.dot(emb1, emb2) for emb1, emb2 in zip(embeddings1, embeddings2)]

        # Compute Pearson and Spearman correlations
        
        print(cosine_scores.shape)
        print(labels.shape)
            
        eval_pearson_cosine, _ = pearsonr(labels, cosine_scores)
        eval_spearman_cosine, _ = spearmanr(labels, cosine_scores)

        eval_pearson_manhattan, _ = pearsonr(labels, manhattan_distances)
        eval_spearman_manhattan, _ = spearmanr(labels, manhattan_distances)

        eval_pearson_euclidean, _ = pearsonr(labels, euclidean_distances)
        eval_spearman_euclidean, _ = spearmanr(labels, euclidean_distances)

        eval_pearson_dot, _ = pearsonr(labels, dot_products)
        eval_spearman_dot, _ = spearmanr(labels, dot_products)

        # Return results as a dictionary
        return {
            "pearson_cosine": eval_pearson_cosine,
            "spearman_cosine": eval_spearman_cosine,
            "pearson_manhattan": eval_pearson_manhattan,
            "spearman_manhattan": eval_spearman_manhattan,
            "pearson_euclidean": eval_pearson_euclidean,
            "spearman_euclidean": eval_spearman_euclidean,
            "pearson_dot": eval_pearson_dot,
            "spearman_dot": eval_spearman_dot,
        }

The key thing to examine here is the arguments of the EvalPrediction class.

As long as the forward function returns a properly structured dict, everything will be handled correctly.

In this example, the model returns a pair of logits (more precisely, pooled embedding vectors), packed as a tuple and returned as logits.

Those values are fed into EvalPrediction as the predictions field, and compute_metrics then processes both predictions and labels.

---

To summarize the key points:

• Wrap your custom model around the PreTrainedModel class.
• Design the forward function of your custom model to return an output class that includes the loss value. Custom loss computation should happen here.
• Subclass the Dataset class and prepare data that can be accessed via dataset[idx]. Each item returned must be a dict.
  (The standard keys the model expects are things like "input_ids", "attention_mask", and "labels". If you customize these, make sure the collator and the model's forward function handle them appropriately.)
• During evaluation, the compute_metrics function passed to the Trainer is called. It receives the logits returned by forward, along with the labels. Adjust these as needed.