Principle:Bigscience workshop Petals Classification Model Loading

Knowledge Sources	Petals: Collaborative Inference and Fine-tuning of Large Models Petals
Domains	NLP, Classification, Distributed_Computing
Last Updated	2026-02-09 14:00 GMT

Overview

Loading a distributed large language model configured for sequence classification with a task-specific linear head and optional prompt tuning embeddings, enabling text classification through the distributed Petals network.

Description

Classification Model Loading extends the distributed model loading principle for supervised classification tasks. Instead of a causal LM head, the model is loaded with a linear classification head (nn.Linear(hidden_size, num_labels)) and is configured for prompt tuning to enable training without modifying frozen remote weights.

The loaded model consists of:

• Local trainable components: Classification score head, prompt embeddings (if ptune enabled)
• Remote frozen components: All transformer blocks via RemoteSequential
• Local frozen components: Token embeddings, layer normalization

This configuration enables the standard supervised learning pipeline: forward pass produces logits per class, cross-entropy loss is computed, and gradients flow back through the distributed autograd to update only the local trainable parameters.

Usage

Use this principle when you need to perform text classification (sentiment analysis, topic classification, NLI, etc.) with a distributed large language model. Configure num_labels for your task and enable prompt tuning via tuning_mode and pre_seq_len in the model config.

Theoretical Basis

Sequence classification with LLMs:

Given input tokens, the model produces hidden states $H\in {\mathbb{ℝ}}^{n\times d}$. The last token's hidden state is used for classification:

$logits=W\cdot h_{last}+b\in {\mathbb{ℝ}}^C$

where W is the classification head weight matrix, C is the number of classes.

With prompt tuning: The effective sequence becomes [prompts; input_tokens], but classification still uses the last non-prompt token's representation.

# Abstract classification model setup
model = load_distributed_model(model_name, task="classification")
model.config.num_labels = num_classes
model.config.tuning_mode = "ptune"
model.config.pre_seq_len = 16

# Trainable: model.score (Linear), model.prompt_embeddings (Embedding)
# Frozen: model.model.layers (RemoteSequential), model.model.embed_tokens

Related Pages

Implemented By

• Implementation:Bigscience_workshop_Petals_DistributedLlamaForSequenceClassification_From_Pretrained