Phase 10 - Lesson 05

Escalonamento: Treinamento Distribuído, FSDP, DeepSpeed

Seu modelo de 124M foi treinado em uma única GPU. Agora tente 7 bilhões de parâmetros. O modelo não cabe na memória. O processamento dos dados leva semanas em uma única máquina. O treinamento distribuído não é opcional em escala. É o único caminho a seguir.

Tipo: Construção Idiomas: Python Pré-requisitos: Fase 10, Lição 04 (Pré-treinamento de um Mini GPT) Tempo: ~120 minutos

Objetivos de aprendizagem

• Explicar os três tipos de paralelismo (de dados, de tensor e de pipeline) e quando cada um é necessário com base no tamanho do modelo e do cluster
• Implementar o treinamento paralelo de dados usando PyTorch DDP com sincronização de gradiente em várias GPUs
• Calcular o orçamento de memória para um determinado tamanho de modelo (pesos + estados do otimizador + gradientes + ativações) para determinar o hardware mínimo
• Configurar FSDP ou estágios do DeepSpeed ZeRO para fragmentar (shard) os estados do modelo nas GPUs e ajustar modelos que excedem a memória de uma única GPU

O problema

Um modelo de 7B parâmetros em FP16 precisa de 14 GB apenas para os pesos. O otimizador Adam armazena duas cópias adicionais de cada parâmetro (estimativas de primeiro e segundo momentos). Isso equivale a mais 28 GB. Os gradientes durante a retropropagação (backpropagation) adicionam mais 14 GB. Você já está em 56 GB antes mesmo de armazenar uma única ativação.

Uma NVIDIA A100 possui 80 GB de memória.

56 GB dos 80 GB consumidos. Isso deixa 24 GB para ativações — os valores intermediários computados durante a etapa direta (forward pass) que devem ser mantidos ativos para a retropropagação. Para uma sequência de 2048 tokens com um modelo de 4096 dimensões, as ativações de uma única camada usam cerca de 64 MB. Com 32 camadas, você precisa de 2 GB por amostra. Um tamanho de lote (batch size) de 8 requer 16 GB. Você tem 24 GB. Um tamanho de lote de 12 causa estouro de memória (blow up).

Agora tente 70B parâmetros. Pesos sozinhos: 140 GB em FP16. Não cabe em uma GPU. Você precisa de pelo menos 2 A100s (2 x 80 GB = 160 GB) apenas para conter os pesos. Adicione os estados do otimizador e gradientes e você precisará de muito mais: no mínimo 3+ GPUs, e realisticamente de 8 a 16 dependendo da estratégia de fragmentação (sharding).

O Llama 3 405B foi treinado em 16.384 GPUs NVIDIA H100. A execução do treinamento custou cerca de 100 milhões de dólares em computação. O DeepSeek V3 treinou um modelo comparável por aproximadamente 5,6 milhões de dólares ao ser inteligente em relação à arquitetura (a Mistura de Especialistas, ou Mixture of Experts, significa que apenas uma fração dos parâmetros é ativada por token) e à eficiência de treinamento.

Esta lição cobre as quatro estratégias que tornam possível o treinamento em grande escala: paralelismo de dados, paralelismo de tensor, paralelismo de pipeline e paralelismo de dados totalmente fragmentado (Fully Sharded Data Parallelism). Você simulará cada um deles em Python puro para entender a mecânica antes de tocar em qualquer framework de treinamento distribuído.

O Conceito

Por que a distribuição é necessária

Aqui está a matemática de memória para modelos reais. Cada número é calculado, não estimado.

A coluna "Estados do Adam" é a vilã. O Adam armazena uma média móvel (m) e uma variância móvel (v) para cada parâmetro, ambas em FP32. Para um modelo de 70B, isso equivale a 70B x 4 bytes x 2 = 560 GB. O otimizador sozinho precisa de sete A100s.

Uma única H100 possui 80 GB. O Llama 3 405B precisa de pelo menos 61 H100s para conter os pesos, o otimizador e os gradientes. Adicione as ativações e o número aumenta ainda mais. A Meta usou 16.384 GPUs não porque quis — mas porque foi necessário.

Paralelismo de dados

