先定坐标。这里的 3D parallelism 指 Tensor Parallel × Pipeline Parallel × Data Parallel。 加上长序列维度后,常见 4D 形态是 TP × PP × DP × CP。 Sequence Parallel 通常依附在 TP 组内切 activation 序列维度,不一定让 world size 多乘一个新维度; Context Parallel 则常作为独立的 context_parallel_size 进入并行乘积。

world size 约束

world_size = tensor_model_parallel_size
           * pipeline_model_parallel_size
           * data_parallel_size
           * context_parallel_size   # 如果启用 CP

# 单机 8 卡例子
TP=2, PP=2, CP=1  => DP = 8 / (2*2*1) = 2
TP=4, PP=2, CP=1  => DP = 8 / (4*2*1) = 1
TP=2, PP=2, CP=2  => DP = 8 / (2*2*2) = 1

假设 world_size=8、TP=2、PP=2、DP=2。 一个典型 Megatron rank order 下,TP group 是相邻 rank,PP group 跨 pipeline stage, DP group 是同一 TP/PP 位置上的数据副本。你不需要死背顺序,但必须能看出每个 collective 发生在哪个 group。

Megatron-LM 版本迭代很快,命令细节可能随分支变化。 今天使用的是经典 pretrain_gpt.py 教学路径: 先准备一个极小文本数据集,用 Megatron 的预处理脚本生成 indexed dataset, 再用 torchrun 或 Megatron 启动脚本跑小 GPT。 如果你的分支提供了新的 recipe 脚本,以仓库 README 和脚本参数为准。

# 1. 获取代码。官方仓库包含 Megatron-LM 示例和 Megatron Core
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM

# 2. 安装。官方 README 当前推荐 uv pip install -e .
python -m pip install uv
uv pip install -e .

# 3. 检查 GPU 与 PyTorch
python - <<'PY'
import torch
print(torch.__version__)
print(torch.cuda.device_count())
print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "no cuda")
PY

准备一个 tiny corpus

mkdir -p data/tiny
cat > data/tiny/corpus.jsonl <<'EOF'
{"text": "AI infrastructure connects models, data, kernels, networks, and clusters."}
{"text": "Tensor parallelism splits matrices inside a transformer layer."}
{"text": "Pipeline parallelism splits transformer layers into stages."}
{"text": "Data parallelism replicates workers and synchronizes gradients."}
{"text": "Context parallelism helps long sequence training scale across devices."}
EOF

# GPT-2 BPE tokenizer files, used by the classic GPT preprocessing path
wget -O data/tiny/gpt2-vocab.json https://huggingface.co/gpt2/resolve/main/vocab.json
wget -O data/tiny/gpt2-merges.txt https://huggingface.co/gpt2/resolve/main/merges.txt

python tools/preprocess_data.py \
  --input data/tiny/corpus.jsonl \
  --output-prefix data/tiny/gpt_text \
  --tokenizer-type GPT2BPETokenizer \
  --vocab-file data/tiny/gpt2-vocab.json \
  --merge-file data/tiny/gpt2-merges.txt \
  --append-eod \
  --json-keys text

先跑最简单的 DP-only baseline,确认数据、tokenizer、模型配置和 launcher 都没问题。 下面的模型故意很小:8 层、hidden 512、8 个 attention heads。 这样 TP=1/2/4/8 都能整除 head 数,PP=1/2/4 也能整除层数,方便做实验矩阵。

cat > run_tiny_gpt.sh <<'EOF'
#!/usr/bin/env bash
set -euo pipefail

TP=${TP:-1}
PP=${PP:-1}
CP=${CP:-1}
GPUS=${GPUS:-8}
SEQ=${SEQ:-512}
MB=${MB:-2}
GB=${GB:-32}
ITERS=${ITERS:-50}

DATA_PATH=${DATA_PATH:-data/tiny/gpt_text_text_document}
VOCAB_FILE=${VOCAB_FILE:-data/tiny/gpt2-vocab.json}
MERGE_FILE=${MERGE_FILE:-data/tiny/gpt2-merges.txt}

