概要

前回の記事で Megatron-LM (megatron-core) を Miyabi GH200 上にビルドしました。本記事では、その上に Megatron-Bridge をビルドする手順をまとめます。

Megatron-Bridge は Megatron-LM ベースの LLM/VLM 学習レシピを提供するパッケージで、TransformerEngine や Mamba-SSM といった CUDA 拡張を含みます。そのぶんビルドにはいくつか追加の対応が必要になります。

前提

前回の記事のビルドが完了していること（torch 2.10.0+cu129 + megatron-core が venv にインストール済み）
Megatron-Bridge リポジトリがクローン済みであること

環境（前回と同じ）

項目	値
CPU	ARM Grace (aarch64)
GPU	GH200 (Hopper, sm_90)
CUDA	12.9
Python	3.12.12 (uv 管理)
torch	2.10.0+cu129（前回インストール済み）

ビルド対象

パッケージ	バージョン	備考
megatron-bridge	0.3.0	commit id 21b02e0cfb2f8ff907e0d8baee0c5205876c6812
transformer-engine	2.12.0	git ソースビルド
mamba-ssm	2.2.5	CUDA ソースビルド
causal-conv1d	1.5.3	CUDA ソースビルド
nv-grouped-gemm	1.1.4.post8	CUDA ソースビルド
flash-linear-attention	0.4.1
+ その他 140 以上の依存パッケージ

ビルドスクリプト

build-megatron-bridge.sh

#!/bin/bash
# build-megatron-bridge.sh  Megatron-Bridge v0.3.0
# 前提: build-megatron-lm.sh 実行済み
# 使い方: bash build-megatron-bridge.sh [MEGATRON_BRIDGE_DIR] [VENV_DIR]
set -euo pipefail

MEGATRON_BRIDGE_DIR="$(cd "${1:-$(pwd)/Megatron-Bridge}" && pwd)"
VENV_DIR="$(cd "${2:-$(pwd)/Megatron-LM/.venv}" && pwd)"
MEGATRON_LM_DIR="$(cd "$(dirname "$VENV_DIR")" && pwd)"

[[ -f "$MEGATRON_BRIDGE_DIR/pyproject.toml" ]] \
    || { echo "ERROR: pyproject.toml not found: $MEGATRON_BRIDGE_DIR" >&2; exit 1; }
