祝融说。

In the previous chapters, we have delved into the various modes and parallelism strategies of large model training. Now, we face the ultimate question that every AI Infra engineer must confront: "How many resources do I need?"

This question comes in many forms:

• An algorithm engineer asks: "I want to fine-tune a 70B model. How many A800s do I need?"
• Your boss asks: "We plan to pre-train a hundred-billion-parameter model from scratch. How many servers do we need to purchase? What is the approximate budget? How long will it take?"
• You ask yourself: "Why did this training task OOM? Which part is blowing up the memory?"

Answering these questions can no longer rely on experience and gut feeling. What you need is a set of scientific, rigorous quantitative analysis methods. In this chapter, we will become "compute power actuaries." Starting from first principles, we will derive and establish the two core mathematical models for LLM training: the memory footprint model and the training throughput model. We will use formulas and calculators to transform vague "resource requirements" into precise numbers. Finally, through a real-world hardware selection case study, we will apply these theories to the decision-making of cluster planning.

Mastering this chapter will give you the "superpower" to see through the resource consumption of AI training. The next time you face resource planning, you will no longer be a passive executor, but a decision-maker armed with data and strategy.

7.1 Deriving Memory Formulas: Parameters, Gradients, Optimizer States, Activations = ?

GPU memory is the most precious and easily bottlenecked resource in AI training. Whether a training task can run depends first and foremost on whether its peak memory usage exceeds the physical memory capacity of a single GPU. Understanding how memory is consumed is the first step in troubleshooting and performance optimization.

The memory footprint of a Transformer large model during training consists of four main components: Model Parameters, Gradients, Optimizer States, and Activations. Additionally, there are miscellaneous overheads like temporary buffers and the CUDA runtime itself.

7.1.1 Basics: Data Types and Byte Sizes

Before calculating, we must clarify the byte sizes occupied by different data precisions:

• FP32 (Single Precision Floating Point): 4 bytes
• FP16 (Half Precision Floating Point): 2 bytes
• BF16 (bfloat16): 2 bytes
• INT8 (8-bit Integer): 1 byte
• 1B (Billion) parameters = 1,000,000,000 parameters

7.1.2 The Four Components of Memory Footprint

Assume we have a model with parameter count P (in Billions).

Model Parameters

This is the most intuitive part. The model itself must be loaded into memory.

• In standard FP32 training, model parameters occupy P * 10^9 * 4 bytes.
• In the mainstream mixed precision training, typically one copy of FP32 "master weights" (for precise gradient updates) and one copy of FP16 weights (for efficient forward and backward computation) are kept.
  • Model parameters occupy: P * 10^9 * (4 + 2) = 6P GB (approximately, since 10^9 bytes is roughly 1 GB).
  • However, in many modern framework implementations (e.g., PyTorch AMP), for optimization purposes, only FP16 weights and FP32 gradients may be kept in memory, with parameter updates done on CPU or temporarily on GPU. A more common arrangement is that one copy of FP16 weights is kept in memory for computation.
  • To simplify the model, we typically focus on the main memory overhead. Under mixed precision, the weights used for computation are FP16, so at least P * 2 GB is occupied.

Gradients

The backpropagation algorithm computes a gradient for every trainable parameter in the model. Therefore, the number of gradients is exactly the same as the number of parameters.

In mixed precision training, gradients are typically stored in FP32 format to maintain update precision.

Gradient footprint: P * 10^9 * 4 = 4P GB

Optimizer States

This is the most easily overlooked and most common cause of OOM -- the "memory killer." Optimizers like Adam or AdamW maintain their own "states" when updating model parameters.

• Adam/AdamW Optimizer: It maintains two states for each model parameter:
  • First Moment (Momentum): Stores the exponential moving average of past gradients, typically FP32.
  • Second Moment (Variance): Stores the exponential moving average of past squared gradients, typically FP32.
• Therefore, the Adam optimizer introduces an additional memory overhead of 2 times the model parameter count.
  • Optimizer state footprint: P * 10^9 * 4 (momentum) + P * 10^9 * 4 (variance) = 8P GB
