mHC 与 TE 在 H200 上的环境排障#

最短成功路径#

先核对本文的环境记录，再依次完成以下步骤。路径与配置来自本次 DSW 运行记录，换环境时需要重新验证。

1. 建立独立 Venv，核对实际 Python 与依赖环境。
2. 按顺序安装依赖，再编译 CUDA extension。
3. 检查配置组合，结合最小配置核对开关。
4. 完成启动前检查，再按启动命令运行。
5. 用性能记录对照本次配置，出现异常时回查对应步骤中的故障说明。

记录把 examples/apodex/a400m_base_8b.yaml (TE + MLA + mHC + 8B 切片) 从 0 跑通的全部坑。下次新会话 / 新人接手按这份顺序排查能少走弯路。

1. 环境信息（关键背景）#

判定 GPU 唯一可信源：torch.cuda.get_device_properties(0).major/minor，别信 nvidia-smi --query-gpu=compute_cap。

含义：

• 编 CUDA extension（apex / cpp_extension 系）必须 TORCH_CUDA_ARCH_LIST="9.0"
• Megatron 报的 MFU 是 ÷ 670 (L20X peak)，实际硬件 MFU 要 ÷ 989 (H200 peak)，再乘 670/989 ≈ 0.68 折算

2. Venv 配置#

2.1 Venv 选择#

.venv-megatron-0521（py3.12.3，uv 创建）是当前可用 venv。

2.2 用 VENV= 切 venv 而不是改 .venv 软链#

start.sh:27 加了 hook：

source "${VENV:-.venv}/bin/activate"

用法：

VENV=.venv-megatron-0521 MCORE_DPA_USE_SDPA=... examples/apodex/start.sh <yaml>

2.3 Venv 复制陷阱#

如果你从别的 venv（如 /mnt/pre_train/public/venv-megatron-te2141-nvrx）复制 venv 内容过来，bin/activate 和所有 entry-point script 的 shebang 里 VIRTUAL_ENV 是写死的绝对路径，需要 sed 替换：

find <new_venv>/bin/ -type f -exec grep -l <old_venv_path> {} \; | \
  xargs sed -i "s|<old_venv_path>|<new_venv_path>|g"

正解：用 uv venv 现造，不要复制。

3. 装包顺序（重要）#

3.1 uv sync 处理 lock + project deps#

# 注意:UV_PROJECT_ENVIRONMENT 若设了会盖掉 VIRTUAL_ENV
unset UV_PROJECT_ENVIRONMENT
VIRTUAL_ENV=.venv-megatron-0521 uv sync --no-group test --no-group docs

坑 1：pyproject.toml 里 test group 拉 nemo-run 这个 git dep，上游 pyproject 格式错让 build 挂。用 --no-group test --no-group docs 跳过。

坑 2：repo 根有 pyproject.toml + uv.lock，uv pip install 在 repo 根目录被视作 project mode，静默 no-op。要装包必须 cd /tmp 切出 project root，或者用 --active。

3.2 torch（uv sync 不会装）#

uv.lock 有 manifest override：

[manifest]
overrides = [
    { name = "torch", marker = "sys_platform == 'never'" },
]

uv sync 在任何 sys_platform 都 skip torch（假设 image 提供）。DSW 上得手装：

cd /tmp && unset UV_PROJECT_ENVIRONMENT && \
  VIRTUAL_ENV=.venv-megatron-0521 UV_HTTP_TIMEOUT=600 \
  uv pip install torch==2.11.0+cu130 \
    --index-url https://download.pytorch.org/whl/cu130 \
    --extra-index-url https://pypi.nvidia.com

坑：pypi.nvidia.com 经常 timeout（nvidia-cuda-* 这堆运行时包从那拉），用 UV_HTTP_TIMEOUT=600 给足重试。

3.3 nvidia-cublas 升级（TE 2.14 要新 cublasLt symbol）#

RuntimeError: undefined symbol: cublasLtGroupedMatrixLayoutInit_internal, version libcublasLt.so.13

