Estimate Complexity

Estimating the total number of floating-point operations (FLOPs) required to train GPT-3 involves breaking down the training process into its fundamental components: forward pass, backward pass, and parameter updates. Here's a step-by-step breakdown:

---

1. Key Parameters of GPT-3

GPT-3 has the following key specifications (from the paper or known details):

• Model size (parameters): ~175 billion (1.75 × 10¹¹) parameters.
• Training tokens: ~300 billion tokens (3 × 10¹¹).
• Batch size: ~3.2 million tokens per batch (3.2 × 10⁶).
• Sequence length: ~2048 tokens.
• Number of layers: 96 (for the 175B model).
• Hidden dimension: ~12,288 (1.2288 × 10⁴).
• Number of attention heads: 96.
• Optimizer: Adam (or variant), which requires storing additional state (e.g., momentum and variance).

---

2. FLOPs per Forward Pass

The forward pass of a transformer model involves:

• Embedding lookup: Negligible compared to other operations.
• Self-attention: For each layer, the self-attention mechanism computes:
  • Query, Key, Value projections: 3 × (sequence_length × hidden_dim × hidden_dim) = 3 × (2048 × 12288 × 12288) FLOPs.
  • Attention scores: sequence_length × sequence_length × hidden_dim = 2048 × 2048 × 12288 FLOPs.
  • Softmax and weighted sum: ~sequence_length × sequence_length × hidden_dim FLOPs.
  • Output projection: sequence_length × hidden_dim × hidden_dim = 2048 × 12288 × 12288 FLOPs.
• Feed-forward network (FFN): For each layer, the FFN has two linear layers with an expansion factor of 4:
  • First layer: sequence_length × hidden_dim × (4 × hidden_dim) = 2048 × 12288 × 49152 FLOPs.
  • Second layer: sequence_length × (4 × hidden_dim) × hidden_dim = 2048 × 49152 × 12288 FLOPs.
• Layer normalization and residual connections: Negligible.

Simplified FLOPs per Layer:

For one layer, the dominant terms are:

• Self-attention: ~6 × (sequence_length × hidden_dim²) = 6 × (2048 × 12288²) ≈ 1.8 × 10¹² FLOPs.
• FFN: ~8 × (sequence_length × hidden_dim²) = 8 × (2048 × 12288²) ≈ 2.4 × 10¹² FLOPs.
• Total per layer: ~4.2 × 10¹² FLOPs.

For 96 layers: 96 × 4.2 × 10¹² ≈ 4.0 × 10¹⁴ FLOPs per forward pass.

FLOPs per Token:

Since the sequence length is 2048, the FLOPs per token is: 4.0 × 10¹⁴ / 2048 ≈ 2.0 × 10¹¹ FLOPs per token.

---

3. FLOPs per Backward Pass

The backward pass requires roughly 2× the FLOPs of the forward pass because:

• Gradients must be computed for all operations (same as forward).
• Additional operations for weight updates (e.g., gradient accumulation).

Thus, backward pass FLOPs ≈ 2 × forward pass FLOPs ≈ 8.0 × 10¹⁴ FLOPs per batch.

---

4. Total FLOPs per Batch

• Forward pass: 4.0 × 10¹⁴ FLOPs.
• Backward pass: 8.0 × 10¹⁴ FLOPs.
• Total per batch: 1.2 × 10¹⁵ FLOPs.

---

5. Total Training FLOPs

• Number of batches: Total tokens / tokens per batch = 3 × 10¹¹ / 3.2 × 10⁶ ≈ 9.4 × 10⁴ batches.
• Total FLOPs: 9.4 × 10⁴ × 1.2 × 10¹⁵ ≈ 1.1 × 10²⁰ FLOPs.

However, this is an underestimate because:

1. Optimizer overhead: Adam requires storing and updating momentum and variance, which adds ~2× more FLOPs (since each parameter update involves 2 additional operations).
   • Total FLOPs with Adam: ~3 × forward + backward FLOPs ≈ 3 × 1.2 × 10¹⁵ = 3.6 × 10¹⁵ per batch.
   • Total training FLOPs: 9.4 × 10⁴ × 3.6 × 10¹⁵ ≈ 3.4 × 10²⁰ FLOPs.
2. Mixed precision training: Often uses FP16/FP32, which may reduce FLOPs slightly, but we'll ignore this for simplicity.

---

6. Refining the Estimate

A more precise formula for transformer FLOPs (from Kaplan et al., 2020) is: [ \text{Total FLOPs} = 6 \times N \times D ] where:

• ( N ) = number of parameters (~1.75 × 10¹¹).
• ( D ) = number of training tokens (~3 × 10¹¹).

This gives: 6 × 1.75 × 10¹¹ × 3 × 10¹¹ ≈ 3.15 × 10²³ FLOPs.

This is much larger than our previous estimate because:

• The formula accounts for the fact that each parameter is updated for every token (not just per batch).
• It includes the full backward pass and optimizer overhead.

However, this seems too high. A better approximation is: [ \text{Total FLOPs} \approx 2 \times N \times D \times \text{seq_len} ] But this is also not precise.

Empirical Scaling Laws:

From Chinchilla scaling laws, the FLOPs for training a transformer can be approximated as: [ \text{FLOPs} \approx 6 \times N \times D ] For GPT-3: 6 × 1.75 × 10¹¹ × 3 × 10¹¹ ≈ 3.15 × 10²³ FLOPs.

But this is likely an overestimate. The actual number is closer to 3 × 10²³ FLOPs (as reported in the GPT-3 paper).

---

7. Final Estimate

Combining all the above, the most reasonable estimate is: [ \text{Total training FLOPs for GPT-3} \approx 3.14 \times 10^{23} \text{ FLOPs} ] This aligns with:

• The 6 × N × D rule of thumb.
• Reported estimates for GPT-3 (e.g., ~3.14 × 10²³ FLOPs in the paper).

---

Summary of Steps:

1. Identify model parameters (N = 1.75 × 10¹¹).
2. Identify training tokens (D = 3 × 10¹¹).
3. Use the empirical scaling law: FLOPs ≈ 6 × N × D.
4. Compute: 6 × 1.75 × 10¹¹ × 3 × 10¹¹ ≈ 3.15 × 10²³ FLOPs.

This is the most reliable estimate without looking up the exact number. The actual number may vary slightly due to implementation details (e.g., mixed precision, gradient checkpointing).