TPU Architecture — Types, Topology, Multi-host Model

Tại sao TPU khác biệt so với GPU

GPU và TPU cùng là accelerators, nhưng scheduling model hoàn toàn khác:

GPU scheduling:
  "Pod yêu cầu nvidia.com/gpu: 1"
  → Kubernetes allocate 1 GPU (từ 2-8 GPU trên node)
  → Pod runs trên GPU đó

TPU scheduling:
  "Pod yêu cầu tpu slice 4x4x4 topology (64 chips)"
  → Kubernetes must allocate **toàn bộ 64 chips đồng thời**
  → Pod không chạy tới khi toàn bộ topology provisioned
  → If 1 chip fails → entire topology considered failed

TPU là "supercomputer-like" — các chip phải hoạt động nguyên trạng, không thể incrementally thêm chips vào training.

---

TPU Types: Generations và characteristics

TPU v4

Hardware:

• 8 chips per TPU VM
• Each chip: 16GB HBM2 (High Bandwidth Memory)
• Topology: customizable (4x4x4, 2x2x2, 1x1x8, etc.)
• Peak throughput: ~420 TFLOPS (float32)

Mental model: "Specialized for transformer training". Excellent matrix-matrix ops, weak on irregular/sparse operations.

TPU v5p

Hardware:

• 8 chips per TPU VM (same as v4)
• Each chip: 16GB HBM2e (improved memory bandwidth)
• Topology: same as v4
• Peak throughput: ~596 TFLOPS (float32) — ~40% faster than v4

When to use: Larger models (70B+ parameters), requiring improved memory bandwidth.

TPU v6e (Lite)

Hardware:

• Variable: 1, 4, or 8 chips per TPU VM (flexible)
• Each chip: 16GB HBM2e
• Topology: 2x2x2 (single VM) or larger multi-host
• Peak throughput: improved compute density

Advantage: Can run single-host topologies (1 TPU VM = training ready), vs v4/v5p requiring multi-VM for meaningful topologies.

TPU Ironwood (TPU7x)

Hardware (newest, 2026):

• 8 chips per TPU VM
• Each chip: 24GB HBM3 (highest bandwidth + density)
• Topology: 4x4x4 and larger (8x8x8 multi-host)
• Peak throughput: ~1.5x faster than v6e

When to use: State-of-the-art training, inference at scale. Latest models (Gemini, future LLMs) optimized for Ironwood.

GKE Support: Ironwood available in GKE Standard v1.34.0+ and Autopilot v1.34.1+.

---

3D Topology: The Constraint Model

Why topology matters

TPU chips không tất cả connected tương nhau đều. Chúng được arranged trong 3D logical grid:

Topology: 4x4x4 (64 chips)

Physical arrangement (ICI = Inter-Chip Interconnect):
  
  Dimension X (4 chips)  Dimension Y (4 chips)  Dimension Z (4 chips)
  
  Chip[0,0,0]           ← same X, Y, diff Z
      ↓ (Z-link)
  Chip[0,0,1]
      ↓
  Chip[0,0,2]
      ↓
  Chip[0,0,3]
  
  + Chip[0,1,0], Chip[0,2,0], ... (Y-links)
  + Chip[1,0,0], Chip[2,0,0], ... (X-links)

ICI bandwidth:

• Intra-dimension (X, Y, Z): ~600GB/sec per link
• Cross-dimension: routed through intermediate chips

Why this matters for training:

Model parallelism strategy phải respect topology. Nếu chia model theo X-dimension:

Layers 0-16: Chip[0,*,*] (X=0)
Layers 17-32: Chip[1,*,*] (X=1)
→ All-reduce giữa X=0 và X=1: high-bandwidth link

Nếu chia không align:

Layers 0-16: Chip[0,0,0], Chip[1,0,0], Chip[2,0,0], Chip[3,0,0] (scattered across X)
→ All-reduce requires routing through many hops
→ Bandwidth-limited, training slows

Topology selection constraints

Topology must be AxBxC where:

For v4/v5p/Ironwood:
- A ≤ B ≤ C
- All values either ≤4 OR all integer multiples of 4
- Total chips: A×B×C ≤ max_chips_available

Valid:        Invalid:
1x1x8 ✓       1x2x8 ✗ (mixed ≤4 and non-multiple)
2x2x8 ✓       2x3x8 ✗ (3 is neither ≤4 nor multiple of 4)
4x4x4 ✓       4x5x4 ✗ (5 invalid)
4x8x8 ✓       3x3x3 ✓ (all ≤4)

---

Multi-host vs Single-host TPU Slices

Single-host TPU slice

Definition: 1 TPU VM = 1 TPU slice

1 × TPU v5p VM:
- 8 TPU chips
- Topology: up to 2x2x2 (8 chips), 1x1x8, etc.
- Communication: all intra-VM (no network needed)