[[ -d "$VENV_DIR" ]] \
    || { echo "ERROR: venv not found: $VENV_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

source "$VENV_DIR/bin/activate"
export MAX_JOBS="${MAX_JOBS:-$(nproc)}"

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 をコンパイラに見せる (transformer-engine ビルド用)
NCCL_INCLUDE="$(python -c 'import nvidia.nccl, pathlib; print(pathlib.Path(nvidia.nccl.__path__[0]) / "include")')"
export CPATH="$NCCL_INCLUDE:${CPATH:-}"

# wheel を事前インストール (--no-build-isolation 時に必要)
uv pip install wheel

cd "$MEGATRON_BRIDGE_DIR"
uv pip install -e . \
    --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 \
    --no-build-isolation-package flash-linear-attention \
    || true

# transformer-engine: Bridge インストールで入った cu12/torch をアンインストール
# (git TE 2.12.0 は PyPI 版に置換できないため、不整合なパッケージを除去)
uv pip uninstall transformer-engine-cu12 transformer-engine-torch 2>/dev/null || true

# git TE を再インストールして Python ソースを復元
TE_GIT_REV="5671fd3675906cda1ade26c24a65d3dedd88eb89"
uv pip install "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@${TE_GIT_REV}" \
    --no-build-isolation --reinstall-package transformer-engine

# スタンドアロン版 Megatron-LM を editable install
# (Bridge の 3rdparty 版ではなく、前回ビルドした Megatron-LM を使用)
uv pip install -e "$MEGATRON_LM_DIR[mlm]"

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

Megatron-Bridge の構成を理解する

ビルドの前に依存関係を把握しておくと、つまずいたときに原因を追いやすくなります。

pyproject.toml のポイント

[project]
dependencies = [
    "megatron-core[dev,mlm]>=0.15.0a0,<0.17.0",
    "transformer-engine[pytorch,core_cu13]>=2.10.0a0,<2.13.0",
    "mamba-ssm",
    "causal-conv1d",
    # ... 他多数
]

[tool.uv]
# torch は事前インストール前提 (uv の依存解決から除外)
override-dependencies = [
    "torch; sys_platform == 'never'",
    # git TE を強制使用
    "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@5671fd36...",
]

# CUDA 拡張はビルド時に torch が必要なので分離ビルド不可
no-build-isolation-package = [
    "transformer-engine", "transformer-engine-torch",
    "mamba-ssm", "causal-conv1d", "nv-grouped-gemm",
    "flash_mla", "flash-linear-attention",
]

[tool.uv.sources]
# 3rdparty の megatron-core を editable で使用
megatron-core = { path = "3rdparty/Megatron-LM/", editable = true }
# v0.3.0 では mamba-ssm / causal-conv1d の git 指定なし (PyPI 版を使用)

torch が sys_platform == 'never' で依存解決から除外されているため、事前に手動インストールが必須になります。no-build-isolation-package に列挙されているパッケージは torch が入った状態でビルドする必要があり、--no-build-isolation を付けて個別に処理します。

TransformerEngine については、PyPI に 2.12.0 は存在しますが、pyproject.toml の override-dependencies で特定の git コミットにピン留めされています。これは PyPI 版が分割パッケージ構成（transformer-engine メタ + transformer-engine-cu13 + transformer-engine-torch）であるのに対し、git 版は単一パッケージとしてビルドされるためです。両者が混在すると不整合が起きるため、git 版に統一する必要があります。

Megatron-LM のバージョンについて

前回ビルドした Megatron-LM と Bridge 同梱の 3rdparty/Megatron-LM は、どちらもバージョン 0.16.0、コミット b0cc2706 で同一です。

ただし Bridge をインストールすると megatron-core のソースが 3rdparty 側に切り替わるため、最終ステップで uv pip install -e ./Megatron-LM[mlm] を実行してスタンドアロン版に戻します。
こうすることで、改造した Megatron-LM を Bridge と組み合わせて使うことができます。

# Bridge インストール直後
megatron-core==0.16.0 (from .../Megatron-Bridge/3rdparty/Megatron-LM)

# スタンドアロン版を再インストール後
megatron-core==0.16.0 (from .../Megatron-LM)  ← こちらを使用

つまずいた箇所

`No module named 'wheel'`

causal-conv1d のビルド時にこのエラーが出ることがあります。

ModuleNotFoundError: No module named 'wheel'

--no-build-isolation でビルドするため、通常はビルド環境に自動インストールされる wheel が存在しない状態になります。明示的にインストールすることで解消します。

$ uv pip install wheel ninja

TransformerEngine の PyPI 版と git 版の不整合

PyPI には transformer-engine 2.12.0 が存在しますが、pyproject.toml は override-dependencies で git 版を強制しています。

PyPI 版の TE は分割パッケージ構成（transformer-engine メタ + transformer-engine-cu13 CUDA バイナリ + transformer-engine-torch PyTorch バインディング）ですが、git 版は全てを含む単一パッケージとしてビルドされます。Bridge のインストール時にこれらが混在すると、サブパッケージ間でバージョン不整合が起きます。

Bridge ビルド後、PyPI 由来のサブパッケージをアンインストールしてから git TE を再インストールする必要があります。

# 不要なサブパッケージを削除する
$ uv pip uninstall transformer-engine-cu12 transformer-engine-torch

# git TE を再インストールする (git 版は単一パッケージ)
$ TE_GIT_REV="5671fd3675906cda1ade26c24a65d3dedd88eb89"
$ uv pip install "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@${TE_GIT_REV}" --no-build-isolation --reinstall-package transformer-engine

transformer-engine-cu12 をアンインストールしておくと、TE の sanity_checks_for_pypi_installation() がスキップされるため Root-Is-Purelib の不整合問題を回避できます。

許容される警告

ビルド後のインポート時に以下の警告が出ることがありますが、いずれも正常動作に影響しないと思います。

Triton is not supported on current platform, roll back to CPU.

aarch64 で Triton の一部機能が制限されますが、flash-linear-attention の CPU フォールバックが使われるだけです。

The pynvml package is deprecated. Please install nvidia-ml-py instead.

警告のみで、nvidia-ml-py は自動的にインストールされています。

最終的なパッケージ一覧

$ uv pip list | grep -iE "megatron|torch|transformer-engine|mamba|causal|flash|triton"
causal-conv1d              1.5.3
flash-linear-attention     0.4.1
mamba-ssm                  2.2.5
megatron-bridge            0.3.0     (editable)
megatron-core              0.16.0    (editable)
megatron-energon           6.0.1
torch                      2.10.0+cu129
torchvision                0.25.0+cu129
transformer-engine         2.12.0+5671fd36
triton                     3.6.0

使い方

$ module purge
$ module load cuda/12.9

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

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

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 torch
print(f'CUDA available: {torch.cuda.is_available()}')
print(f'GPU: {torch.cuda.get_device_name(0)}')
import megatron.bridge
print(f'Bridge: {megatron.bridge.__version__}')
"

全体の依存関係は以下のようになります。

Megatron-LM ビルド（前回の記事）
  └─ torch 2.10.0+cu129
  └─ megatron-core 0.16.0
       │
       ▼
Megatron-Bridge ビルド（本記事）
  └─ megatron-bridge 0.3.0
  └─ transformer-engine 2.12.0 (git build)
  └─ mamba-ssm 2.2.5
  └─ causal-conv1d 1.5.3
  └─ その他

おわりに

Megatron-Bridge の構成を見ると、Docker や Singularity でのビルドを前提に設計されている印象を受けます。torch を依存解決から除外して事前インストールを必須にしているのも、コンテナ内に PyTorch 入りのベースイメージがある前提で組まれているためだと思います。
今回初めて --no-build-isolation を使いましたが、ビルド環境と実行環境を分離するという考え方は便利な反面、依存関係の管理が少し複雑になると感じました。特に CUDA 拡張を含むパッケージは torch の存在が前提になるため、ビルド順序やインストール手順に注意が必要です。同じ環境で試行錯誤している方の参考になれば幸いです。