A estratégia distribuída mais simples. Copie o modelo inteiro para N GPUs. Divida cada lote de treinamento em N partes iguais. Cada GPU executa uma etapa direta (forward) e uma reversa (backward) em seu fragmento dos dados. Após a etapa reversa, faça a média dos gradientes entre todas as GPUs. Cada GPU atualiza sua cópia dos pesos com os mesmos gradientes calculados na média, mantendo todas as cópias sincronizadas.

O lado bom: Escalonamento linear de taxa de processamento (throughput). N GPUs processam N vezes mais dados por etapa. A comunicação é limitada à média dos gradientes, que se sobrepõe à computação.

O lado ruim: Cada GPU mantém uma cópia completa do modelo, dos estados do otimizador e dos gradientes. Para um modelo de 70B, cada GPU precisa de 840 GB. O paralelismo de dados não faz nada para reduzir a memória por GPU. Ele apenas reduz o tempo de treinamento.

A matemática: Tamanho do lote efetivo = tamanho_do_lote_por_gpu x N. Para N=64 GPUs com lote por GPU de 16, o lote efetivo é 1.024. O Llama 3 usou um tamanho de lote efetivo de 16 milhões de tokens por etapa.

graph TD
    subgraph DataParallel["Paralelismo de dados (N=4 GPUs)"]
        B["Lote Completo\n(1024 amostras)"] --> S["Dividir"]
        S --> G1["GPU 1\nCópia Completa do Modelo\n256 amostras"]
        S --> G2["GPU 2\nCópia Completa do Modelo\n256 amostras"]
        S --> G3["GPU 3\nCópia Completa do Modelo\n256 amostras"]
        S --> G4["GPU 4\nCópia Completa do Modelo\n256 amostras"]
        G1 --> AR["AllReduce\nMédia dos Gradientes"]
        G2 --> AR
        G3 --> AR
        G4 --> AR
        AR --> U["Atualização\n(idêntica em todas as GPUs)"]
    end

    style B fill:#1a1a2e,stroke:#e94560,color:#fff
    style G1 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G2 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G3 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G4 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style AR fill:#1a1a2e,stroke:#51cf66,color:#fff
    style U fill:#1a1a2e,stroke:#51cf66,color:#fff

Paralelismo de tensor

Divida camadas individuais entre as GPUs. Uma única multiplicação de matriz é dividida entre as GPUs, cada uma computando parte del resultado.

Considere uma matriz de pesos de formato (8192, 8192) em uma camada feedforward. Com paralelismo de tensor de 4 vias, cada GPU mantém um fragmento de (8192, 2048). Cada GPU multiplica a entrada pelo seu fragmento, produzindo um resultado parcial. Os resultados parciais são combinados (via all-reduce ou all-gather) para produzir a saída completa.

O lado bom: Reduz a memória por GPU para os pesos do modelo. Um modelo de 70B dividido em 8 GPUs significa que cada GPU mantém aproximadamente ~8.75B de parâmetros em pesos.

O lado ruim: Requer comunicação rápida entre as GPUs após cada camada. O all-reduce após cada multiplicação de matriz (matmul) adiciona latência. Isso funciona bem com NVLink (900 GB/s entre GPUs no mesmo nó), mas mal entre nós conectados por InfiniBand (400 Gb/s, cerca de 50 GB/s). O paralelismo de tensor é quase sempre limitado ao ambiente interno de um único nó (8 GPUs).

Uso real: O Megatron-LM foi pioneiro no paralelismo de tensor. O Llama 3 405B usa paralelismo de tensor de 8 vias dentro de cada nó.

Paralelismo de pipeline

Divida o modelo por camadas. A GPU 1 executa as camadas 1-8. A GPU 2 executa as camadas 9-16. A GPU 3 executa as camadas 17-24. A GPU 4 executa as camadas 25-32. Os dados fluem pelo pipeline: a GPU 1 calcula suas camadas e envia as ativações para a GPU 2, que calcula suas camadas e envia para a GPU 3, e assim por diante.

O lado bom: Comunicação mínima entre as GPUs — apenas as ativações nas fronteiras das camadas, que são pequenas em comparação com os gradientes ou pesos. Funciona entre nós porque os requisitos de largura de banda são baixos.

