Model parallelism: how Kimi K3 distributes trillions of parameters

⏱️ Reading time: 13 min

Kimi K3 has 3 trillion parameters. Stored in 4-bit precision, that weighs 1.5 terabytes, far more than the 44 GB of memory a single Cerebras wafer-scale chip has. The only way to serve a model like this is model parallelism: splitting it across dozens of chips at once.

📑 En este artículo

1. TL;DR
2. What is model parallelism?
3. Why it matters
4. How weight sharding and attention work across chips
   1. Pipeline: each chip runs a stretch of layers
   2. Tensor: splitting a single layer into pieces
   3. Experts: activating only a fraction of the model
5. The Kimi K3 case on Cerebras WSE-3T chips
6. Getting started: testing weight sharding with Hugging Face Accelerate
7. Real-world use cases
8. Common mistakes and best practices
9. Comparison with alternatives
10. Going deeper
11. Frequently Asked Questions
    1. Does model parallelism replace data parallelism?
    2. Why does Kimi K3 need 59 Cerebras WSE-3T wafers?
    3. Does Hugging Face’s device_map=”auto” do real weight sharding?
    4. Does pipeline parallelism make a chip wait for the previous request to finish?
    5. What’s the difference between tensor parallelism and expert parallelism?
12. References

This isn’t exclusive to Cerebras: the same principles apply to GPU clusters running Megatron-LM, DeepSpeed, or vLLM. This guide explains how weights and attention get split across chips, using Kimi K3 on Cerebras wafers as a concrete case, not as a product announcement.

TL;DR

• Kimi K3 splits its 93 layers (69 linear attention and 24 dense) across 59 Cerebras WSE-3T wafers.
• Each Kimi K3 layer weighs about 16 GB in 4-bit precision, so no single 44 GB chip can hold the entire model.
• Dense layers reserve memory for the KV cache; linear layers don’t, which is why two fit per wafer instead of one.
• Hugging Face’s device_map=”auto” applies the same sharding principle on a home GPU, without a wafer-scale cluster.
• A chip pipeline keeps several requests in flight at once, so no wafer sits idle waiting for its turn.

What is model parallelism?

Model parallelism is the set of techniques that splits a neural network’s weights and computations across multiple chips when the model doesn’t fit in a single accelerator’s memory, dividing layers, portions of each matrix, or groups of experts depending on the available hardware architecture.

There are three main variants, and large systems almost always combine more than one: pipeline parallelism (each chip runs a different group of layers), tensor parallelism (each chip runs a portion of the same layer), and expert parallelism (each chip stores a subset of specialized networks in Mixture of Experts models). Kimi K3 on Cerebras mostly uses the first one.

Why it matters

Model parallelism isn’t an optional optimization, it’s a necessary condition for the model to load at all. A 3-trillion-parameter model in 4-bit precision weighs 1.5 TB, and the largest GPU on the market has tens of gigabytes of memory, not terabytes. Without splitting the model, it simply doesn’t fit on the chip.

The same happens at a smaller scale. A 70-billion-parameter model at full precision already weighs about 280 GB, more than the combined memory of several consumer GPUs. That’s why any team serving a large model, whether with Cerebras wafers or a GPU cluster, ends up solving the same distribution problem.

How weight sharding and attention work across chips

Before splitting anything, you need to know how much each piece weighs. Cerebras’s compiler maps each model layer to a physical wafer region, and the runtime decides the path tokens take between those regions. The following diagram summarizes how Kimi K3’s 93 layers end up distributed across 44 GB wafers:

flowchart TD
    A["69 linear attention layers"] --> B["2 linear layers per wafer (32 GB)"]
    C["24 dense layers"] --> D["1 dense layer + KV cache per wafer (44 GB)"]
    B --> E["35 wafers for linear layers"]
    D --> F["24 wafers for dense layers"]
    E --> G["59 WSE-3T wafers total"]
    F --> G

Pipeline: each chip runs a stretch of layers

In pipeline parallelism, each chip holds a group of consecutive layers and acts like a station on an assembly line. Wafer 1 runs layers 1 and 2 and sends its activations to wafer 2, which runs layers 3 and 4 and passes them to wafer 3, and so on through the last stage. Only activations cross from one wafer to another, never the weights, because each entire layer already fits in the memory of the wafer that hosts it.

In Cerebras’s case, those activations travel over RDMA over Ethernet, with a latency of about 2 microseconds per hop. The following diagram shows how the pipeline fills up with several requests at the same time:

