Tensor Parallelism in oMLX — Complete Visual Guide

Correction — 22 Aug 2026

The 1.78x figure throughout this post is oMLX’s own published number for a 27B model across two Macs over Thunderbolt RDMA — it is not a general property of tensor parallelism and was not reproduced here. The Phi / Gemma / Mistral rows previously read “Partial”; all of them in fact return _supports_tensor_parallel = False. Separately, the error no tensor-parallel degree divides N nodes across architecture dimensions (...) is unreliable — it can fire when those dimensions divide the node count exactly. See MODELS.md and upstream issues jundot/omlx#3021 and #3022.

Complete Visual Guide — Understanding Attention Heads, Model Configuration, and Distributed Inference

1. Overview 2. Attention Architecture 3. Pipeline Parallelism 4. Supported Models 5. Config.json Guide

What is Tensor Parallelism?

Split a model’s layers across multiple nodes so they compute every token together. oMLX reports 1.78x on a 27B across two Macs over Thunderbolt RDMA.

A — Single Node vs. Tensor Parallel

When to use: When your model has many attention heads that divide evenly by the number of nodes, and you have a fast link (Thunderbolt, RDMA).

The Trade-off

✓ Tensor Parallelism

Faster responses. oMLX’s published figure is 1.78x, measured on a 27B over Thunderbolt RDMA — not a general guarantee.

⚠ Requires fast link between nodes and even head division.

Best for speed

→ Pipeline Parallelism

Bigger models: Each node holds different layers, no model size limit.

✓ Works on any link, any model.

Best for capacity

How Attention Heads Work

The dimensions that matter for tensor parallelism, explained visually.

B — Multi-Head Attention Architecture

Key Concept: Tensor parallelism requires all architecture dimensions to divide evenly by the number of nodes. If you have 24 heads and 2 nodes, that works. If you have 23 heads and 2 nodes, it doesn’t.

C — Sharding a Layer Across Two Nodes

Pipeline Parallelism — The Alternative

When tensor parallelism won’t work, split the model’s layers across nodes instead.

D — Pipeline vs. Tensor Parallelism

How Pipeline Parallelism Works

1. Node 1 processes token through layers 1-3

   All attention heads and weights for those layers stay on Node 1

2. Hidden states are sent to Node 2

   Only the intermediate results (small tensors) cross the Thunderbolt link

3. Node 2 processes layers 4-5

   Completes the forward pass

4. One full token generated per iteration

   Takes longer than tensor parallelism, but uses less memory

E — Pipeline Parallelism Timeline

Tensor vs. Pipeline: When to Use Each

✓ Use Tensor Parallelism When:

• Model has many attention heads that divide evenly
• You need fast token latency
• You have a fast link (Thunderbolt, RDMA)
• Qwen3, Qwen3.5, Nemotron-H, Llama 3

✓ Use Pipeline Parallelism When:

• Model doesn’t support tensor sharding (like Qwen3.8)
• You need larger context windows (fits bigger models)
• Any architecture, any link
• Works on slow links, even Wi-Fi

Pipeline stages add latency per token: While tensor parallelism speeds up responses (oMLX reports 1.78x on a 27B over TB RDMA), pipeline parallelism may be slightly slower than a single node. The trade-off is fitting larger models that wouldn’t run on one Mac alone.

Your Qwen3.8 Question: Yes, on 2 Nodes

✓ Can you deploy Qwen3.8-27B on 2 Mac minis?

Yes — using pipeline parallelism.

• Tensor parallelism: Not yet supported (MLX hasn’t added shard() for qwen3_8)
• Pipeline parallelism: Splits layers across nodes ✓ Works immediately
• Speed: Slightly slower than one node, but you can run the full 27B model
• Memory: ~14 GB per Mac (fits 24GB MacBook Pro)

How to Deploy Qwen3.8 on 2 Nodes:

1. Load the model in oMLX dashboard
2. Select “Deploy on 2 Macs”
3. Choose strategy: “Pipeline” (not Automatic, which tries Tensor first)
4. oMLX splits layers: roughly 16 layers on each Mac
5. Done — model runs across both nodes

Which Models Support Tensor vs. Pipeline Parallelism?

Only models with a native MLX shard() implementation support distributed tensor parallelism.

Model Family	Examples	Tensor TP Support	Key Dimensions
Qwen3	qwen3-32b, qwen3-27b	✓ Yes	num_attention_heads, num_key_value_heads
Qwen3-Next	qwen3-next-32b, qwen3-next-7b	✓ Yes	num_attention_heads, num_key_value_heads, linear_num_key_heads, linear_num_value_heads
Qwen3.5	qwen3.5-32b, qwen3.5-7b	✓ Yes	num_attention_heads, num_key_value_heads, linear_num_key_heads, linear_num_value_heads
Qwen3-MoE	qwen3-moe-128b	✓ Yes	num_attention_heads, num_key_value_heads
Nemotron-H	nemotron-h-405b	✓ Yes	num_attention_heads, mamba_num_heads, n_groups
Llama	llama-3-70b, llama-2-7b	✓ Yes	num_attention_heads, num_key_value_heads
Phi	phi-3-mini, phi-4	No	_supports_tensor_parallel returns False (phi3, phi3small, phimoe)
Gemma	gemma-2-27b	No	_supports_tensor_parallel returns False (gemma2, gemma3)
Mistral	mistral-7b	No	_supports_tensor_parallel returns False (mistral, mixtral)

Important: Just because a model isn’t in the “full support” list doesn’t mean it won’t work. MLX can dynamically detect if a model has a shard() method. Check the oMLX dashboard — it will tell you if a specific model supports tensor parallelism when you try to load it.