torch 2.11 pin 了 nvidia-cublas==13.1.0.3，但 TE 2.14.1 要求 13.4.1.1+ 才有 grouped matrix init symbol。手动升：

cd /tmp && unset UV_PROJECT_ENVIRONMENT && \
  VIRTUAL_ENV=.venv-megatron-0521 \
  uv pip install --upgrade nvidia-cublas --extra-index-url https://pypi.nvidia.com

3.4 transformers 缺隐藏依赖#

File "huggingface_tokenizer.py", line 91:
ValueError: ... Exception: name 'AutoTokenizer' is not defined

try: from transformers import AutoTokenizer / except ModuleNotFoundError 吞了 regex 包缺失的 ImportError，AutoTokenizer 没定义。uv pip install --reinstall transformers 拉齐所有 deps（regex / tokenizers / safetensors / tqdm / shellingham 等）。

3.5 wandb 链式 deps（click / platformdirs / requests / urllib3）#

也是 --reinstall wandb 一次性补齐。

3.6 装 pip（apex build 需要）#

VIRTUAL_ENV=.venv-megatron-0521 uv pip install pip

uv-managed venv 默认不带 pip。uv pip install 不认 apex 要的 --build-option 参数（这是 pip 自己的），所以编 apex 必须有 pip。

4. 编 CUDA extension#

4.1 helpers_cpp (Megatron 自带 pybind11 module)#

仓里预装的 megatron/core/datasets/helpers_cpp.cpython-312-*.so 跟 py3.12 venv 直接 ABI 兼容不用重编。

如果换 py 版本（比如 py3.11 venv），现成 .so 不 work，要重编：

EXT_SUFFIX=$(.venv-X/bin/python -c "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX'))")
INCS=$(.venv-X/bin/python -m pybind11 --includes)
cd megatron/core/datasets && \
  g++ -O3 -Wall -shared -std=c++17 -fPIC -fdiagnostics-color \
    $INCS helpers.cpp -o "helpers_cpp${EXT_SUFFIX}"

绕开 Makefile 那个 python3-config 依赖（uv venv 里不存在）。

4.2 apex (FusedAdam / fused_weight_gradient_mlp_cuda 等)#

编译时长：H200 单卡 / 16 cores parallel，~15 分钟全部 extension。

git clone --depth 1 https://github.com/NVIDIA/apex.git /tmp/apex
cd /tmp/apex && \
  VIRTUAL_ENV=.venv-megatron-0521 \
  TORCH_CUDA_ARCH_LIST="9.0" \
  MAX_JOBS=$(nproc) NVCC_THREADS=2 \
  .venv-megatron-0521/bin/python -m pip install -v \
    --no-build-isolation --no-cache-dir \
    --config-settings "--build-option=--cpp_ext --cuda_ext" \
    .

关键：TORCH_CUDA_ARCH_LIST="9.0"，不是 8.9。apex 默认从 torch.cuda.get_arch_list() 挑最高 SM 编。torch 2.11+cu130 的 arch list 是 [75, 80, 86, 90, 100, 120]，没 89。挑 90 是正确的（硬件是 H200）。

踩过的坑：

• 一开始信了 nvidia-smi 的 8.9，强制 TORCH_CUDA_ARCH_LIST=8.9，build 出来的 kernel 跑不动（PTX 89 倒还行，但默认 sm_89 SASS 在 H200 上没意义）。重新 build sm_90 才对。

5. yaml: mHC + TE 互斥矩阵#

参考 PR #99 / #100 / #101。当前 a400m_base_8b.yaml 通过的 mutex 检查：

5.1 触发 mHC bug 的两个 fine-grained schedule flag#

# 跟 mHC 在本分支冲突,触发 combined_1f1b_schedule_for_no_pipelining 路径
# 绕过 transformer_block.forward 的 HyperConnectionModule.input_expand
# → mapping_proj weight (5120) vs hidden_states (1280) shape mismatch
# --overlap-moe-expert-parallel-comm  # 关
# --delay-wgrad-compute               # 关

错误现场：

RuntimeError: mat1 and mat2 shapes cannot be multiplied (16384x1280 and 5120x24)

