Principle:Bigscience workshop Petals Distributed Model Configuration
| Knowledge Sources | |
|---|---|
| Domains | Distributed_Computing, NLP, Model_Configuration |
| Last Updated | 2026-02-09 14:00 GMT |
Overview
A configuration pattern that extends HuggingFace model configs with distributed infrastructure attributes required for peer-to-peer inference through Petals.
Description
Distributed Model Configuration is the pattern by which Petals adapts standard HuggingFace model configuration classes (BloomConfig, LlamaConfig, FalconConfig, MixtralConfig) for use in a decentralized inference network. Each model family requires a Distributed*Config class that combines the base model configuration with three Petals-specific mixin classes:
- ClientConfig: Settings for connecting to remote servers (DHT prefix, initial peers)
- PTuneConfig: Parameters for prompt tuning support
- LMHeadConfig: Language model head configuration
The distributed config also defines critical per-model metadata:
- block_class: The wrapped transformer block class to instantiate on servers
- attn_class: The attention implementation (determines KV cache sizing)
- block_prefix: The naming prefix for transformer layers in the state dict
- num_key_value_groups: Number of KV groups for grouped-query or multi-query attention
The from_pretrained override in each distributed config automatically derives a DHT prefix from the model name, enabling peer discovery without manual configuration.
Usage
Use this principle when adding support for a new model architecture to Petals. Each new model family requires a distributed config class that maps the base HuggingFace config to Petals' distributed execution requirements. The config class determines how blocks are split, how KV caches are sized, and how peers discover each other.
Theoretical Basis
The configuration bridge pattern:
In standard HuggingFace, a model config describes the architecture (hidden size, number of layers, attention heads). In Petals, this must be extended with:
- Block routing metadata: Which class implements each transformer block and how layers are named in the state dict
- Attention type metadata: Whether the model uses multi-head, multi-query, or grouped-query attention (affects KV cache memory)
- DHT discovery: A prefix that allows peers to find servers hosting blocks for the same model
Pseudo-code logic:
# Abstract pattern for distributed config
class DistributedModelConfig(BaseModelConfig, ClientConfig, PTuneConfig, LMHeadConfig):
block_class = WrappedBlock # Server-side block implementation
attn_class = AttentionImpl # Attention class for cache sizing
block_prefix = "model.layers" # State dict layer prefix
num_key_value_groups = ... # KV groups (architecture-dependent)
@classmethod
def from_pretrained(cls, model_name, *, dht_prefix=None, **kwargs):
if dht_prefix is None:
dht_prefix = derive_prefix(model_name)
config = super().from_pretrained(model_name, dht_prefix=dht_prefix, **kwargs)
config.pad_token_id = config.pad_token_id or 0
return config
Design trade-offs:
- Multiple inheritance vs. composition: Petals uses multiple inheritance (mixin pattern) to keep config classes compatible with HuggingFace serialization while adding distributed attributes
- Static vs. dynamic KV groups: Some models (Falcon) compute num_key_value_groups dynamically based on architecture flags; others (Bloom, Mixtral) use a fixed value