LeaderWorkerSet: Serving Models Too Big for One Node
โ kubernetes, lws, llm, inference, gpu, distributed โ 10 min read
Some models no longer fit on one machine. Llama 3.1 405B needs on the order of 800 GB of GPU memory, which is more than a single 8-GPU node holds. To serve it you have to split one model across several nodes and make those nodes behave like one inference server. Kubernetes has no built-in object for "these N pods are one replica." LeaderWorkerSet (LWS) is the API that fills that gap. This post explains the problem it solves, how it works underneath, and how to deploy a real multi-node model with it.
Here is the whole picture before the walkthrough. A big model is split across GPUs with tensor parallelism within a node and across nodes with pipeline parallelism, so multi-node serving means several pods must run as one server. Core Kubernetes objects do not model that group: a Deployment's pods are independent, and a StatefulSet's pods are peers, so neither says "1 leader plus M workers is one replica." LWS fills the gap by making a group (1 leader + size-1 workers) the unit of replication and running N such groups that schedule, scale, restart, and roll out together. Under the hood it composes StatefulSets, gives every pod a stable DNS name, and injects LWS_LEADER_ADDRESS and LWS_GROUP_SIZE so members find each other, while restartPolicy: RecreateGroupOnPodRestart restarts the whole group if any pod dies, which is what collective communication libraries need. Everything below is checked against the LWS documentation.
The problem: one model, many nodes
To understand LWS you first need the two ways a model gets split, because they decide how many pods a group needs.
Tensor parallelism (TP) splits an individual layer's math across several GPUs. Each GPU holds a slice of every weight matrix, and the GPUs exchange partial results on every forward pass. This chatter is constant and latency-sensitive, so TP wants GPUs connected by a fast link (NVLink inside one node). TP therefore usually stays within a single node.
Pipeline parallelism (PP) splits the model by layers. Node A runs layers 1 to 40, node B runs layers 41 to 80, and activations pass from A to B once per stage. The traffic between stages is smaller and less latency-sensitive than TP traffic, so PP tolerates a network hop and is how you span nodes.
Tensor parallelism (within a node): Pipeline parallelism (across nodes):
GPU0 GPU1 GPU2 GPU3 Node A Node B \ fast NVLink / layers 1-40 โโโถ layers 41-80 split one layer (activations cross the network)Caption: TP splits a layer across GPUs in a node; PP splits layers across nodes.
A 405B model uses both: tensor parallelism across the 8 GPUs inside each node, and pipeline parallelism across two nodes. That is two pods, 16 GPUs, one model, one server. Now the Kubernetes question: how do you express "two pods that are one server" as a workload?
Why Deployment and StatefulSet do not fit
A Deployment runs identical, independent pods behind a Service. Any pod can serve any request, and pods know nothing about each other. That is wrong here: the two pods are not interchangeable replicas, they are halves of one server.
A StatefulSet runs pods with stable identities and ordered startup, which is closer, but its pods are still peers of one role. It has no notion of a leader plus workers, no notion of "these size pods form one replica, and I want N replicas," and no group-level restart or rollout.
What multi-node inference actually needs is a group as the unit: a fixed set of pods that start together, get scheduled together, fail together, and scale as a whole. That is exactly what LWS adds (concepts).
What LeaderWorkerSet is
An LWS defines a group made of one leader and size - 1 workers, and runs replicas copies of that group. The group is the unit of replication. Scaling to 3 replicas means three independent copies of the whole leader-plus-workers server, not three loose pods.
LeaderWorkerSet (replicas = 2, size = 2)
Group 0 Group 1 โโโ leader vllm-0 โโโ leader vllm-1 โโโ worker vllm-0-1 โโโ worker vllm-1-1 (one model server) (a second, independent model server)Caption: each group is one multi-node model server; replicas are independent copies.
Two templates let the leader and workers differ. leaderTemplate defines the leader pod, workerTemplate defines the workers, and if you omit leaderTemplate the worker template applies to both. This matters because the leader often does extra work, like running the API server and the coordination head, while workers just contribute GPUs.
How LWS works under the hood
LWS does not manage pods directly. It composes StatefulSets, which is a deliberate reuse of a battle-tested controller (concepts):
- One leader StatefulSet named
<lws-name>withreplicaspods:<lws-name>-0,<lws-name>-1, and so on. - For each leader, one worker StatefulSet named
<lws-name>-<group-index>withsize - 1pods starting at ordinal 1:<lws-name>-0-1,<lws-name>-0-2, and so on.
Building on StatefulSet gives every pod a stable, predictable DNS name through a headless Service, so ranks and peers are known without a discovery service. It also gives parallel pod startup, so all workers in a group come up at once instead of one at a time, which matters because collective communication (NCCL, MPI) cannot initialize until every member is present.
On top of that, LWS adds the group-level behaviour StatefulSet lacks. It injects environment variables so members find each other: the leader learns the group size from LWS_GROUP_SIZE, and each worker learns where to connect from LWS_LEADER_ADDRESS, the leader's stable DNS name. It supports a startupPolicy that controls whether workers wait for the leader to exist or to become ready. And it supports an all-or-nothing restart policy, which is the next section.
The failure model: all or nothing
This is the part that separates LWS from just running pods. In a distributed inference group, every pod holds a shard of the model and a rank in the collective. If one worker dies, the collective is broken: the survivors block waiting for a peer that is gone, and the server produces nothing. Restarting only the dead pod does not help, because the collective has to be re-established from scratch with all ranks present.
restartPolicy: RecreateGroupOnPodRestart handles this. If any pod in the group restarts, LWS restarts the entire group together, so the collective re-forms cleanly. The alternative, None, leaves individual pod restarts to the StatefulSet, which is fine only for workloads where members are independent. For sharded inference you almost always want the group restart.
worker vllm-0-1 crashes โ RecreateGroupOnPodRestart โผ whole group (vllm-0 + vllm-0-1) restarts together โผ collective re-initializes with all ranks presentCaption: a single pod failure recreates the whole group so the collective can re-form.
Walkthrough: Llama 3.1 405B with vLLM
Here is the shape of the official LWS vLLM example, serving a 405B model across two nodes. replicas: 2 runs two independent copies of the server; size: 2 makes each copy a leader plus one worker.
apiVersion: leaderworkerset.x-k8s.io/v1kind: LeaderWorkerSetmetadata: name: vllmspec: replicas: 2 leaderWorkerTemplate: size: 2 restartPolicy: RecreateGroupOnPodRestartNotice size: 2 plus the two forms of parallelism means each replica uses 16 GPUs: 8 per pod (tensor parallel) across 2 pods (pipeline parallel).
The leader starts the coordination head and the vLLM OpenAI server. It uses LWS_GROUP_SIZE to tell the head how many nodes to expect before serving.
leaderTemplate: spec: containers: - name: vllm-leader image: vllm/vllm-openai:v0.8.5 command: - sh - -c - "bash .../multi-node-serving.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); python3 -m vllm.entrypoints.openai.api_server --port 8080 --model meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline_parallel_size 2" resources: limits: { nvidia.com/gpu: "8" }Notice --ray_cluster_size=$(LWS_GROUP_SIZE) makes the leader wait until all pods have joined before it starts serving, and --tensor-parallel-size 8 --pipeline_parallel_size 2 maps directly onto 8 GPUs per pod across the 2-pod group.
The worker joins the leader's coordination head using the injected leader address, then contributes its GPUs.
workerTemplate: spec: containers: - name: vllm-worker image: vllm/vllm-openai:v0.8.5 command: - sh - -c - "bash .../multi-node-serving.sh worker --ray_address=$(LWS_LEADER_ADDRESS)" resources: limits: { nvidia.com/gpu: "8" }Notice the worker runs no server of its own. It connects to $(LWS_LEADER_ADDRESS) and becomes a node in the leader's cluster. This is the whole leader-worker split: the leader coordinates and serves, the workers lend GPUs.
Finally, a Service targets only the leaders, since the leader is the one running the HTTP server. LWS labels leader pods so the selector can pick them out.
apiVersion: v1kind: Servicemetadata: { name: vllm-leader }spec: selector: leaderworkerset.sigs.k8s.io/name: vllm role: leader ports: [{ port: 8080, targetPort: 8080 }]Notice the selector combines the LWS name label with your own role: leader label, so requests reach a leader, never a worker.
client โโโถ Service (leaders only) โ leader pod โโ coordinates โโโถ worker pod (vLLM server) (GPUs, no server)Caption: traffic hits the leader, which drives its workers to run the sharded model.
Scaling and rollouts
Because the group is the unit, scaling is coarse and correct. Raising replicas from 2 to 3 adds one more complete server, all 16 GPUs at once, not a stray pod. A Horizontal Pod Autoscaler can scale an LWS through the leader's scale subresource.
Rollouts are group-aware too. rolloutStrategy.rollingUpdateConfiguration takes maxUnavailable and maxSurge counted in groups, so an update replaces whole servers at a time and never leaves a half-updated group serving traffic (rollout strategy).
Gotchas worth knowing
- Shared memory. GPU collectives use
/dev/shm, and the 64 MB default is far too small. Mount a memory-backed/dev/shm, as the example does, or startup hangs. - Group size must match your parallelism.
sizehas to equal the number of pods your TP and PP mapping needs. Here PP is 2, sosizeis 2. A mismatch means the leader waits forever for a node that never joins. - Keep the group on the fast fabric. For latency-sensitive setups, LWS supports exclusive 1:1 topology placement so a whole group lands in one topology domain (a rack or block), keeping cross-pod traffic on the fastest network (concepts). Without it the scheduler may spread a group across slow links.
- Restart the group, not the pod. Use
RecreateGroupOnPodRestartfor sharded inference. A lone worker restart cannot rejoin a broken collective. - Readiness is the leader's. The server runs on the leader, so health checks and the Service target the leader; workers have no endpoint.
When to use LWS, and when not to
Use LWS when a single model instance spans more than one node: a model too large for one node's GPUs, or a serving setup that needs pipeline parallelism across nodes. That is its reason to exist.
Do not reach for it when a model fits on one node. If one pod with tensor parallelism across its GPUs serves the model, a plain Deployment (or the vLLM setup from the single-node case) is simpler and has fewer failure modes. LWS earns its complexity only once you are genuinely multi-node.
Serving a model larger than any node
LeaderWorkerSet gives Kubernetes the one thing multi-node inference needs and core objects lack: a group of pods treated as a single replica. It composes StatefulSets for stable identity and parallel startup, injects LWS_LEADER_ADDRESS and LWS_GROUP_SIZE so members find each other, and restarts the whole group on failure so collective communication can re-form. Map your tensor and pipeline parallelism onto size, put the server on the leader, point the Service at leaders, and you can serve a model far larger than any single node.