CS336 | SPACESHIP

Tokenization

Methods
- Character Tokenization: Unicode
- Byte Tokenization: UTF-8, vocabulary size is 256
- Word Tokenization: words are too many
Principle:
- Model (e.g. Transformer) should operate on chunks (abstractions) of the sequence (text, video, DNA, etc.) [OTHERWISE Low Signal-to-Noise]
- Chunks should be variable (allocate more model capacity to interesting chunks)

BPE

Basic idea: train the tokenizer on raw text to construct a vocabulary tailored to the data.
Intuition: common sequences of bytes are represented by a single token, rare sequences are represented by many tokens.
Sketch: start with each byte as a token (256), and successively merge the most common pair of adjacent tokens.
Combine word tokenization and byte tokenization.

Resource Accounting

Memory Types: fp32, fp16, bf16, fp8, fp4
Mixed precision training
- Use bf16 for parameters, activations, and gradients
- Use fp32 for optimizer states
Pytorch has an automatic mixed precision (AMP) library, which can try to cast things into bf16 when safe (matmuls, not exp).

einops

einsum

# Example 1
x = torch.ones(3, 4)
y = torch.ones(4, 3)
z = x @ y                                                  # Old way
z = einsum(x, y, "seq1 hidden, hidden seq2 -> seq1 seq2")  # New way

# Example 2
x = torch.ones(2, 3, 4)
y = torch.ones(2, 3, 4)
z = x @ y.transpose(-2, -1)                                                  # Old way
z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")  # New way
z = einsum(x, y, "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2")

reduce

x = torch.ones(2, 3, 4)
y = x.sum(dim=-1)
y = reduce(x, "... hidden -> ...", "sum") # we can also replace sum to mean, max, min

rearrange

x = torch.ones(3, 8)
w = torch.ones(4, 4)
x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)
x = einsum(x, w, "... hidden1, hidden1 hidden2 -> ... hidden2")
x = rearrange(x, "... heads hidden2 -> ... (heads hidden2)")

FLOP

A floating-point operation (FLOP) is a basic operation like addition (x+y) or multiplication (x y).
- FLOPs: floating-point operations (measure of computation done).
- FLOP/s: floating-point operations per second (also written as FLOPS), which is used to measure the speed of hardware.

Linear Model

if torch.cuda.is_available():
    B = 16384  # Number of points
    D = 32768  # Dimension
    K = 8192   # Number of outputs
else:
    B = 1024
    D = 256
    K = 64
x = torch.ones(B, D, device=cuda_if_available())
w = torch.randn(D, K, device=cuda_if_available())
y = x @ w
# actual_num_flops = 2 * B * D * K (while the correct num is B * (2D-1) * K)