sequenceDiagram
    participant W1 as Wafer 1
    participant W2 as Wafer 2
    participant W3 as Wafer 3
    W1->>W2: activations for request A
    W2->>W3: activations for request A
    W1->>W2: activations for request B
    W2->>W3: activations for request B
    Note over W1,W3: each wafer processes a different request at the same time

💭 Key point: a wafer doesn’t wait for the entire request to finish: it processes a different portion of each request at the same time as its neighbors, so the pipeline never runs at batch size 1.

Tensor: splitting a single layer into pieces

Tensor parallelism splits a single weight matrix into fragments, and each chip computes its fragment of the multiplication. At the end, the partial results need to be combined with a synchronization operation called all-reduce, which requires a very fast network between chips (usually NVLink). This is the approach Megatron-LM uses to train GPT-style models on GPU clusters.

Experts: activating only a fraction of the model

In Mixture of Experts models, each token gets routed to a subset of specialized networks (experts), and each chip can store just a few experts instead of a full copy of the model. Network traffic depends on which experts each token gets routed to, so it isn’t as predictable as in a fixed pipeline.

The Kimi K3 case on Cerebras WSE-3T chips

Kimi K3 has 69 linear attention layers and 24 dense layers, 93 total. If the 3 trillion parameters are stored in 4-bit precision, the raw weight is 1.5 TB; split across 93 layers, each one weighs around 16 GB, well under the 44 GB each wafer provides.

That’s where the difference between layer types shows up. A dense layer needs extra space for the KV cache, because it stores the keys and values for every token of every active request; that’s why only one dense layer fits per wafer, leaving the rest of the memory for that cache. A linear attention layer, on the other hand, holds a fixed-size state per request regardless of how much context it has, so it doesn’t need a KV cache and two fit per wafer (2 x 16 GB = 32 GB out of the 44 GB available).

With 69 linear layers at two per wafer, you need 35 wafers (34.5 rounded up), and with 24 dense layers at one per wafer, you need 24 more wafers. The total is 59 WSE-3T wafers. Since each entire layer fits on a single wafer, the weights are never split across chips: only the activations cross from wafer to wafer.

According to the same Cerebras post, Kimi K3 scores 88.3% on Terminal-Bench 2.1 (versus 88.8% for GPT-5.6 Sol) and 91.2% on BrowseComp (versus 90.4% for GPT-5.6 Sol), results that put it on par with proprietary models on coding and tool-use tasks, despite being an open-weight model.

Getting started: testing weight sharding with Hugging Face Accelerate

Nobody has Cerebras wafers on their desktop, but the same layer-splitting principle can be tested with open tools. Before running anything, install the dependencies (the command is identical on Windows, macOS, or Linux):

pip install transformers accelerate torch

First, load a small model without telling it where to live:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
print(model.device)

Without device_map or a call to .to(), Transformers loads the weights into the process’s RAM. The output is:

cpu

Now let Accelerate decide the split for you:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-1.5B-Instruct",
    device_map="auto",
)
print(model.hf_device_map)

With a small model that fits entirely on a single GPU, Accelerate doesn’t split anything and the output is simply:

{'': 0}

That empty dictionary means “the entire model on device 0”. With a model that doesn’t fit on a single GPU, the same hf_device_map attribute returns a layer-by-layer mapping, something like {'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.14': 1, 'lm_head': 1}: that’s where you actually see weight sharding in action, with each block of layers assigned to a different GPU.

💡 Tip: always print model.hf_device_map after loading with device_map="auto": it’s the only direct way to confirm whether Accelerate split the model or left it whole on one device.

Real-world use cases

• Cerebras + Kimi K3: production inference with a 59-wafer pipeline, each wafer with a fixed group of layers and memory dedicated to the KV cache.
• Training with Megatron-LM: NVIDIA combines tensor, pipeline, and data parallelism in what’s known as 3D parallelism, to train GPT-style models on clusters of thousands of GPUs.
• DeepSpeed ZeRO: instead of splitting layers, ZeRO-3 splits optimizer states and gradients across GPUs, another way to solve the same memory problem during training.
• Serving with vLLM: the vLLM engine exposes the --tensor-parallel-size flag so a large model can be split across several GPUs when serving requests, without touching the model’s code.

Common mistakes and best practices

