Choosing the Right Hardware: GPUs, TPUs, and Cloud Resources

Why Hardware Matters for LLMs

Large language models are computationally intensive, requiring massive parallel processing to train efficiently. The hardware you choose directly impacts training time, cost, and model performance. For example, training a 7-billion-parameter model on a single CPU could take years, while a modern GPU or TPU can reduce this to days or weeks. The key factors to consider are memory bandwidth, parallel processing capabilities, and cost efficiency. Memory bandwidth determines how quickly data can be fed to the processor, while parallel processing allows the model to handle thousands of operations simultaneously. Cost efficiency ensures you’re not overspending on hardware that doesn’t align with your project’s scale. By understanding these principles, you can reason about hardware choices even for models or use cases not explicitly covered here, such as edge deployment or distributed training.

# Example: Checking GPU memory bandwidth (NVIDIA GPU)
import torch

def check_gpu_bandwidth():
    # Create a large tensor to measure memory bandwidth
    size = 1024 * 1024 * 100  # 100MB
    x = torch.rand(size, device='cuda')
    y = torch.rand(size, device='cuda')
    
    # Warm-up
    for _ in range(10):
        z = x + y
    
    # Measure bandwidth
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    z = x + y
    end_event.record()
    torch.cuda.synchronize()
    
    elapsed_time = start_event.elapsed_time(end_event) / 1000  # Convert to seconds
    bandwidth = (size * 4 * 2) / elapsed_time / (1024 ** 3)  # GB/s (float32 = 4 bytes)
    print(f"GPU Memory Bandwidth: {bandwidth:.2f} GB/s")

check_gpu_bandwidth()

GPUs: The Workhorse of LLM Training

Graphics Processing Units (GPUs) are the most common choice for training LLMs due to their balance of performance, flexibility, and accessibility. GPUs excel at parallel processing, which is essential for the matrix multiplications that dominate LLM training. Modern GPUs, like NVIDIA’s A100 or H100, offer high memory bandwidth and large memory capacities (up to 80GB), making them suitable for models with billions of parameters. They also support mixed-precision training, which reduces memory usage and speeds up training without sacrificing model quality. However, GPUs can be expensive, especially for large-scale training, and their power consumption can add to operational costs. When choosing a GPU, consider the model size, batch size, and whether you’ll use distributed training. For example, a single A100 can handle a 7B-parameter model, but larger models may require multiple GPUs or even GPU clusters.

# Example: Profiling GPU memory usage for a small LLM
import torch
from transformers import AutoModelForCausalLM

def profile_gpu_memory(model_name="gpt2"):
    # Load a small LLM and move it to GPU
    model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
    
    # Check GPU memory usage
    print(f"Model: {model_name}")
    print(f"GPU Memory Allocated: {torch.cuda.memory_allocated() / (1024 ** 3):.2f} GB")
    print(f"GPU Memory Reserved: {torch.cuda.memory_reserved() / (1024 ** 3):.2f} GB")
    
    # Clean up
    del model
    torch.cuda.empty_cache()

profile_gpu_memory()

TPUs: Google’s Specialized Accelerators

Tensor Processing Units (TPUs) are custom-designed accelerators developed by Google for machine learning workloads. Unlike GPUs, which are general-purpose, TPUs are optimized specifically for tensor operations, making them highly efficient for training large models. TPUs offer high memory bandwidth and are designed to work seamlessly with Google’s TensorFlow framework, though they can also be used with PyTorch via libraries like JAX. One of the key advantages of TPUs is their scalability: Google’s TPU pods can scale to thousands of chips, making them ideal for training massive models like PaLM or T5. However, TPUs are less flexible than GPUs and may not support all operations or frameworks out of the box. They are also primarily available through Google Cloud, which can limit their accessibility. When considering TPUs, evaluate whether your model and framework are compatible and whether the cost savings justify the reduced flexibility.

# Example: Running a simple model on TPU (using JAX)
import jax
import jax.numpy as jnp
from jax import random

