概要

NVIDIAの大規模言語モデル学習フレームワークMegatron-LMを、GH200上でビルドした手順をまとめます。

通常はDockerやSingularityなどのコンテナ上でのビルドが想定されていますが、今回はそれを無理やり実環境でやった記録になります。同じ環境で苦労している方の参考になれば幸いです。

環境

項目	値
CPU	ARM Grace (aarch64)
GPU	GH200 (Hopper, sm_90)
CUDA Toolkit	12.9
NVIDIA HPC SDK	25.9
Python (system)	3.9.25
uv	0.9.30

ビルド対象

パッケージ	バージョン	commit id
megatron-core	0.16.0	b0cc2706ddc60d2aefd5fff346445b5c013036a8
torch	2.10.0+cu129

ビルドスクリプト

最終的に動作確認済みのスクリプトです。詳しい経緯はつまずき解説のセクションで説明します。

build-megatron-lm.sh

#!/bin/bash
# build-megatron-lm.sh  Megatron-LM(megatron-core)
# 使い方: bash build-megatron-lm.sh [MEGATRON_LM_DIR]
set -euo pipefail

MEGATRON_LM_DIR="$(cd "${1:-$(pwd)/Megatron-LM}" && pwd)"
PYTHON_VERSION="3.12"
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"

[[ -f "$MEGATRON_LM_DIR/pyproject.toml" ]] \
    || { echo "ERROR: pyproject.toml not found: $MEGATRON_LM_DIR" >&2; exit 1; }

module purge 2>/dev/null || true
module load gcc
module load cuda/12.9
module load cudnn/9.10.1.4

export MAX_JOBS="${MAX_JOBS:-$(nproc)}"

rm -rf "$MEGATRON_LM_DIR/.venv" # 環境依存をなくすために毎回削除
uv venv --python "$PYTHON_VERSION" "$MEGATRON_LM_DIR/.venv"
source "$MEGATRON_LM_DIR/.venv/bin/activate"

uv pip install torch torchvision --index-url "$TORCH_INDEX_URL"
uv pip install ninja pybind11 nvidia-mathdx wheel "setuptools<80.0.0" "packaging>=24.2"

TORCH_LIB="$(python -c 'import torch; print(torch.__path__[0])')/lib"
export LD_LIBRARY_PATH="$TORCH_LIB:${LD_LIBRARY_PATH:-}"
export LIBRARY_PATH="$TORCH_LIB:${LIBRARY_PATH:-}"

# nvidia-nccl-cu12 内の nccl.h をコンパイラに見せる
NCCL_INCLUDE="$(python -c 'import nvidia.nccl, pathlib; print(pathlib.Path(nvidia.nccl.__path__[0]) / "include")')"
export CPATH="$NCCL_INCLUDE:${CPATH:-}"

# apex をインストールする (torch と CUDA が必要)
uv pip install --no-build-isolation git+https://github.com/NVIDIA/apex.git

cd "$MEGATRON_LM_DIR"
# torch などのパッケージに依存したものをインストールする時に隔離されないよう --no-build-isolation-package を追加
uv pip install -e ".[mlm,dev]" \
    --no-build-isolation-package transformer-engine \
    --no-build-isolation-package transformer-engine-torch \
    --no-build-isolation-package mamba-ssm \
    --no-build-isolation-package causal-conv1d \
    --no-build-isolation-package nv-grouped-gemm \
    --no-build-isolation-package flash_mla \
    || true

echo "Done: $(python -c 'import megatron.core; print(f"megatron-core {megatron.core.__version__}")')"

つまずいた箇所

uv pip install torch が "Audited 1 package" で何もしない

uv pip install torch --index-url ... を実行すると、ダウンロードが走らずに Audited 1 package と表示されて終わることがあります。先に torchaudio をインストールしていると、uv が依存関係を「解決済み」と誤判断するのが原因です。--reinstall を付けることで回避できます。

$ uv pip install --reinstall torch torchvision \
    --index-url https://download.pytorch.org/whl/cu129

NVIDIA HPC SDK (nvc++) と PyTorch が衝突する

