This note is to review an AI optmization specifically in a low-level to prepare an interview.

Preill and Decode

  • Prefill: compute-bound
  • Decode: memory-bound

Parallelism

  • Tensor Parallelism (TP)
    • split a model into multiple gpus
  • Data Parallelism (DP)
    • multiple models, forwarding a data in parallel
    • ex. K and V are duplicated across gpus
    • However, ZeRo-DP (Zero Redundancy Optimizer for Data Parallelism) is a memory-saving technology, comprising of three stages (optimizer, gradient, and parameter partitioning)
      • https://share.google/aimode/Z8PucMgRKs879XQQA
      • Due to this (communication bottleneck derived by TP), DeepSeek also adopted EP and ZeRo-DP without TP.
  • Pipeline Parallelism
  • Expert Parallelism (EP)
    • Sparse Attention (MoE)

TP 3 and DP 2 means it will provide an inference system within three gpus for each model and two models simultaneously. Therefore, six gpus are required.

cf. In FP8, 70B models = 70 GB GPU memory

AI Optimization 1: Low-level

Distillation, pruning and quantization

  • Distillation
  • Pruning (might cut a wrong, duplicated, or outdated knowledge)
    • Structured pruning
      • Magnitude pruning
      • Activation-based pruning
    • Unstructured pruning
      • Dept-based pruning
      • Width-based pruning
  • Quantization
    • Use both pruning and distillation in a quantization setting. It will make more efficient
    • Mix precision -> Activation-aware Quantization (AWQ)
    • Dynamic range options
      • Per-tensor (layer): good for hardware efficiency, but subject to outliers
      • Per-channel: good for quality, but hardware burden

Kernel

  • T=N; T>N; T<N; and Grid-stride loop
  • cudaMallocManaged, cudaMemcpy, cudaMemcpyHostToDevice, cudaMemcpyDeviceToHost
  • cudaStream
  • someKernel«<blocks, threads, sharedMemSize(0), stream»>
    • overlapping data transfer and kernel execution
  • Map, reduction, scan/prefix, and sort

AI Optimization 2

Attention families

Self-attention

Architecture-level

  • Multi-head attention
  • Multi-query attention
    • saves memory, same compute, less quality
  • Grouped-query attention
    • saves memory, same compute (comparing with multi-head)
    • the middle point

Attentions

Math-based

  • Paged attention
    • naive -> block -> paged
  • Flash attention
    • chain rule