def run_on_tpu():
    # Check if TPU is available
    try:
        devices = jax.devices("tpu")
        print(f"TPU devices found: {devices}")
    except:
        print("TPU not available, running on CPU")
        devices = jax.devices("cpu")
    
    # Create a simple matrix multiplication
    key = random.PRNGKey(0)
    x = random.normal(key, (1024, 1024))
    y = random.normal(key, (1024, 1024))
    
    # Run on TPU/CPU
    z = jnp.dot(x, y)
    print(f"Result shape: {z.shape}")
    print(f"Device used: {z.device()}")

run_on_tpu()

Cloud vs. On-Premises: Cost and Flexibility

Choosing between cloud and on-premises hardware depends on your budget, scalability needs, and long-term goals. Cloud resources, like AWS, Google Cloud, or Azure, offer flexibility and scalability, allowing you to rent hardware on-demand and avoid upfront costs. This is ideal for small teams or projects with variable workloads, as you can scale up or down as needed. However, cloud costs can add up quickly, especially for long-term training jobs. On-premises hardware, on the other hand, requires a significant upfront investment but can be more cost-effective for large-scale or long-term projects. It also gives you full control over your infrastructure, which can be important for security or compliance reasons. When deciding, consider factors like cost per hour, data transfer costs, and whether you need specialized hardware like TPUs. For example, training a 13B-parameter model on the cloud might cost thousands of dollars per week, while purchasing GPUs upfront could save money in the long run.

# Example: Estimating cloud costs for LLM training (AWS)
import math

def estimate_cloud_costs(model_size_billion=7, training_hours=100, gpu_type="p4d.24xlarge"):
    # AWS pricing (as of 2023, us-east-1)
    pricing = {
        "p4d.24xlarge": 32.77,  # $32.77/hour for A100 GPU instance
        "p3.2xlarge": 3.06,     # $3.06/hour for V100 GPU instance
    }
    
    # Estimate number of GPUs needed (rule of thumb: 1 GPU per 1B parameters)
    gpus_needed = math.ceil(model_size_billion / 1)
    
    # Calculate cost
    hourly_cost = pricing.get(gpu_type, 0) * gpus_needed
    total_cost = hourly_cost * training_hours
    
    print(f"Model size: {model_size_billion}B parameters")
    print(f"GPU type: {gpu_type}")
    print(f"GPUs needed: {gpus_needed}")
    print(f"Hourly cost: ${hourly_cost:.2f}")
    print(f"Total cost for {training_hours} hours: ${total_cost:.2f}")

estimate_cloud_costs()

Distributed Training: Scaling Beyond a Single Machine

For very large models or datasets, a single GPU or TPU may not be sufficient, and you’ll need to use distributed training. Distributed training involves splitting the model or data across multiple devices or machines, allowing you to train larger models faster. There are two main approaches: data parallelism and model parallelism. Data parallelism splits the dataset across devices, with each device holding a copy of the model and synchronizing gradients. This is simple to implement but requires high-speed interconnects to avoid bottlenecks. Model parallelism splits the model itself across devices, which is more complex but necessary for models too large to fit on a single device. Frameworks like PyTorch and TensorFlow provide built-in support for distributed training, but you’ll need to consider factors like communication overhead, fault tolerance, and hardware compatibility. For example, training a 175B-parameter model like GPT-3 requires model parallelism and hundreds of GPUs, while a 7B-parameter model can often be trained with data parallelism on a handful of GPUs.

# Example: Data parallel training with PyTorch
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

def data_parallel_training():
    # Check for multiple GPUs
    if torch.cuda.device_count() < 2:
        print("This example requires at least 2 GPUs")
        return
    
    # Create a simple model
    model = nn.Sequential(
        nn.Linear(1024, 512),
        nn.ReLU(),
        nn.Linear(512, 10)
    )
    
    # Wrap model with DataParallel
    model = nn.DataParallel(model).cuda()
    
    # Create dummy data
    x = torch.randn(1000, 1024)
    y = torch.randint(0, 10, (1000,))
    dataset = TensorDataset(x, y)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
    
    # Training loop
    optimizer = optim.Adam(model.parameters())
    criterion = nn.CrossEntropyLoss()
    
    for epoch in range(2):
        for batch_x, batch_y in dataloader:
            batch_x, batch_y = batch_x.cuda(), batch_y.cuda()
            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch + 1}, Loss: {loss.item():.4f}")

data_parallel_training()

Interview questions

All Creating Own LLM interview questions →

Check yourself

Take the full Creating Own LLM quiz →