GKE node pool creation:

gcloud container node-pools create tpu-single-pool \
  --cluster=my-cluster \
  --machine-type=tpu-vm \
  --accelerator=type=tpu-v5p-single,count=1 \
  --num-nodes=5

Characteristics:

• Simple scaling: add more nodes = more independent training jobs
• No topology awareness needed (entire cluster is 1x1x8 per node)
• No data center interconnect (ICI) — intra-node only

Use case: Small models (7B-30B parameters), batch training where model fits in single TPU.

Multi-host TPU slice

Definition: N TPU VMs = 1 TPU slice (forming single large topology)

8 × TPU v5p VMs:
- 64 TPU chips total (8 chips/VM × 8 VMs)
- Topology: 4x4x4 (coordinated across all 8 VMs)
- Communication: intra-VM (fast) + inter-VM ICI (datacenter network)

GKE node pool creation:

gcloud container node-pools create tpu-multi-pool \
  --cluster=my-cluster \
  --machine-type=tpu-vm \
  --accelerator=type=tpu-v5p-multi,count=8,topology=4x4x4 \
  --num-nodes=8  # 8 TPU VMs

Critical constraint: Once TPU slice node pool created, cannot add nodes to it:

Node pool created with 8 nodes (topology 4x4x4 fixed)
→ Cluster admin wants to scale to 16 nodes (topology 8x8x8)
→ CANNOT: must delete pool + recreate with correct size

# Why? Topology is baked into node pool. Kubernetes scheduler
# assumes TPU topology = exactly what was configured.

Scaling behavior:

Node pool: min=0, max=8 (fixed 8 VMs)

