0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@aida0710(Masaki Aida)

Megatron-Bridge v0.3.0をGH200の実環境でビルドする

0
Last updated at Posted at 2026-03-03

概要

前回の記事で 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 の存在が前提になるため、ビルド順序やインストール手順に注意が必要です。同じ環境で試行錯誤している方の参考になれば幸いです。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?