O lado ruim: Bolhas de pipeline (pipeline bubbles). Quando a GPU 4 está computando a etapa direta no micro-lote (micro-batch) 1, as GPUs 1, 2 e 3 estão ociosas (elas já encaminharam suas partes). Durante a etapa reversa, o padrão se inverte. Com o pipeline ingênuo, a utilização da GPU é de apenas 1/N para N estágios de pipeline.

O GPipe e o PipeDream resolvem o problema das bolhas dividindo o lote em micro-lotes. A GPU 1 começa no micro-lote 2 assim que termina de encaminhar o micro-lote 1. Isso sobrepõe a computação entre os estágios do pipeline. Com M micro-lotes e N estágios, a fração da bolha cai para (N-1)/M. Use M=16 micro-lotes com N=4 estágios e a bolha é de 3/16 = 18.75% de tempo ocioso.

FSDP: Fully Sharded Data Parallel

O FSDP combina a escalabilidade do paralelismo de dados com a eficiência de memória da fragmentação (sharding). Em vez de cada GPU manter uma cópia completa do modelo, cada GPU mantém apenas 1/N dos parâmetros, gradientes e estados do otimizador.

Antes da etapa direta de uma camada, o FSDP executa um all-gather para coletar os parâmetros completos de todas as GPUs na memória de cada GPU. Após a etapa direta, cada GPU descarta os parâmetros não locais. Durante a etapa reversa, o all-gather é executado novamente para reconstruir os parâmetros para a computação do gradiente. Após a etapa reversa, um reduce-scatter distribui os fragmentos de gradiente para que cada GPU armazene apenas 1/N dos gradientes.

A matemática para um modelo de 70B em 8 GPUs:

Sem FSDP, você não consegue colocar um modelo de 70B em uma única GPU de 80 GB. Com FSDP em 8 GPUs, cada GPU usa 105 GB — espere, isso ainda não cabe. Você precisa de pelo menos 16 GPUs para ficar abaixo de 80 GB por GPU, ou combinar o FSDP com checkpointing de ativação (recomputar ativações durante a etapa reversa em vez de armazená-las).

O custo de comunicação é maior do que o paralelismo de dados tradicional devido ao all-gather antes de cada camada. Mas a economia de memória torna possíveis execuções de treinamento que antes seriam impossíveis.

graph TD
    subgraph FSDP["FSDP: Fully Sharded Data Parallel (4 GPUs)"]
        direction TB
        S["Modelo: 4 camadas, fragmentado"]

        subgraph GPU1["GPU 1"]
            G1S["Fragmento: 1/4 params\n1/4 otimizador\n1/4 gradientes"]
        end
        subgraph GPU2["GPU 2"]
            G2S["Fragmento: 1/4 params\n1/4 otimizador\n1/4 gradientes"]
        end
        subgraph GPU3["GPU 3"]
            G3S["Fragmento: 1/4 params\n1/4 otimizador\n1/4 gradientes"]
        end
        subgraph GPU4["GPU 4"]
            G4S["Fragmento: 1/4 params\n1/4 otimizador\n1/4 gradientes"]
        end

        AG["All-Gather\n(reconstrói params completos\nantes de cada camada)"]
        FW["Etapa Direta\n(params completos temporariamente)"]
        RS["Reduce-Scatter\n(distribui fragmentos de gradiente\napós a etapa reversa)"]

        S --> GPU1
        S --> GPU2
        S --> GPU3
        S --> GPU4
        GPU1 --> AG
        GPU2 --> AG
        GPU3 --> AG
        GPU4 --> AG
        AG --> FW
        FW --> RS
    end

    style G1S fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G2S fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G3S fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G4S fill:#1a1a2e,stroke:#0f3460,color:#fff
    style AG fill:#1a1a2e,stroke:#e94560,color:#fff
    style FW fill:#1a1a2e,stroke:#51cf66,color:#fff
    style RS fill:#1a1a2e,stroke:#e94560,color:#fff

DeepSpeed ZeRO

O ZeRO (Zero Redundancy Optimizer) da DeepSpeed é conceitualmente idêntico ao FSDP, mas foi desenvolvido de forma independente pela Microsoft. Ele define três estágios, cada um fragmentando de forma mais agressiva:

O ZeRO-3 é equivalente ao FSDP. A nomenclatura é diferente, o mecanismo é o mesmo. O PyTorch adicionou o FSDP como uma implementação nativa depois que o DeepSpeed provou o conceito.