Interpretation:
- B is the number of data points
- (D K) is the number of parameters
- FLOPs for forward pass is 2 (# tokens) (# parameters)

Time Calculation

def time_matmul(a: torch.Tensor, b: torch.Tensor) -> float:
    """Return the number of seconds required to perform 'a @ b'."""
    # Wait until previous CUDA threads are done
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    def run():
        a @ b
        if torch.cuda.is_available():
            torch.cuda.synchronize()
    num_trials = 5
    total_time = timeit.timeit(run, number=num_trials)
    return total_time / num_trials
actual_time = time_matmul(x, w)
actual_flop_per_sec = actual_num_flops / actual_time

# Note that the FLOP/s depends heavily on the data type
promised_flop_per_sec = get_promised_flop_per_sec(x.dtype)

mfu = actual_flop_per_sec / promised_flop_per_sec

Model FLOPs utilization (MFU)
- Definition: (actual FLOP/s) / (promised FLOP/s) [ignore communication/overhead]
- Usually, MFU >= 0.5 is quite good (and will be higher if matmuls dominate)
Bottleneck
- Memory-bound: communication time > computation time or arithmetic intensity < accelerator intensity
- Compute-bound: computation time > communication time or arithmetic intensity > accelerator intensity
  
  (h100_accelerator_intensity = h100_flop_per_sec / h100_bytes_per_sec = 295.3731; arithmetic_intensity = flops / bytes; intensity depends on precision)
- ReLU is memory-bound (where x.shape = n, bytes = 2 * 2 * n, flops = n)
- In order to improve arithmetic intensity, we can use gelu whose intensity is 5, but it is still smaller than 300. That’s to say that the choice of relu or gelu doesn’t influence the speed because the bottleneck is not it.
- Dot Product is memory-bound (~0.5)
```
n = 1024
x = torch.ones(n, dtype=torch.bfloat16, device=cuda_if_available())
w = torch.ones(n, dtype=torch.bfloat16, device=cuda_if_available())
y = x @ w

bytes = (2 * n) + (2 * n) + 2
flops = 2 * n - 1
```
- Matrix Vector Product is memory-bound (~1)
- Matmul is compute-bound (~n/3) => The larger the matrix, the higher the GPU utilization rate.
- Training transformer is compute-bound but the inference is memory-bound (matrix-vector product)
Backward Pass
- Backward pass is 2x more expensive than the forward pass which is 2 (# data points) (# parameters) FLOPs
- Total is 6 (# data points) (# parameters) FLOPs per training step
Optimizer: Customary to use fp32 for stability
Gradient accumulation, activation checkpointing: reduce memory to use bigger batch sizes.

Architectures

Pre-vs-post norm

removing warmup
stability and larger LRs for large networks

Architecture Update in 2024

LayerNorm vs RMSNorm

RMSNorm is faster (and just as good)
- Fewer operations (no mean calculation)
- Fewer parameters (no bias term to store)
More generally: dropping bias terms in FFN for memory and optimization stability

Gating

Gated activations (*GLU)
Archive

Serial vs Parallel layers

Position Embeddings

Position Embeddings 1

RoPE

Position Embeddings 2

Hyperparameters

$d_{ff}=4d_{model}$
$d_{model}=n_{heads}\cdot d_{heads}$
$\frac{d_{model}}{n_{layer}}=(100,200)$
Regularization: You still ‘regularize’ LMs but its effects are primarily on optimization dynamics.

Hyperparameters 1

Hyperparameters 1.1

Hyperparameters 1.2

Hyperparameters 2

Hyperparameters 3

Hyperparameters 4

Stability tricks

Softmax: can be ill-behaved due to exponentials / division by zero
- Output softmax stability - the ‘z-loss’
- Attention softmax stability - the ‘QK norm’
  - Logit soft-capping: Soft-capping the logits to some maximum value via Tanh

Attention heads

GQA/MQA - Reducing attention head cost
Sparse / Sliding Window Attention

Attention Alternatives

Linear Attention

Linear Attention 1

Linear Attention 2

SSM

Intuition: If we can regard linear attention as an RNN, then we can draw inspiration from LSTM to improve it.

Mamba-2

Gated delta net (and friends)

Gated_Delta_Net

DSA (Deepseek Sparse Attention)

Mixture of experts

Routing function
- Routing type
  - Token chooses expert (popular)
  - Expert chooses token
  - Global routing via optimization
- Common routing variants in detail
  - Top-k: Used in most MoEs; Switch Transformer(k=1)
    - shared expert
  - Hashing: Common baseline
  - RL to learn routes: Used in some of the earliest work; Bengio 2013, not common now
  - Solve a matching problem: Linear assignment for routing; Used in various papers like Clark 22
Training
- Major challenge: we need sparsity for training-time efficiency, but sparse gating decisions are not differentiable.
- Solutions:
  - Reinforcement learning to optimize gating policies
  - Stochastic perturbations
  - Heuristic ‘balancing’ losses (popular)
    - per-expert balancing
    - per-device balancing
    - auxiliary loss free balancing
    - load balancing
- the systems side
  - MoEs parallelize nicely - Each FFN can fit in a device
  - Hardware architecture design to use sparse MMs
- Stability
  - Solution: Use Float 32 just for the expert router (sometimes with an aux z-loss)
- fine-tuning
- upcycling: copy trained MLP as experts and add a router (example: MiniCPM, Qwen MoE)
Expert sizes
Training objectives