import megatron.core を実行すると、以下のエラーが出ることがあります。

Your compiler (nvc++) is not compatible with the compiler Pytorch was
built with for this platform, which is g++ on linux.

Miyabi ではデフォルトで nvidia/25.9 モジュールがロードされており、C++ コンパイラが nvc++ (NVIDIA HPC Compiler) になっています。PyTorch の C++ 拡張は g++ を前提としているため、両者が衝突してしまいます。module purge で HPC SDK をアンロードし、CUDA Toolkit だけをロードすることで解消します。

$ module purge
$ module load cuda/12.9

$ which c++
/usr/bin/g++
$ echo $CUDA_HOME
/work/opt/local/aarch64/cores/cuda/12.9

GPU のないログインノードで arch 検出が失敗する

import megatron.core を実行すると、以下のエラーが出ることがあります。

File "torch/utils/cpp_extension.py", line 2533, in _get_cuda_arch_flags
    arch_list[-1] += '+PTX'
IndexError: list index out of range

Megatron-LM の初回インポート時に unified_memory.py が C++ コードを JIT コンパイルしようとします。その際 PyTorch が GPU アーキテクチャを自動検出しようとしますが、ログインノードには GPU がないため arch_list が空になってしまいます。環境変数で明示することで問題を回避できます。

$ export TORCH_CUDA_ARCH_LIST="9.0"  # GH200 = Hopper (sm_90)

GPU があるノードでは自動検出されるため設定は不要ですが、ログインノードとの差異をなくすために常に設定しておくと良いと思います。

実行ログ

クリーンな環境で build-megatron-lm.sh を実行した際のログです（GPU ノード上）。

$ bash build-megatron-lm.sh ./Megatron-LM

Using CPython 3.12.12
Creating virtual environment at: Megatron-LM/.venv
Activate with: source Megatron-LM/.venv/bin/activate

# torch インストール
Using Python 3.12.12 environment at: Megatron-LM/.venv
Resolved 31 packages in 3.02s
Installed 31 packages in 11.83s
 + torch==2.10.0+cu129
 + torchvision==0.25.0+cu129
 + triton==3.6.0
 + nvidia-cudnn-cu12==9.10.2.21
 + nvidia-nccl-cu12==2.27.5
 ...

# ビルドツール
Resolved 6 packages in 5ms
Installed 5 packages in 1.55s
 + ninja==1.13.0
 + nvidia-mathdx==25.6.0
 + packaging==26.0
 + pybind11==3.0.2
 + wheel==0.46.3

# apex
Resolved 2 packages in 469ms
Installed 1 package in 174ms
 + apex==0.1 (from git+https://github.com/NVIDIA/apex.git@1f7b2922...)

# megatron-core + 全依存パッケージ
Resolved 145 packages in 2.93s
   Building megatron-core @ file:///.../Megatron-LM
      Built megatron-core @ file:///.../Megatron-LM
Prepared 1 package in 8.58s
Uninstalled 3 packages in 297ms
Installed 134 packages in 17.42s
 + megatron-core==0.16.0 (from file:///.../Megatron-LM)
 + transformer-engine==2.12.0+5671fd36 (from git+https://...)
 + causal-conv1d==1.5.3
 + mamba-ssm==2.2.5
 + nv-grouped-gemm==1.1.4.post8
 + flash-linear-attention==0.4.1
 + flashinfer-python==0.5.3
...(その他)

Done: megatron-core 0.16.0

所要時間は約 1.5 分です（Python・依存パッケージのダウンロード含む、GPU ノード）。

使い方

ビルド後にセッションを開き直した場合のセットアップは以下の通りです。

$ module purge
$ module load cuda/12.9

$ export TORCH_CUDA_ARCH_LIST="9.0"

$ source /path/to/Megatron-LM/.venv/bin/activate

$ python -c "import megatron.core; print(megatron.core.__version__)"
# => 0.16.0

次の記事では、この環境の上に Megatron-Bridge をビルドする手順を解説します。
https://qiita.com/aida0710/items/2df9e9205997d0ef08bb