Petals

Enables running inference on models larger than any single machine's memory by splitting transformer blocks across a peer-to-peer network discovered via DHT. The client queries the DHT to locate servers hosting different model blocks, then routes input sequentially through the network with RemoteSequenceManager determining optimal paths. Attention states are cached across servers to optimize multi-token generation, eliminating redundant computation.

Implements a Distributed Hash Table (DHT) for decentralized peer discovery where servers register themselves and clients query to locate which peers host which model blocks. The DHT stores mappings of model block identifiers to peer addresses and connection metadata. RemoteSequenceManager uses DHT lookups to construct optimal routing paths through the network, handling peer churn by re-querying when connections fail.

Manages server startup, block loading, DHT registration, and graceful shutdown. When a server starts, it loads assigned transformer blocks into memory, registers itself in the DHT with block availability metadata, and begins accepting inference requests. On shutdown, it deregisters from DHT and releases resources. The Server class orchestrates this lifecycle with health monitoring.

Implements TransformerBackend that executes individual transformer blocks (attention, MLP, layer norm) on server hardware. The backend handles forward passes, backward passes (for fine-tuning), and optimization of block execution (kernel fusion, quantization). ModuleContainer wraps blocks and manages their lifecycle on the server.

Implements MemoryCache component that manages attention key-value caches and intermediate activations on servers with configurable eviction policies. When cache memory exceeds limits, the system evicts least-recently-used entries or uses other strategies to free space. This prevents out-of-memory errors during high-throughput inference with many concurrent sessions.

Supports running multiple model architectures (BLOOM, Llama, Falcon, Mixtral) with different precision formats (float32, float16, bfloat16, int8 quantization). The system automatically handles precision conversion at peer boundaries and optimizes computation for the target precision. This enables flexibility in model choice and memory/speed trade-offs.

Enables fine-tuning of large distributed models using parameter-efficient methods (LoRA, prefix tuning, etc.) where only a small fraction of parameters are updated while frozen base model blocks remain distributed across peers. The fine-tuning adapters are stored locally on the client, and gradients are computed only for adapter parameters during backpropagation through the frozen distributed blocks.

Optimizes multi-token generation by caching intermediate attention states (key-value pairs) across distributed servers, eliminating redundant computation of previously processed tokens. When generating the next token, only the new token is processed through the full network, and cached attention states from prior tokens are reused. This reduces per-token latency by 30-50% in typical generation workloads.

Provides AutoDistributedModelForCausalLM and RemoteGenerationMixin classes that mimic HuggingFace Transformers API exactly, allowing users to load and run distributed models with zero code changes from standard Transformers usage. The model loading, tokenization, and generation interfaces are identical to local models, abstracting away all distributed complexity.

Supports both public swarms (shared community resources for general inference) and private swarms (isolated networks for sensitive data or proprietary models). Private swarms can be created by running servers with custom DHT bootstrap nodes, isolating them from the public network. Access control is enforced at the server level through authentication tokens or IP whitelisting.

Implements BlockSelection and swarm balancing mechanisms to distribute inference load across peers hosting the same model blocks. When multiple peers host redundant blocks, the system selects the least-loaded peer based on current queue depth and response latency. This prevents bottlenecks where a single peer becomes overloaded while others remain idle.

Implements RemoteSequential class that manages execution of transformer blocks distributed across multiple peers as a sequential pipeline. Each block's forward pass is executed on its hosting peer, with intermediate activations transmitted between peers. The RemoteSequenceManager determines optimal routing through the block sequence, handling peer failures by re-routing around unavailable blocks.

Manages InferenceSession objects that maintain stateful connections and cached state (attention KV caches, position embeddings) across multiple inference steps. Sessions persist peer connections and cache metadata, enabling efficient multi-token generation without re-establishing connections or recomputing attention for prior tokens.

Manages distribution of transformer model blocks across peers, determining which blocks are hosted on which servers. The system supports flexible block assignment strategies (contiguous blocks per peer, interleaved distribution, etc.) and handles block replication for redundancy. ModuleContainer on each server manages the assigned blocks and their execution.