Cluster autoscaler logic:
  if Pod pending for 3 minutes:
    if all 8 nodes running: wait (cannot scale up)
    if <8 nodes: provision 1 more node
    
  When 8 nodes + all running:
    Pod still pending → stuck forever (autoscaler can't help)
    → must manually delete low-priority Pod or recreate node pool

Atomic scheduling: all-or-nothing

Multi-host TPU slice must be scheduled as atomic unit:

Pod requesting tpu: 4x4x4 slice

Kubernetes scheduler checks:
- Do all 8 TPU nodes have resources?
  - Node 0: available ✓
  - Node 1: available ✓
  - ...
  - Node 7: available ✓
  
If any node unavailable → Pod pending

Not partial: "let's schedule on 7 nodes + wait for 1" ← not allowed

Implication for multi-tenant clusters:

Cluster có 2 multi-host TPU node pools:
- Pool A: 4x4x4 (8 nodes, team-a)
- Pool B: 4x4x4 (8 nodes, team-b)

If team-a Pod pending for long time:
→ Cluster autoscaler can't scale Pool A (already max)
→ But Pool B nodes idle
→ Wasted resources (no work-stealing between pools)

→ Solution: use single Pool + NodeSelector/Affinity ← complex

ICI (Inter-Chip Interconnect) topology

Multi-host TPU slices connected via ICI — Google's proprietary datacenter network (different from regular Ethernet):

TPU VM 0              TPU VM 1             TPU VM 2
[8 chips]       →[ICI]→  [8 chips]  →[ICI]→  [8 chips]
   ↓ (Z-link)                                 ↓
 Chip[0,0,0]                              Chip[3,0,0]
 (all-reduce happens across ICI)

ICI bandwidth: ~600GB/sec per link (between VMs).

Network requirements:

• Standard GCP VPC Ethernet: ~100GB/sec max
• Must use ICI-enabled network topology
• VMs cannot be in different zones (ICI latency breaks)
• Dataplane V2 recommended (better hardware offload)

Implication: Multi-host TPU slice nodes must be zone-local or region-local (same zone preferred).

---

TPU-to-Pod mapping rules

Container-to-TPU binding

Hard rule: "Any container requesting TPUs must consume ALL TPU chips in the node".

# ❌ INVALID
spec:
  containers:
  - name: inference
    resources:
      limits:
        google.com/tpu: 4  # Request 4 chips
  # → REJECTED (must request entire slice)

# ✓ VALID
spec:
  containers:
  - name: training
    resources:
      limits:
        google.com/tpu: 8  # 8 chips = entire TPU VM

Why: TPU drivers + topology initialization is per-node, not per-container. Kubernetes scheduler can't slice TPU within a node.

Pod-to-TPU slice mapping

Multi-host TPU requires special Pod spec:

apiVersion: v1
kind: Pod
metadata:
  name: llm-training
spec:
  # Must NOT run on non-TPU nodes
  nodeSelector:
    cloud.google.com/gke-accelerator: tpu-v5p-multi-4x4x4
  
  # Distributed training framework (e.g., jax.distributed) requires
  # coordinated startup on ALL nodes simultaneously
  # GKE handles Pod replication across multi-host slice
  
  containers:
  - name: training
    image: llm-trainer:latest
    resources:
      limits:
        google.com/tpu: 8

Behind the scenes:

• Pod spec: define once
• GKE scheduler: replicate to all 8 nodes
• All 8 Pods start simultaneously (coordinated startup)
• Training framework (JAX, PyTorch) auto-discovers TPU topology from environment

---

Topology-aware training strategies

Data parallelism (no topology awareness needed)

Single-host training on 8 chips (1x1x8 topology):
- Replica 0: runs on Chip[0]
- Replica 1: runs on Chip[1]
- ...
- Replica 7: runs on Chip[7]

All-reduce (gradient sync): local within VM (fast)

Tensor parallelism (topology-aware)

Model split across dimensions:

Topology: 4x4x4 (64 chips total)

Model layers sharded by X-dimension:
- Layers 0-16:  Chip[0,*,*] (all Y, Z)
- Layers 17-32: Chip[1,*,*] (all Y, Z)
- Layers 33-48: Chip[2,*,*] (all Y, Z)
- Layers 49-64: Chip[3,*,*] (all Y, Z)

All-reduce for next layer: between X=i and X=i+1 (high-BW X-link)

Pipeline parallelism (stage mapping)

Topology: 4x4x4

Pipeline stages mapped to slices:
- Stage 0: Chip[0,0:4,0:4] (16 chips)
- Stage 1: Chip[1,0:4,0:4] (16 chips)
- Stage 2: Chip[2,0:4,0:4] (16 chips)
- Stage 3: Chip[3,0:4,0:4] (16 chips)

Stage-to-stage communication: along X-dimension (1 hop)

---

Failure modes and recovery

Single chip failure in multi-host slice

TPU slice 4x4x4 (64 chips), Chip[2,1,3] fails

From GKE perspective:
→ Node running Chip[2,1,3] becomes unhealthy
→ Kubelet marks node NotReady
→ All Pods on that node → Evicted/Pending

Training framework (JAX):
→ Loses Chip[2,1,3] out of 64
→ Cannot continue (collective ops blocked)
→ Training job must restart from checkpoint

Recovery: Manual intervention or automatic job restart

Node eviction (preemption, maintenance)

Node 0 (containing 8 TPU v5p chips) scheduled for maintenance

GKE behavior:
1. Pod disruption budget: if none specified, allow eviction
2. Evict Pod on Node 0
3. Training job interrupted (no crash recovery in TPU alloc)
4. New Pod pending on full 8-node slice

Scheduling: must wait for all 8 nodes to be available again
→ If Node 5 is running other workload, entire training blocked

Mitigation: Use Spot TPU (preemptible, cheaper), with checkpoint/resume logic.

Topology mismatch

GKE node pool: topology 4x4x4 configured
Training job: requests topology 2x2x2

→ Kubernetes scheduler:
   - Cannot down-size topology
   - Pod pending forever

→ Fix: either update node pool topology, or resubmit job for 4x4x4

---

GKE TPU management specifics

Autopilot vs Standard

Standard:

• User create node pool, specify topology (e.g., 4x4x4)
• User responsible for autoscaling policy
• More control, more operational overhead

Autopilot:

• User specify resource request: google.com/tpu: 8
• GKE auto-selects TPU type + topology
• Simpler, but less control

Node pool scaling for multi-host

# Scale from 8 nodes to 16 nodes
gcloud container node-pools scale tpu-pool \
  --num-nodes=16 \
  --cluster=my-cluster

# For 4x4x4 topology (8 nodes):
#   min_nodes=8, max_nodes=8, target=8 (FIXED)
#   Scale command → no-op OR error

# To change topology, must recreate:
gcloud container node-pools delete tpu-pool --cluster=my-cluster
gcloud container node-pools create tpu-pool \
  --cluster=my-cluster \
  --accelerator=type=tpu-v5p-multi,count=16,topology=8x8x8 \
  --num-nodes=16

GKE-TPU device plugin

Similar to GPU plugin, but:

• Exposes google.com/tpu (not nvidia.com/gpu)
• No fractional allocation (request = all chips in slice)
• Topology validation (rejects invalid topologies)

---

Mental model: TPU as "owned supercomputer"

Key insight:

GPU cluster: "bag of GPU compute" — can dynamically
            distribute single GPU to 10 different Pods

TPU slice:  "dedicated supercomputer" — if you request it,
            you own it (all chips, all nodes) until Pod terminates

Scheduling implication:

GPU cluster: high utilization (multiplexing), but queuing delays
TPU cluster: lower utilization (dedicated slices), predictable latency

---

References

• GKE About TPUs
• GKE Plan TPUs
• GKE Deploy TPU Workloads in Standard
• GKE Deploy TPU Workloads in Autopilot
• TPU v4 Documentation
• TPU v6e Documentation
• TPU Ironwood (TPU7x)
• JAX Multi-host Training Guide