How to Check a Specific Model

1. Find the model’s config.json — Usually in the model directory on Hugging Face or your local storage.
2. Look for architecture dimensions:
   • num_attention_heads — Primary dimension, must be divisible
   • num_key_value_heads — GQA dimension (optional, usually same or smaller)
   • linear_num_key_heads — Qwen3-Next only
   • mamba_num_heads — Nemotron-H only
3. Divide by node count: If you have 2 nodes, all dimensions must be divisible by 2 with no remainder.

D — Why This Model Failed: 16 ÷ 2 = 8, But 4 ÷ 2 = 2, And 32 ÷ 2 = 16

supports_tensor_parallel=True in config.json ✗ MLX implementation for this model type doesn't have a shard() method ✗ The sharding method isn't layer-safe (can't shard individual pipeline stages) Solutions 1. Use pipeline parallelism (split layers, not tensors) 2. Run on 1 node only (model fits in one Mac) 3. Try a different model (one known to support TP)

Reading Model Config.json Files

The configuration file that tells oMLX how a model is structured and whether it can be parallelized.

Essential Fields for Tensor Parallelism

{ “model_type”: “qwen3”, “hidden_size”: 3584, “num_attention_heads”: 24, // Primary dimension for sharding “num_key_value_heads”: 8, // GQA dimension (optional) “num_hidden_layers”: 32, “intermediate_size”: 18944, “vocab_size”: 128000, “max_position_embeddings”: 32768, }

model_type

Architecture name: qwen3, llama, nemotron_h, etc. oMLX uses this to find the MLX implementation and check for shard().

num_attention_heads

Total query attention heads. Must divide evenly by node count. Example: 24 heads ÷ 2 nodes = 12 per node ✓

num_key_value_heads

Grouped Query Attention: Fewer KV heads than query heads to save memory. Must also divide evenly.

hidden_size

Total embedding dimension. Usually num_attention_heads × head_size. Example: 24 × 149.33 ≈ 3584

Extended Fields (Model-Specific)

{ “linear_num_key_heads”: 16, // Qwen3-Next: linear attention dimension “linear_num_value_heads”: 4, // Must also divide evenly “mamba_num_heads”: 4, // Nemotron-H: Mamba layer heads “n_groups”: 32, // Nemotron-H: group dimension }

Qwen3-Next & Qwen3.5: These models combine standard attention with linear attention. All four dimensions must divide evenly.

Nemotron-H: Hybrid Mamba + Attention architecture. Uses mamba_num_heads and n_groups instead of traditional heads.

Quick Check: Can This Model Use Tensor Parallelism?

1. Open the model’s config.json

   Find it in your model directory or on Hugging Face

2. Find the architecture dimensions

   Look for: num_attention_heads, num_key_value_heads, and any linear_* or mamba_* fields

3. Divide by your node count

   Example: 2 nodes → all dimensions must divide by 2 with remainder 0

4. Check the verdict

   All divide evenly → Tensor parallelism possible (if MLX has shard())
   At least one doesn’t divide → Use pipeline parallelism or 1 node

E — Real-World Example: Qwen3-32B Config

Real Model: Qwen3-32B with 24 heads and 8 KV heads works with 2 or 4 nodes, but not 3 nodes. This is exactly why oMLX’s autoconfigure picks the best split automatically.

Model Folder Structure and File Types

What each file in a model directory does and why it matters for oMLX.

model-directory/ ├── config.json ← Architecture and metadata ├── model.safetensors ← Weights (MLX format) ├── model-00001-of-00005.safetensors ├── model-00002-of-00005.safetensors ├── … ├── model.safetensors.index.json ← Weight sharding map ├── tokenizer.model ← Token vocabulary ├── tokenizer.json ├── generation_config.json ← Inference defaults ├── README.md └── .gitattributes

config.json

Architecture blueprint. Defines model structure (attention heads, layers, hidden size). oMLX reads this to determine if tensor parallelism is possible.

Required for: Parallelism detection

model.safetensors

or
model-00001-of-N.safetensors

Model weights. The actual parameters (~32 GB for Qwen3-32B). Stored in MLX native format for fast loading. Large files are sharded into parts (00001, 00002, etc.) to avoid disk overhead.

Required for: Running inference

model.safetensors.index.json

Weight shard map. Lists which weights live in which shard file. Tells oMLX how to load layers progressively without loading the entire model.

Required for: Progressive layer loading

tokenizer.model

and
tokenizer.json

Vocabulary and encoding. Converts text to tokens (numbers). Model-specific; not interchangeable across architectures.

Required for: Text input processing

generation_config.json

Inference defaults. Settings like max tokens, temperature, top_p. oMLX uses these as fallbacks if you don’t override them.

Optional: Defaults if not specified

README.md

Model documentation. License, usage notes, training data, known limitations. Always read this before deployment.

Important for: Legal and operational context

F — How oMLX Uses Model Files

Why the safetensors format? Modern models use safetensors (safe, fast, and designed for MLX). Older models might use `.bin` (PyTorch) or `.gguf` (GGML), but oMLX converts to safetensors on load. Always prefer safetensors — they load 10-20x faster.

Quick File Checklist Before Running a Model

✓ Required Files (Model Won’t Load Without These)

• config.json — Architecture definition
• model*.safetensors (or all shards) — Weights
• tokenizer.model or tokenizer.json — Token encoding

⚠ Important (Deployment Will Fail Without Proper Values)

• model.safetensors.index.json (if weights are sharded) — Weight location map
• All shard files present if index.json references them

✓ Optional (Nice to Have)

• generation_config.json — Inference defaults
• README.md — Documentation