O DeepSpeed também introduziu o ZeRO-Offload (descarregar estados do otimizador para a memória RAM da CPU, que é mais barata e maior) e o ZeRO-Infinity (descarregar para SSDs NVMe). Eles trocam a velocidade de computação por capacidade de memória — as operações descarregadas são mais lentas, mas liberam memória da GPU.

Treinamento com precisão mista

O treinamento moderno usa múltiplos formatos de ponto flutuante simultaneamente:

• Etapa direta: FP16 ou BF16 (16 bits). Metade da memória do FP32. Matmuls rodam 2x mais rápido em Tensor Cores.
• Pesos mestre: FP32 (32 bits). Mantidos pelo otimizador para precisão numérica durante as atualizações dos pesos.
• Escalonamento de perda (loss scaling): Multiplica a perda por uma constante grande antes da etapa reversa para evitar que os gradientes em FP16 sofram de subfluxo (underflow) para zero. Divide pela mesma constante antes da etapa do otimizador.

BF16 (Brain Float 16) possui a mesma faixa de expoente do FP32 (8 bits de expoente), mas precisão reduzida (7 bits de mantissa contra 23 do FP32). Ele raramente precisa de escalonamento de perda porque pode representar a mesma faixa de valores. O FP16 possui 5 bits de expoente e 10 bits de mantissa — ele pode representar valores detalhados, mas estoura ou zera em magnitudes extremas.

As TPUs do Google usam BF16 nativamente. As A100 e H100 da NVIDIA suportam tanto FP16 quanto BF16. A indústria migrou amplamente para o BF16 porque ele elimina as dores de cabeça do escalonamento de perda.

Comparação de memória para um modelo de 7B:

A precisão mista economiza 28 GB neste modelo. Os estados do otimizador permanecem em FP32 de qualquer forma — é aqui que a maior parte da memória é consumida.

Megatron-LM e paralelismo 3D

O treinamento real em grande escala combina os três tipos de paralelismo:

• Paralelismo de dados entre grupos de nós (escalar o tamanho do lote)
• Paralelismo de tensor dentro de um nó (dividir camadas em 8 GPUs)
• Paralelismo de pipeline entre nós (dividir grupos de camadas entre máquinas)

Llama 3 405B em 16.384 H100s:

• 8-way tensor parallelism within each node (8 GPUs per node)
• 16-way pipeline parallelism across nodes (16 pipeline stages)
• 128-way data parallelism across the remaining dimension (16,384 / 8 / 16 = 128)

Essa decomposição 3D (8 x 16 x 128 = 16.384) é a forma de escalar para milhares de GPUs. Cada GPU vê um fragmento de dados diferente (paralelo de dados), mantém uma fatia de cada camada (paralelo de tensor) e computa um conjunto diferente de camadas (paralelo de pipeline).

O DeepSeek V3 adotou uma abordagem diferente. Sua arquitetura de Mistura de Especialistas (MoE) ativa apenas 37B de 671B parâmetros por token. Isso significa que cada GPU precisa computar apenas (e armazenar ativações para) os parâmetros ativos. Eles treinaram em 2.048 GPUs H800 — menos de 1/8 da contagem de GPUs da Meta — por $5,6M vs os estimados

graph TD
    subgraph ThreeD["Paralelismo 3D (Llama 3 405B)"]
        direction TB
        subgraph DP["Paralelismo de Dados (128 vias)\nDividir lote em 128 grupos"]
            subgraph PP["Paralelismo de Pipeline (16 vias)\nDividir camadas em 16 estágios"]
                subgraph TP["Paralelismo de Tensor (8 vias)\nDividir cada camada em 8 GPUs"]
                    G1["GPU 1\nFatia das camadas 1-N"]
                    G2["GPU 2\nFatia das camadas 1-N"]
                    G8["GPU 8\nFatia das camadas 1-N"]
                end
            end
        end
    end

    N1["Total: 8 x 16 x 128 = 16,384 GPUs"]

    style G1 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G2 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G8 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style N1 fill:#1a1a2e,stroke:#e94560,color:#fff

import numpy as np

