How to Choose AI Hardware for Training, Fine-tuning and Inference

The idea in one minute#

Choosing AI hardware is three questions in order:

What stage is this? It tells you which resource binds (lesson 01).
What size is it? It tells you which level of scale you need (lesson 02).
What can you actually run and buy? Software support and available capacity eliminate most options before price is even considered (Beyond NVIDIA, lesson 01).

Only then compare price per unit of your work. This lesson turns the topic into a decision guide and works through four real situations. Checked on 16 October 2026.

A picture#

A first cut, not an answer. Every branch ends in the same three checks at the bottom.

How it really works#

The guide, branch by branch#

Pretraining. Buy the largest fast domain you can get and a network that will not stall it: NVL72 racks (GB300 or Vera Rubin), TPU pods or Trainium3 UltraServers, on a rail-optimised 800G-1.6T fabric with fast checkpoint storage. Between these, the deciding factors are your team’s software (CUDA and PyTorch, or JAX and XLA) and how much capacity a provider can actually commit to on your dates.

Post-training and RL. The same cluster, divided into a trainer pool and a rollout pool. Size the rollout pool like a serving fleet, because generation dominates RL time. Keep weight synchronisation from trainers to generators on the fastest links you have.

Fine-tuning. Capacity decides. Up to ~8B parameters, QLoRA on one 24-48 GB GPU (RTX 5090, L40S, or a DGX Spark if slow is acceptable). From 8B to 70B, one to eight H100/H200/B200-class GPUs, or an RTX PRO 6000’s 96 GB. Full fine-tuning of 70B and above needs a node with FSDP.

Serving. Start with fit: if the model fits one GPU at FP8 with room for its KV cache, serve on single GPUs (H100, H200, B200, MI355X) or the equivalent in your cloud (TPU, Inferentia, Trainium). Large MoE models (≥ 400B parameters) want wide expert parallelism: NVL72, Helios or a TPU slice. If per-user speed is what you sell, add an SRAM decode tier (Groq LPX, Cerebras) and keep GPUs for prefill.

Edge. The device decides: phone and PC NPUs, Jetson Thor for robots, DGX Spark, a Mac Studio or an RTX PRO workstation for a private desk.

Four worked examples#

1. A startup fine-tunes an 8B model on 50,000 support tickets. Stage: fine-tuning; binding resource: capacity. QLoRA on an 8B model needs about 10-16 GB, so one 24-32 GB card does it in hours. Renting one H100 for a day costs less than a week of engineering time; buying a 5090 makes sense only if this repeats weekly. Do not rent a node.

2. A bank serves a 70B model at FP8 to 2,000 employees, 8k context, in its own data center. Stage: serving; binding: capacity, then bandwidth. Weights 70 GB; KV cache at 163,840 bytes per token × 8k × (say) 200 concurrent users ≈ 270 GB. One 8-GPU H200 server (1.1 TB) holds this with room to spare and air cooling suffices; two give redundancy. An NVL72 would be wasted. If budget is tight, RTX PRO 6000 Server Edition cards (96 GB each) in four-card tensor-parallel groups are an option, at lower decode speed.

3. A lab runs RL post-training on a 1T-parameter MoE with long reasoning rollouts. Stage: post- training; binding: generation bandwidth, then expert all-to-all. Rollouts on NVL72 racks with wide expert parallelism; trainers on more NVL72 racks; weight sync over the scale-out fabric after every step. Budget most of the GPUs for rollouts. A TPU-pod equivalent works if the code base is JAX.

4. A robotics company puts a vision-language-action model on a warehouse robot. Stage: edge; binding: power and capacity. A Jetson AGX Thor (128 GB, 40-130 W) runs a multi-billion-parameter model at FP4 on the robot; training happens in the cloud on the same CUDA stack, which is half of why the choice is easy.

The three checks, always#

Regime. Measure whether the workload is compute-, bandwidth- or capacity-bound on your current hardware (4.a). The matching lever predicts your speed-up; the others do not.
Run it. Run your exact model, precision and serving or training stack on the candidate. Count the days it took to make it work.
Unit cost. Convert to dollars per million tokens, per training step or per fine-tune, at the batch size you will really run, including power if you own the hardware.

Code#

The guide as a function: describe the job, get a first-cut recommendation and the reason.

// choose.go — the decision guide as code.
package main

import "fmt"

