nanoGPT

To reproduce GPT-2 results, follow these steps to tokenize OpenWebText and run distributed training using torchrun.

Single Node (8x A100): Use torchrun with --nproc_per_node=8 to utilize all GPUs on a single node.

Multi-Node Cluster: When training across multiple nodes, you must specify the --nnodes, --node_rank, and the --master_addr of the first node. If you do not have Infiniband, prepend NCCL_IB_DISABLE=1 to prevent slow interconnect performance.

# 1. Prepare OpenWebText
python data/openwebtext/prepare.py

# 2. Single node (8 GPUs)
torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py

# 3. Multi-node (Example: 2 nodes)
# On Master Node (rank 0):
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr=123.456.123.456 --master_port=1234 train.py

# On Worker Node (rank 1):
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr=123.456.123.456 --master_port=1234 train.py

You can estimate the total number of parameters in a Transformer model by summing the parameters in the embeddings, attention blocks, MLP blocks, and the final layer normalization.

Note that in nanoGPT, the dense layer (the final linear projection) typically shares weights with the embedding layer, resulting in 0 additional parameters for that specific component in the calculation.

Key components included in the estimate:

• Embeddings: Token and position embeddings.
• Attention Blocks: LayerNorm, K/Q/V projections, and the output projection.
• MLP Blocks: LayerNorm, the feed-forward expansion, and the projection back to n_embd.
• Transformer: The sum of all blocks multiplied by n_layer.
• Final LayerNorm: ln_f.

def params():
    """ estimates the number of parameters in the model"""
    out = OrderedDict()

    # token and position embeddings
    out['emebedding/position'] = n_embd * block_size
    out['embedding/token'] = n_embd * vocab_size
    out['embedding'] = out['emebedding/position'] + out['embedding/token']

    # attention blocks
    out['attention/ln'] = n_embd
    out['attention/kqv'] = n_embd * 3*n_embd
    out['attention/proj'] = n_embd**2
    out['attention'] = out['attention/ln'] + out['attention/kqv'] + out['attention/proj']

    # MLP blocks
    ffw_size = 4*n_embd
    out['mlp/ln'] = n_embd
    out['mlp/ffw'] = n_embd * ffw_size
    out['mlp/proj'] = ffw_size * n_embd
    out['mlp'] = out['mlp/ln'] + out['mlp/ffw'] + out['mlp/proj']

    # the transformer and the rest of it
    out['block'] = out['attention'] + out['mlp']
    out['transformer'] = n_layer * out['block']
    out['ln_f'] = n_embd
    out['dense'] = 0 

    out['total'] = out['embedding'] + out['transformer'] + out['ln_f'] + out['dense']
    return out

To estimate the storage required for a model checkpoint and the memory required for training, consider the following:

1. Checkpoint Size: Parameters are typically stored in fp32 (4 bytes). If using the AdamW optimizer, you must account for 2 additional buffers per parameter for optimizer statistics (momentum and variance).

   • params_bytes = total_params * 4
   • params_and_buffers_bytes = params_bytes + 2 * params_bytes

2. GPU Memory: While parameters and optimizer buffers occupy significant space, for many models, the majority of GPU memory is consumed by activations during the forward and backward passes. This ratio increases significantly as models grow larger.

# Estimate checkpoint size (fp32 + AdamW buffers)
params_bytes = params_total * 4
params_and_buffers_bytes = params_bytes + 2 * params_bytes
print(f"est checkpoint size: {params_and_buffers_bytes/1e9:.2f} GB")

# Estimate GPU memory usage for parameters/buffers
gpu_memory = 40e9 # e.g., 40 GB A100
print(f"memory ratio taken up just for parameters: {params_and_buffers_bytes / gpu_memory * 100:.2f}%")

Estimating FLOPs (Floating Point Operations) is critical for calculating Model FLOPs Utilization (MFU). This notebook uses a method that counts actual FLOPs (not MACs), meaning matrix multiplications of $(B imes C) imes (C imes D)$ are calculated as $2 imes B imes C imes D$.

Key FLOP components:

• Attention: K/Q/V projection, attention score calculation, value reduction, and final projection.
• MLP: Feed-forward expansion and projection.
• Total Pass: The sum of forward and backward passes. A common estimate for the backward pass is $2 imes ext{forward pass}$ cost.

PaLM Formula: For a more standardized estimate, the PaLM paper formula can be used: $mf_{ ext{per token}} = 6N + 12LHQ T$ where $N$ is non-embedding parameters, $L$ is layers, $H$ is heads, $Q$ is head dimension, and $T$ is block size.

def flops():
    # ... (implementation details) ...
    # 1) the projection to key, query, values
    out['attention/kqv'] = 2 * block_size * (n_embd * 3*n_embd)
    # 2) calculating the attention scores
    out['attention/scores'] = 2 * block_size * block_size * n_embd
    # 3) the reduction of the values
    out['attention/reduce'] = 2 * n_head * (block_size * block_size * head_size)
    # 4) the final linear projection
    out['attention/proj'] = 2 * block_size * (n_embd * n_embd)
    # ... (MLP and total calculations) ...
    return out