def simulate_data_parallelism(data, num_gpus, model_fn):
    batch_size = len(data)
    shard_size = batch_size // num_gpus
    remainder = batch_size % num_gpus

    gpu_losses = []
    gpu_gradients = []

    offset = 0
    for gpu_id in range(num_gpus):
        extra = 1 if gpu_id < remainder else 0
        shard = data[offset:offset + shard_size + extra]
        offset += shard_size + extra

        loss, grad = model_fn(shard)
        gpu_losses.append(loss)
        gpu_gradients.append(grad)

    avg_loss = np.mean(gpu_losses)
    avg_gradient = np.mean(gpu_gradients, axis=0)

    return avg_loss, avg_gradient

def simulate_tensor_parallelism(input_data, weight_matrix, num_gpus):
    d_in, d_out = weight_matrix.shape
    assert d_out % num_gpus == 0, f"d_out {d_out} not divisible by num_gpus {num_gpus}"
    shard_size = d_out // num_gpus

    partial_results = []
    for gpu_id in range(num_gpus):
        start = gpu_id * shard_size
        end = start + shard_size
        weight_shard = weight_matrix[:, start:end]

        partial = input_data @ weight_shard
        partial_results.append(partial)

    full_output = np.concatenate(partial_results, axis=-1)

    direct_output = input_data @ weight_matrix
    error = np.abs(full_output - direct_output).max()

    return full_output, error

def simulate_pipeline_parallelism(num_layers, num_stages, num_microbatches):
    layers_per_stage = num_layers // num_stages

    timeline = {}
    clock = 0

    for mb in range(num_microbatches):
        for stage in range(num_stages):
            start_time = max(
                timeline.get((stage, mb - 1, "fwd"), (0, 0))[1] if mb > 0 else 0,
                timeline.get((stage - 1, mb, "fwd"), (0, 0))[1] if stage > 0 else 0,
            )
            end_time = start_time + layers_per_stage
            timeline[(stage, mb, "fwd")] = (start_time, end_time)

    last_fwd_end = max(v[1] for v in timeline.values())

    for mb in range(num_microbatches - 1, -1, -1):
        for stage in range(num_stages - 1, -1, -1):
            deps = [last_fwd_end]
            if mb < num_microbatches - 1 and (stage, mb + 1, "bwd") in timeline:
                deps.append(timeline[(stage, mb + 1, "bwd")][1])
            if stage < num_stages - 1 and (stage + 1, mb, "bwd") in timeline:
                deps.append(timeline[(stage + 1, mb, "bwd")][1])
            start_time = max(deps)
            end_time = start_time + layers_per_stage
            timeline[(stage, mb, "bwd")] = (start_time, end_time)

    total_time = max(v[1] for v in timeline.values())
    compute_time = num_microbatches * num_stages * layers_per_stage * 2
    bubble_fraction = 1.0 - compute_time / (total_time * num_stages)

    return timeline, total_time, bubble_fraction

def memory_calculator(
    params_billions,
    precision_bytes=2,
    optimizer="adam",
    num_gpus=1,
    sharding="none",
    sequence_length=2048,
    batch_size_per_gpu=1,
    hidden_dim=None,
    num_layers=None,
):
    params = params_billions * 1e9

    weight_memory = params * precision_bytes

    if optimizer == "adam":
        optimizer_memory = params * 4 * 2
    elif optimizer == "sgd":
        optimizer_memory = params * 4
    else:
        optimizer_memory = 0

    gradient_memory = params * precision_bytes

    total_no_activation = weight_memory + optimizer_memory + gradient_memory

    if hidden_dim and num_layers:
        activation_per_layer = (
            sequence_length * batch_size_per_gpu * hidden_dim * precision_bytes * 4
        )
        activation_memory = activation_per_layer * num_layers
    else:
        activation_memory = params * precision_bytes * 0.5

    if sharding == "fsdp" or sharding == "zero3":
        weight_memory /= num_gpus
        optimizer_memory /= num_gpus
        gradient_memory /= num_gpus
    elif sharding == "zero2":
        optimizer_memory /= num_gpus
        gradient_memory /= num_gpus
    elif sharding == "zero1":
        optimizer_memory /= num_gpus

    per_gpu_total = weight_memory + optimizer_memory + gradient_memory + activation_memory

    return {
        "params_billions": params_billions,
        "weights_gb": weight_memory / 1e9,
        "optimizer_gb": optimizer_memory / 1e9,
        "gradients_gb": gradient_memory / 1e9,
        "activations_gb": activation_memory / 1e9,
        "per_gpu_total_gb": per_gpu_total / 1e9,
        "total_across_gpus_gb": per_gpu_total * num_gpus / 1e9,
        "fits_on_80gb": per_gpu_total / 1e9 <= 80,
        "num_gpus": num_gpus,
        "sharding": sharding,
    }