COMMON_ARGS=(
  --tensor-model-parallel-size ${TP}
  --pipeline-model-parallel-size ${PP}
  --num-layers 8
  --hidden-size 512
  --ffn-hidden-size 2048
  --num-attention-heads 8
  --seq-length ${SEQ}
  --max-position-embeddings ${SEQ}
  --micro-batch-size ${MB}
  --global-batch-size ${GB}
  --train-iters ${ITERS}
  --lr 1.0e-4
  --min-lr 1.0e-5
  --lr-decay-style cosine
  --lr-decay-iters ${ITERS}
  --lr-warmup-fraction 0.01
  --weight-decay 0.1
  --clip-grad 1.0
  --tokenizer-type GPT2BPETokenizer
  --vocab-file ${VOCAB_FILE}
  --merge-file ${MERGE_FILE}
  --data-path ${DATA_PATH}
  --split 949,50,1
  --log-interval 1
  --eval-interval 25
  --eval-iters 5
  --save-interval 1000
  --no-save-optim
  --no-save-rng
  --bf16
)

if [[ "${CP}" != "1" ]]; then
  COMMON_ARGS+=(--context-parallel-size ${CP})
fi

torchrun --standalone --nproc_per_node=${GPUS} pretrain_gpt.py "${COMMON_ARGS[@]}"
EOF
chmod +x run_tiny_gpt.sh

# baseline: TP=1, PP=1, DP=8
GPUS=8 TP=1 PP=1 ./run_tiny_gpt.sh

RoadMap 点名要求改 --tensor-model-parallel-size 和 --pipeline-model-parallel-size。 这一步不要只看能不能跑,要写下你的解释: TP 增大后单层 GEMM 被切小,通信更频繁; PP 增大后每卡层数减少,但 pipeline bubble 和 send/recv 增加。

自动跑矩阵并保存日志

mkdir -p logs/day22

for cfg in "1 1 1" "2 1 1" "4 1 1" "2 2 1" "1 4 1" "2 2 2"; do
  read -r TP PP CP <<< "${cfg}"
  name="tp${TP}_pp${PP}_cp${CP}"
  echo "running ${name}"
  GPUS=8 TP=${TP} PP=${PP} CP=${CP} ITERS=30 \
    ./run_tiny_gpt.sh 2>&1 | tee "logs/day22/${name}.log"
done

这类实战最容易变成"跑了一堆命令,然后不知道学到了什么"。 你要把每次配置拆成三层判断: 模型是否能放下,计算是否仍然高效,通信是否被隐藏。 TP/PP 不是越大越好,更不是卡数越多越好。

记录模板

| config | TP | PP | CP | DP | micro-batches | max mem/rank | step time | tokens/s | 观察 |
|--------|----|----|----|----|---------------|--------------|-----------|----------|------|
| A      | 1  | 1  | 1  | 8  |               |              |           |          |      |
| B      | 2  | 1  | 1  | 4  |               |              |           |          |      |
| C      | 4  | 1  | 1  | 2  |               |              |           |          |      |
| D      | 2  | 2  | 1  | 2  |               |              |           |          |      |
| E      | 1  | 4  | 1  | 2  |               |              |           |          |      |
| F      | 2  | 2  | 2  | 1  |               |              |           |          |      |

1. 给定 world_size=8、TP=2、PP=2、CP=1,计算 DP size,并画出 TP group、PP group、DP group。
2. 解释 --tensor-model-parallel-size 增大后,单层 GEMM、参数显存和通信频率分别如何变化。
3. 解释 --pipeline-model-parallel-size 增大后,每卡层数、activation 通信和 pipeline bubble 分别如何变化。
4. 为什么 num_attention_heads 要能被 TP size 整除?为什么 num_layers 最好能被 PP size 整除?
5. 用 Day20 的 bubble 公式,比较 PP=2 和 PP=4 在 micro-batches=8 时的 bubble fraction。
6. 从实验 A-F 中选出你认为最合理的一组配置,用显存、吞吐、通信、复杂度四个角度说明理由。