• In some frameworks, FP16 may be used to store optimizer states to save memory, but this sacrifices some precision and stability. By default, we calculate with FP32.

Summary 1: Static Memory (Independent of Batch Size)

We call the three components -- model parameters, gradients, and optimizer states -- "static memory" because their footprint depends only on the model parameter count P, not on the Batch Size used for training.

In a standard full fine-tuning (or pre-training) run using the Adam optimizer, the static memory footprint is:

Memory_static = (P * 2) [FP16 params] + (P * 4) [FP32 gradients] + (P * 8) [Adam states]

Memory_static = 14P GB

If using DeepSpeed ZeRO-2, which shards gradients and optimizer states across N GPUs, the per-GPU static memory becomes:

Memory_static_ZeRO2 = (P * 2) + (P * 4 / N) + (P * 8 / N) = 2P + 12P / N GB

If using DeepSpeed ZeRO-3, which also shards model parameters, the per-GPU static memory becomes:

Memory_static_ZeRO3 = (P * 2 / N) + (P * 4 / N) + (P * 8 / N) = 14P / N GB

LoRA's Memory Advantage: Why does LoRA save memory? Because it only trains a very small number of parameters (let us call it P'). Therefore, gradients and optimizer states exist only for these P' parameters. Static memory is approximately P * 2 (original model FP16 weights, frozen) + 14 * P' (P' is much smaller than P). This dramatically reduces memory overhead.

Activations

This is the most complex and dynamic part of memory usage.

• What are Activations? During forward propagation, the output of each layer of the network must be saved (cached in memory) because they are needed during backward propagation to compute gradients (according to the chain rule). These cached intermediate results are activations.
• What Determines Activation Size?
  • Sequence Length (s): The longer the input text, the larger the intermediate activation matrices.
  • Batch Size (b): The more samples processed at once, the more activations need to be cached.
  • Hidden Dimension (h): The "width" of the model.
  • Number of Attention Heads (a).
  • Number of Layers (L): The "depth" of the model.
• Estimation Formula (Very Important):
  • For a standard Transformer model, there is an approximate estimation formula for activation memory: Memory_activations = s * b * h * L * (10 + 24/T) / T (when using sequence parallelism).
  • A more practical, widely cited simplified formula (without sequence parallelism and using mixed precision FP16) is: Memory_activations_per_gpu = s * b * h * L * (1 + a/(h*T)) * 2.
  • A rougher but easy-to-remember empirical formula (from the BLOOM paper): Memory_activations = 2 * s * b * P / L GB.
  • Let us adopt a widely cited formula derived by Anyscale engineers (assuming FP16/BF16, i.e., 2 bytes per value): Memory_activations_GB = (s * b * h * L * (34 + 5 * a * s / (h * T))) / 10^9, where T is the tensor parallelism size. This formula originates from the activation-memory analysis in the Megatron-LM paper and holds under specific assumptions (no sequence parallelism, no selective recomputation, standard Transformer architecture).

The derivation of this formula is complex, accounting for all details of the attention mechanism: K,V cache, Softmax output, Dropout masks, and so on. As an AI Infra engineer, we do not need to derive it from scratch, but we must be able to use it and understand the meaning of its variables.

Key Insights:

• Activations are proportional to s and b. Doubling the sequence length doubles the activation memory.
• Activations are proportional to model size (h, L).
• Using tensor parallelism (T > 1) can significantly reduce activation memory, as some of the model's hidden states are also sharded.
• Activation Recomputation/Checkpointing: This is a technique that trades computation for memory. During forward propagation, it no longer saves all activations, only a small subset of "key points." When backward propagation needs an activation that was not saved, it recomputes it from the nearest "key point" in a forward pass. This increases computation time (typically on the order of 20-33%, depending on the implementation and strategy) but can reduce activation memory by an order of magnitude. DeepSpeed and Megatron-LM both support this technique.

7.1.3 The Ultimate Formula and Toolbox

Total Per-GPU Memory Footprint (Full Fine-Tuning, FP16, Adam): Memory_Total = (2P) [params] + (4P) [gradients] + (8P) [optimizer] + Memory_activations Memory_Total = 14P + (s * b * h * L * ...)

This formula gives the theoretical peak. In practice, through parallelism strategies like ZeRO, the gradient and optimizer parts can be distributed across N cards.

Total Per-GPU Memory Footprint (ZeRO-3 + Activation Recomputation): Memory_Total_optimized = 14P / N + Memory_activations_recompute

Build Your Excel/Python Compute Calculator: Every AI Infra engineer should have this "artifact."

• Input:
  • Model parameter count P (Billions)
  • Number of layers L
  • Hidden dimension h
  • Number of attention heads a
  • Sequence length s
  • Per-GPU batch size b (Micro Batch Size)
  • Parallelism strategy: DP size, TP size, PP size
  • Optimization techniques: ZeRO enabled (which stage), activation recomputation enabled
• Output:
  • Per-GPU static memory (GB)
  • Per-GPU activation memory (GB)
  • Per-GPU total peak memory (GB)

This calculator will be your most powerful weapon for analyzing OOM issues and evaluating resource requests.

7.2 Compute Estimation Model: Estimating Training Days Based on Token Count and FLOPs

Having solved the memory problem (whether it can run), we now need to solve the time problem (how long it will take). This directly impacts project scheduling and cost budgeting.

The core of estimating training time lies in calculating two things: the total amount of computation needed (Total FLOPs), and the effective compute power our cluster can provide per second (Effective TFLOPS).

Total Training Time = Total FLOPs / Cluster Effective Compute

7.2.1 Estimating Total FLOPs

For Transformer models, there is a widely accepted empirical formula for the computation during training (FLOPs):

Total FLOPs = 6 * P * D

• P: Model parameter count (note: this is the actual number of parameters, not in Billions). For example, a 70B model would be 70 * 10^9.
• D: Total number of tokens in the training dataset.
• Where does the factor 6 come from?
  • Forward Pass: For a model with P parameters, processing 1 token requires approximately 2 * P FLOPs. The 2 is because each parameter participates in one multiplication and one addition in a matrix multiplication.
  • Backward Pass: According to the chain rule, the computation for the backward pass is approximately twice that of the forward pass, or 4 * P FLOPs.
  • Total: 2P + 4P = 6P FLOPs.

Example: Pre-training a 70B model using a dataset of 2T (2 trillion) tokens.

• P = 70 * 10^9
• D = 2 * 10^12
• Total FLOPs = 6 * (70 * 10^9) * (2 * 10^12) = 8.4 * 10^23 FLOPs. This is an astronomical number!

7.2.2 Cluster Effective Compute (Effective TFLOPS)

This is the part of the estimation that most relies on experience. It is not simply the sum of all GPUs' theoretical peak performance. Cluster Effective Compute = Single GPU Theoretical Peak * GPU Count * MFU

• Single GPU Theoretical Peak TFLOPS:

  • Obtained from the hardware specification sheet. Key: Use the compute value that matches the training precision.
  • Example: NVIDIA A100 theoretical peak at FP16/BF16 is 312 TFLOPS.
  • Ascend 910B theoretical peak at FP16 is 320 TFLOPS.

• GPU Count (N): The total number of GPUs in the cluster used for this task.

• MFU (Model FLOPs Utilization):

  • This is the most important and most uncertain variable. It represents how much of the hardware's theoretical performance we are actually extracting.
  • MFU is affected by countless factors: network communication overhead (especially All-Reduce), data loading and preprocessing bottlenecks, unoptimized CUDA kernels, memory wall, and so on.
  • MFU Value Ranges (Empirical):
    • Poor (10-20%): Severe network bottleneck or poor code optimization.
    • Average (20-30%): Common level, still room for improvement.
    • Good (30-50%): Well-optimized cluster and code, e.g., using InfiniBand network, communication overlapping with computation.
    • Excellent (50%+): Industry-leading level, typically achieved by vertically integrated companies like NVIDIA and Google.
  • When doing upfront planning, using a relatively conservative MFU value (such as 30%) is prudent.

7.2.3 Training Time Formula

Total Training Time (seconds) = (6 * P * D) / (Single GPU Peak TFLOPS * N * MFU * 10^12)

For convenience, we usually convert time to days: Total Training Time (days) = Total Training Time (seconds) / (3600 * 24)

Build Your Excel/Python Training Time Calculator:

• Input:
  • Model parameter count P (Billions)
  • Training data size D (Trillions of Tokens)
  • Single GPU peak compute (TFLOPS)
  • GPU count N
  • Estimated MFU (%)
• Output:
  • Estimated total training days
  • Estimated total cost (if per-GPU-hour cost is provided)

This tool will be a powerful support for reporting project timelines and budgets to management.

7.3 Hardware Selection Practice: How Many Cards for a 70B Model? How to Plan Cluster Scale?

Now, let us apply the formulas from the previous two sections to a selection decision scenario (a synthetic computational example demonstrating the method; recalculate the figures for your own model architecture, hardware models, and measured MFU).

Scenario: A company decides to follow the open-source trend and pre-train a foundation model comparable to Llama 2 70B from scratch, based on a high-quality Chinese dataset. The CEO has tasked you with planning the required compute cluster and providing a budget and timeline.

Known Conditions and Constraints:

• Target Model: P = 70B (70 billion parameters). Assume its structure is L=80 layers, h=8192, a=64.
• Training Data: D = 2T (2 trillion) tokens.
• Hardware Options:
  • Option A: NVIDIA A800 (80 GB memory, 312 TFLOPS@FP16)
  • Option B: Huawei Ascend 910B (64 GB memory, 320 TFLOPS@FP16)
• Timeline Requirement: Complete training within 3 months (approximately 90 days).

Step 1: Memory Analysis (Determining Intra-Node Parallelism and Minimum GPU Count)

We need to find a parallelism strategy that allows the 70B model training task to fit into a single GPU. We will use mixed precision training (BF16) and the AdamW optimizer. Sequence length s is typically set to 4096.

• Calculating Activation Memory (Rough Estimate): Memory_activations = 2 * s * b * P / L = 2 * 4096 * b * 70 / 80 = 7168 * b MB To keep activation memory at the GB level, the micro batch size b must be very small, typically set to 1 or 2. Assuming b=1, activation memory is approximately 7.2 GB. This seems manageable, but this is without activation recomputation. It is a significant burden.

  If we enable activation recomputation, we can reduce this part of memory to well below 1 GB. Therefore, for training large models, activation recomputation is mandatory.

• Calculating Static Memory (Full Fine-Tuning): Memory_static = 14 * P = 14 * 70 = 980 GB. This is a staggering number; no single GPU can hold it.

• Introducing Parallelism to Reduce Memory: We must use ZeRO, TP, PP, and related techniques. A typical, efficient configuration is 8-way TP + ZeRO-1 (since TP already significantly shards model parameters and optimizer states).

  • Tensor Parallelism (TP=8): Shards model parameters and optimizer states across 8 cards. This is typically done within a single 8-card node.
  • Per-GPU Static Memory (TP=8): 14 * P / 8 = 980 / 8 = 122.5 GB.
  • This number still exceeds both the A800's 80 GB and the 910B's 64 GB!

• Combination: TP + PP + ZeRO

  • We need to introduce Pipeline Parallelism (PP) to further split the model.
  • Assume we use PP=4.
  • Per-GPU model parameters: P / (TP * PP) = 70 / (8 * 4) = 2.18B.
  • Per-GPU static memory (rough, not rigorous): Is it 14 * P / (TP * PP)? Not exactly. The sharding mechanisms of ZeRO and PP differ.
  • A more engineering-oriented approach is to use existing open-source configurations as a reference. For example, Megatron-LM typically uses TP=8, PP=16 or higher when training models of similar scale.
  • Let us try a configuration: TP=8, PP=8, DP=N'.
    • Each pipeline stage handles 80/8=10 layers.
    • Per-GPU model parameters and optimizer states are primarily sharded by TP, approximately 122.5 GB. Still too large.

• Ultimate Solution: Reference Mature Frameworks' Memory Calculators

  • The documentation for DeepSpeed and Megatron-LM typically provides more precise memory calculators or empirical formulas. According to NVIDIA's public material, a 70B model, with TP=8, PP=1, using BF16 and activation recomputation, has a peak memory of about 105 GB. This still exceeds the A800.
  • To fit it within the 80 GB A800, PP is needed. A feasible configuration is:
    • TP = 4, PP = 4
    • In this configuration, each GPU is responsible for about 1/16 of the model, and memory usage can be controlled within 80 GB.

Conclusion 1: Training a 70B model requires at least a 16-card cluster (TP=4, PP=4) to run a single replica. For 8-card/node servers, this means at least 2 servers.

Step 2: Time Analysis (Determining Total Cluster Size)

We aim to complete training within 90 days. Now let us calculate how many cards are needed.

• Total FLOPs: Total FLOPs = 6 * (70 * 10^9) * (2 * 10^12) = 8.4 * 10^23 FLOPs
• Target Effective Compute: Required Effective TFLOPS = Total FLOPs / (90 * 24 * 3600 * 10^12) = 107.5 * 10^3 TFLOPS = 107,500 TFLOPS
• Calculating Required GPU Count (N):
  • N = Required Effective TFLOPS / (Single GPU Peak TFLOPS * MFU)
  • We use a relatively optimistic but achievable MFU = 40% (0.4) for planning.
  • For A800 (312 TFLOPS): N = 107500 / (312 * 0.4) = 861 cards.
  • For Ascend 910B (320 TFLOPS): N = 107500 / (320 * 0.4) = 840 cards.
• Rounding Up and Considering Redundancy:
  • The calculation yields 861 cards. Considering hardware failures, maintenance, and other factors, some redundancy is needed. Also, to facilitate forming standard parallel units (for example, one DP replica is 16 cards), the cluster size should be an integer multiple of the parallel unit.
  • We can plan a cluster with 861 / 16 = 54 DP replicas.
  • 16 * 54 = 864 cards. Add some redundancy, say 5%: 864 * 1.05 = 907 cards.
  • To round up and simplify network topology design, we can ultimately plan a cluster of 1024 cards (e.g., 128 servers with 8 cards each).

Step 3: Final Plan and Budget

• Cluster Size: 1024 A800 (or 910B) cards, consisting of 128 servers with 8 cards each.
• Parallelism Strategy: TP=4, PP=4, DP=64.
• Network Requirements: A high-performance RDMA network is mandatory. For a scale of 128 servers, a Fat-Tree architecture with InfiniBand or RoCE v2 is necessary.
• Storage Requirements: A parallel file system capable of providing at least 100 GB/s aggregate bandwidth, with capacity in the PB range.
• Timeline Review: Actual Training Days = (8.4e23) / (312 * 1024 * 0.4 * 1e12 * 3600 * 24) = 76 days. This result is within the 90-day target, leaving some buffer time.
• Budget Estimate (order-of-magnitude reference, fluctuating with market conditions):
  • Hardware Cost: 128 servers with 8 cards each (including CPU, memory, hard drives) + 1024 A800/910B cards + high-performance switches + racks, PDUs, and so on. This is an investment on the order of several hundred million RMB.
  • Power Cost: Assuming 40 KW per rack, the 1024 cards require about 32-40 racks, with total power exceeding 1 megawatt. Electricity bills over three years would also be on the order of tens of millions.
  • Personnel Cost: A professional AI Infra team is needed for construction and operations.

Final Conclusion to Present to the CEO (based on this example's assumptions):

"To complete the pre-training of a 70B model within 3 months, we recommend building a dedicated compute cluster of 1024 A800 (or Ascend 910B) cards. Under the assumption of 40% MFU, the estimated training duration for this plan is 76 days. The hardware investment is on the order of several hundred million RMB, with significant ongoing operational costs. As an alternative, we could also consider renting public cloud compute, which, despite a higher unit price, avoids the massive upfront capital expenditure and operational burden."

This conclusion, based on quantitative analysis, is far more powerful than a vague statement like "we need a lot of cards." This is the power that mathematics gives to the AI Infra engineer.