• Confusing model parallelism with data parallelism. The latter copies the entire model onto each chip to handle more requests in parallel; the former splits pieces of a single model that’s too large for one chip. They’re complementary, not alternatives.
• Not measuring the actual size per layer. A dense layer with a KV cache weighs differently than a linear attention layer without one; planning the split without that distinction overestimates how many chips fit per wafer.
• A poorly fed pipeline. If there’s only one request in flight, each chip sits idle waiting for its turn (the so-called pipeline bubble). The solution is to keep several requests or microbatches running at once.
• Ignoring network latency between chips. 2 microseconds per hop sounds insignificant, but added up across many hops in a long pipeline, it does weigh on a request’s total latency.
• Assuming device_map=”auto” always splits something. With a small model, it doesn’t do real weight sharding: it only kicks in when the model doesn’t fit on one device.

Comparison with alternatives

No single technique, pipeline, tensor, or expert parallelism, always wins; each solves a different bottleneck of the available hardware.

Going deeper

The number of concurrent requests a wafer with a dense layer can support depends on the available KV cache. If the wafer has 44 GB and the dense layer uses 16 GB, about 28 GB remain free for the KV cache. The general formula is available cache memory divided by the KV cache size of a single request (which depends on context length, the number of attention heads, and cache precision, 8-bit in Kimi K3’s case). There’s no way to confirm Kimi K3’s exact number of concurrent requests without running the actual model: the available excerpt from the Cerebras post describes the method but doesn’t publish the final number.

The number of in-flight requests also depends on how they’re grouped into microbatches. According to Cerebras, the number of in-flight requests is roughly the number of pipeline stages multiplied by the microbatch size: small microbatches lower each individual request’s latency, large microbatches raise the system’s total throughput.

The following diagram contrasts how pipeline parallelism is structured versus tensor parallelism:

flowchart LR
    subgraph Pipeline["Pipeline parallelism"]
    P1["Layers 1 and 2 on Chip A"] --> P2["Layers 3 and 4 on Chip B"]
    end
    subgraph Tensor["Tensor parallelism"]
    T1["Half the matrix on Chip A"] --> T3["Final all-reduce"]
    T2["Half the matrix on Chip B"] --> T3
    end

Tensor parallelism requires an extremely fast network between chips, because every layer ends with a synchronization. Pipeline parallelism tolerates slower networks, but adds per-request latency, because every token has to pass through all stages before coming out. The choice depends on the actual bottleneck of the available hardware, no single option always wins; that’s why all three model parallelism variants coexist in the same cluster when the model and the hardware call for it.

Your next step: install transformers and accelerate, load an open model with device_map=”auto”, and print model.hf_device_map to see with your own eyes how Accelerate decides the split on your hardware.

Frequently Asked Questions

Does model parallelism replace data parallelism?

No, they’re complementary. Data parallelism copies the entire model onto each chip to handle more requests in parallel; model parallelism splits pieces of a single model that’s too large for one chip. Large systems like Megatron-LM combine both.

Why does Kimi K3 need 59 Cerebras WSE-3T wafers?

Because its 93 layers don’t fit on a single 44 GB wafer. Each layer weighs about 16 GB in 4-bit precision, so the split requires 35 wafers for the 69 linear attention layers and 24 more wafers for the dense layers.

Does Hugging Face’s device_map=”auto” do real weight sharding?

Yes, it applies the same principle at a smaller scale: it distributes the model’s layers across available GPUs when the model doesn’t fit on one. With a small model, on the other hand, it doesn’t split anything.

Does pipeline parallelism make a chip wait for the previous request to finish?

No. Each pipeline stage processes a different request at the same time as its neighbors, so a well-fed pipeline never runs at batch size 1.

What’s the difference between tensor parallelism and expert parallelism?

Tensor parallelism splits a single dense layer into pieces that run in parallel with a final synchronization (all-reduce). Expert parallelism activates only a subset of specialized networks per token, without touching the rest of the model.

References

• Cerebras: official post detailing how they split Kimi K3 across WSE-3T wafers.
• Hugging Face Accelerate: documentation for device_map=”auto” and automatic layer distribution across devices.
• NVIDIA Megatron-LM: repository with the reference implementation of combined tensor and pipeline parallelism.
• vLLM: documentation for the inference engine that exposes –tensor-parallel-size to serve large models in production.

📱 Do you like this content? Follow @programacion on Telegram for daily tech content in Spanish: quick summaries, fresh content every day.

Featured image: Foto de Laura Ockel en Unsplash