def mixed_precision_comparison(params_billions):
    params = params_billions * 1e9

    fp32_weights = params * 4
    fp32_optimizer = params * 4 * 2
    fp32_gradients = params * 4
    fp32_total = fp32_weights + fp32_optimizer + fp32_gradients

    fp16_weights = params * 2
    fp16_master = params * 4
    fp16_optimizer = params * 4 * 2
    fp16_gradients = params * 2
    fp16_total = fp16_weights + fp16_master + fp16_optimizer + fp16_gradients

    mixed_weights = params * 2
    mixed_optimizer = params * 4 * 2
    mixed_gradients = params * 2
    mixed_total = mixed_weights + mixed_optimizer + mixed_gradients

    return {
        "fp32_total_gb": fp32_total / 1e9,
        "fp16_with_master_gb": fp16_total / 1e9,
        "mixed_bf16_gb": mixed_total / 1e9,
        "savings_vs_fp32": 1 - mixed_total / fp32_total,
    }

def run_all_demos():
    print("=" * 70)
    print("DATA PARALLELISM SIMULATION")
    print("=" * 70)

    np.random.seed(42)
    data = np.random.randn(64, 32)
    weight = np.random.randn(32, 16)

    def model_fn(batch):
        output = batch @ weight
        loss = np.mean(output ** 2)
        grad = 2 * batch.T @ (batch @ weight) / len(batch)
        return loss, grad

    for n_gpus in [1, 2, 4, 8]:
        loss, grad = simulate_data_parallelism(data, n_gpus, model_fn)
        print(f"  {n_gpus} GPUs: loss={loss:.4f}, grad_norm={np.linalg.norm(grad):.4f}")

    print()
    print("=" * 70)
    print("TENSOR PARALLELISM SIMULATION")
    print("=" * 70)

    x = np.random.randn(4, 8192)
    W = np.random.randn(8192, 8192)

    for n_gpus in [1, 2, 4, 8]:
        output, error = simulate_tensor_parallelism(x, W, n_gpus)
        print(f"  {n_gpus} GPUs: output_shape={output.shape}, max_error={error:.2e}")

    print()
    print("=" * 70)
    print("PIPELINE PARALLELISM SIMULATION")
    print("=" * 70)

    for n_mb in [1, 4, 8, 16, 32]:
        _, total_t, bubble = simulate_pipeline_parallelism(32, 4, n_mb)
        print(f"  {n_mb:2d} micro-batches: total_time={total_t:4d}, bubble={bubble:.1%}")

    print()
    print("=" * 70)
    print("MEMORY CALCULATOR")
    print("=" * 70)

    configs = [
        (7, "none", 1),
        (7, "fsdp", 8),
        (70, "none", 1),
        (70, "fsdp", 8),
        (70, "fsdp", 16),
        (405, "fsdp", 64),
        (405, "fsdp", 128),
    ]

    print(f"  {'Model':>8} {'Sharding':>8} {'GPUs':>5} {'Per-GPU':>10} {'Fits 80GB':>10}")
    print("  " + "-" * 50)
    for params, shard, gpus in configs:
        result = memory_calculator(params, num_gpus=gpus, sharding=shard)
        fits = "Yes" if result["fits_on_80gb"] else "No"
        print(f"  {params:>6}B {shard:>8} {gpus:>5} {result['per_gpu_total_gb']:>8.1f}GB {fits:>10}")

    print()
    print("=" * 70)
    print("MIXED PRECISION COMPARISON")
    print("=" * 70)

    for params_b in [7, 13, 70, 405]:
        result = mixed_precision_comparison(params_b)
        print(f"  {params_b}B: FP32={result['fp32_total_gb']:.0f}GB, "
              f"Mixed BF16={result['mixed_bf16_gb']:.0f}GB, "
              f"Savings={result['savings_vs_fp32']:.0%}")