5120 = 4 × 1280 = num_residual_streams × hidden_size —— mHC 期望 hidden 已经 expand 到 n*C，但 fine-grained schedule 走的 f_layer.attn.forward 直接接 raw [s, b, C]。

5.2 MoE token dispatcher + CUDA Graph 互斥#

token_dispatcher._maybe_dtoh_and_synchronize 在 capturing stream 里 .synchronize() 会触发 cudaErrorCapturedEvent。原 a400m_base.yaml 用 scope=attn 只 capture attention 避开。mHC 用 scope=moe_router 是专门为 HC partial-replay 设计的，由 HyperConnectionTransformerLayer._te_cuda_graph_replay 重写处理。

6. apex 没装就关 --no-gradient-accumulation-fusion#

RuntimeError: ColumnParallelLinear was called with gradient_accumulation_fusion set to True 
but the custom CUDA extension fused_weight_gradient_mlp_cuda module is not found.

gradient_accumulation_fusion 默认 True，需要 apex 的 cu 扩展。装了 apex 就可以删掉这个 flag，让 fusion 生效（几个百分点 perf）。

7. 跑前 nvrx 检查#

start.sh:128-139 的 nvrx 存在性检查改用 python -c "import" 不是 pip show，因为 uv venv 没 pip。

如果 venv 没 nvrx：

cd /tmp && VIRTUAL_ENV=.venv-megatron-0521 \
  uv pip install nvidia-resiliency-ext

只是 async ckpt 用，不装的话有 warning 但不影响训练。

8. 最小可工作的 a400m_base_8b.yaml 关键 flag 集#

# 架构 (MLA, TE-native)
--multi-latent-attention
--q-lora-rank 768
--kv-lora-rank 256
--qk-head-dim 128
--qk-pos-emb-head-dim 64
--v-head-dim 128

# TE backend (前提:venv 有 TE + apex + cublas 13.4+)
--transformer-impl transformer_engine
--attention-backend flash
--te-rng-tracker

# mHC (PR #99 + Triton fused PR #100)
--enable-hyper-connections
--num-residual-streams 4
--use-fused-mhc

# CUDA Graph (mHC 专用 partial-replay scope, PR #101)
--cuda-graph-impl transformer_engine
--cuda-graph-scope moe_router
--moe-pad-experts-for-cuda-graph-inference

# 其他 TE fused kernels (跟 mHC 不冲突)
--moe-permute-fusion
--moe-router-fusion
--cross-entropy-loss-fusion
--cross-entropy-fusion-impl te

# ❌ 不能加 (跟 mHC 互斥):
# --overlap-moe-expert-parallel-comm
# --delay-wgrad-compute
# --cuda-graph-scope attn / full-layer (default) / mlp / moe / moe_preprocess

9. 启动命令#

cd /mnt/VerifiableAILab/zenan.li/Megatron-LM
VENV=.venv-megatron-0521 \
  examples/apodex/start.sh examples/apodex/a400m_base_8b.yaml 2>&1 | tee /tmp/a400m_8b.log

不要设 MCORE_DPA_USE_SDPA=1 —— 那是 _local 路径（local + SDPA + 绕 NVIDIA #1589）的开关，TE 路径不需要。

10. 性能基线（DSW 8×H200 单节点）#

mHC + 丢失 2 个 overlap flag 大约吃掉 35% 的 baseline MFU。

11. 附：调试 import chain 缺包的统一方法#

Megatron 里很多 try: import X except ImportError 会吞错让 X 变 None，调用时报 name 'X' is not defined。逐层补：

.venv-megatron-0521/bin/python -c "import <module>"

直接 import 看错信息，缺啥装啥。常见缺包候选：regex, click, platformdirs, requests, urllib3, tokenizers, safetensors, tqdm, pybind11, ninja。

复现范围：本篇是特定 H200 / DSW 环境的排障记录，不是通用安装指南。Python ABI 相同不足以单独保证二进制扩展兼容；还需匹配架构、编译器／libstdc++、CUDA 和相关框架 ABI。升级依赖前应保存版本清单，文中的 --reinstall 不保证恢复当时环境。