type Job struct {
	Name          string
	Stage         string  // pretrain | rl | finetune | serve | edge
	ParamsB       float64 // total parameters, billions
	MoE           bool
	BytesPerParam float64 // serving precision
	KVGB          float64 // KV cache budget for serving
	LatencyIsKey  bool    // per-user tokens/s is the product
	Device        string  // for edge: phone | robot | desk
}

func choose(j Job) (string, string) {
	switch j.Stage {
	case "pretrain":
		return "NVL72 racks, TPU pods or Trainium3 UltraServers on an 800G+ rail fabric",
			"bound by FLOPs, network and uptime; pick by software and committed capacity"
	case "rl":
		return "the training cluster split into trainer and rollout pools",
			"rollouts are decode: size them like a serving fleet"
	case "finetune":
		switch {
		case j.ParamsB <= 8:
			return "one 24-48 GB GPU with QLoRA (RTX 5090, L40S)", "capacity-bound; 4-bit base fits"
		case j.ParamsB <= 70:
			return "1-8 H100/H200/B200, or one RTX PRO 6000 with QLoRA", "capacity-bound"
		default:
			return "an 8-GPU node or more with FSDP", "16 bytes/param of optimizer state"
		}
	case "serve":
		need := j.ParamsB*j.BytesPerParam + j.KVGB
		switch {
		case j.MoE && j.ParamsB >= 400:
			return "NVL72, Helios or a TPU slice with wide expert parallelism", fmt.Sprintf("%.0f GB and an all-to-all every layer", need)
		case need <= 0.85*141:
			rec := "single GPUs: H100/H200/B200/MI355X, or TPU/Inferentia in that cloud"
			if j.LatencyIsKey {
				rec += ", plus an SRAM decode tier"
			}
			return rec, fmt.Sprintf("%.0f GB fits one GPU", need)
		default:
			return "an 8-GPU node with tensor parallelism", fmt.Sprintf("%.0f GB needs several GPUs", need)
		}
	case "edge":
		return map[string]string{"phone": "the device NPU", "robot": "Jetson AGX Thor", "desk": "DGX Spark, Mac Studio or an RTX PRO workstation"}[j.Device],
			"bound by power and capacity"
	}
	return "unknown stage", ""
}

func main() {
	jobs := []Job{
		{Name: "support-ticket fine-tune", Stage: "finetune", ParamsB: 8},
		{Name: "bank assistant, 70B FP8", Stage: "serve", ParamsB: 70, BytesPerParam: 1, KVGB: 270},
		{Name: "coding agent, 32B FP8", Stage: "serve", ParamsB: 32, BytesPerParam: 1, KVGB: 40, LatencyIsKey: true},
		{Name: "1T MoE serving", Stage: "serve", ParamsB: 1000, MoE: true, BytesPerParam: 0.56, KVGB: 1500},
		{Name: "reasoning RL", Stage: "rl", ParamsB: 1000, MoE: true},
		{Name: "frontier pretraining", Stage: "pretrain", ParamsB: 2000, MoE: true},
		{Name: "warehouse robot", Stage: "edge", Device: "robot"},
	}
	for _, j := range jobs {
		rec, why := choose(j)
		fmt.Printf("%-26s → %s\n%28s(%s)\n", j.Name, rec, "", why)
	}
	fmt.Println("\nThen: find the regime, run the exact model on the candidate, compare cost per unit of work.")
}

The function is deliberately crude: real decisions also weigh price, existing contracts, data location and team skills. Its value is the order of the questions, which is the same at any budget.

Remember this#

Ask in order: stage (binding resource), size (level of scale), feasibility (software and capacity), then price.
Pretraining: biggest fast domain plus network; RL: split trainer and rollout pools; fine-tuning: capacity; serving: fit first, then MoE width, then latency; edge: the device.
Always finish with the three checks: regime, run it, unit cost.

Try it#

Run choose.go. Add a job of your own and check whether the recommendation matches what you would have bought.
Change the bank example’s KV budget to 1 TB (long documents). What does the guide say now, and do you agree?
Extend choose with a Cloud field so that serving inside Google Cloud or AWS can recommend TPUs or Trainium first.

Check yourself#

What three questions come before price?
Why does RL post-training need an inference-shaped rollout pool?
In worked example 2, why would an NVL72 be wasted?
What are the three checks to run before signing for any hardware?

Sources#

The lessons of this topic, the NVIDIA hardware line and Beyond NVIDIA, and their sources
vLLM supported hardware
QLoRA (Dettmers et al., 2023)