メインコンテンツまでスキップ
バージョン: v2603

Megatronプリセット

Megatron-LM / ms-swift によるマルチGPU分散学習のパラメータを自動チューニングするプリセットです。 並列化構成・マイクロバッチサイズ・活性化再計算の設定を探索し、学習スループット（TFLOP/s/GPU）を最大化します。

CLIの基本操作（optimize / apply / analyze / タイムアウト等）はノーコードチューニングの使用方法を参照してください。

クイックスタート​

ZenithTuneはフレームワークのトレーニングコマンドを直接ラップします。bash train.shのようなシェルスクリプトではなく、実際にCLIオプションを受け付けるコマンドを指定してください。

Megatron-LMの例:

zenithtune optimize --preset megatron \
    --args n_gpus=8,gbs=64 \
    --n-trials 50 --timeout-dynamic \
    -- python pretrain_gpt.py --num-layers 30 --hidden-size 4096 --train-iters 3 ...

ms-swiftの例:

zenithtune optimize --preset megatron \
    --args n_gpus=8,gbs=64 \
    --n-trials 50 --timeout-dynamic \
    -- swift sft --model Qwen/Qwen3-30B-A3B --max_steps 3 ...
警告

ZenithTuneが--tensor-model-parallel-size等のオプションを自動で注入・上書きするため、トレーニングコマンド側での既存オプションの変更／削除は不要です。

ただし、チューニング時間を削減するために、トレーニングコマンド側でイテレーション／エポック／データセットを減らすなどの工夫は効果的です。 しかし、意図的にエポック数などを減らしたトレーニングでチューニングした場合に本番学習との挙動が変わる場合もあり注意が必要です。

トレーニングコマンドの要件​

ZenithTuneに渡すコマンドは以下を満たす必要があります。

  1. Megatron形式のCLIオプションを受け付ける — ZenithTuneは --tensor-model-parallel-size 等のオプションをコマンドに注入します。既存のオプションは上書きされ、無いものは末尾に追加されます
  2. スループットを標準出力に出力する — throughput per GPU (TFLOP/s/GPU): <値> の形式。複数行ある場合は最後の値が採用されます。Megatronはこの形式で出力するため追加対応は不要です
警告

bash train.shのようなラッパースクリプトを指定した場合、ZenithTuneが注入するCLIオプションはシェルスクリプトの引数となり、フレームワーク本体には渡されません。python pretrain_gpt.pyやswift sftなど、CLIオプションを直接受け付けるコマンドを指定してください。

args​

--argsでカンマ区切りのkey=value形式で渡します。

引数説明デフォルト
n_gpusGPU総数（必須）。MegatronHeuristicStrategy使用時は2のべき乗-
gbsグローバルバッチサイズ（必須）-
modelHuggingFaceモデルID。レイヤー数・KVヘッド数・VLMフラグを自動取得しプルーニングに利用なし
num_layersレイヤー数。modelの自動検出を上書きなし
num_query_groupsKVヘッド数。modelの自動検出を上書きなし
mbs_maxマイクロバッチサイズの探索上限8
stage1_trialsmegatron-staged-blackbox戦略のStage 1試行回数30
ヒント

model、num_layers、num_query_groupsはいずれも任意です。 これらを指定することでGQA制約・レイヤー数制約・VLMモデル制約によるプルーニングが有効になり、探索効率が向上します。 modelを指定するとHuggingFaceからレイヤー数・KVヘッド数・VLMフラグを自動取得します。 modelの代わりにnum_layers・num_query_groupsを指定することでGQA制約・レイヤー数制約が有効になります。ただしVLMモデル制約はmodel指定時のみ有効です。

探索パラメータ​

並列化パラメータ​

パラメータCLIオプション選択肢説明
TP--tensor-model-parallel-size1, 2, 4, ... n_gpusテンソル並列
PP--pipeline-model-parallel-size同上パイプライン並列
CP--context-parallel-size同上コンテキスト並列
EP--expert-model-parallel-size同上エキスパート並列（MoE）
ETP--expert-tensor-parallel-size同上エキスパートテンソル並列（MoE、TP以下）

計算効率パラメータ​

パラメータCLIオプション選択肢説明
MBS--micro-batch-size1, 2, 4, ... mbs_max（GBSの約数のみ）マイクロバッチサイズ
GBS--global-batch-size固定値（gbsで指定）グローバルバッチサイズ
活性化再計算の粒度--recompute-granularityfull, selective活性化再計算の方式
活性化再計算の方法--recompute-methoduniform, blockfull粒度時のみ有効
活性化再計算のレイヤー数--recompute-num-layers1〜8full粒度時のみ有効
ViT勾配チェックポイント--vit-gradient-checkpointingtrue, falseVLMモデル用
備考

MegatronHeuristicStrategy（デフォルト）では PP=1、CP=1、ETP=1、recompute_method=uniform、recompute_num_layers=1、vit_gradient_checkpointing=true に固定されます。--strategy megatron-staged-blackbox を指定した場合はこれらも探索されます。

探索戦略​

戦略CLIオプション説明
MegatronHeuristicStrategyデフォルトドメイン知識に基づく決定的探索
MegatronStagedBlackboxStrategy--strategy megatron-staged-blackbox段階的なベイズ最適化による探索
# staged-blackbox戦略の例
zenithtune optimize --preset megatron \
    --args n_gpus=8,gbs=64,stage1_trials=20 \
    --strategy megatron-staged-blackbox --n-trials 50 --timeout-dynamic \
    -- python pretrain_gpt.py --num-layers 30 --hidden-size 4096 --train-iters 3

評価器​

デフォルトの評価器はMegatronThroughputEvaluatorで、標準出力から throughput per GPU (TFLOP/s/GPU): <値> を抽出し最大化します。複数行ある場合は最後の値が採用されます。

評価器の変更方法についてはEvaluatorとはを参照してください。

実行例​

# テスト実行（ダミースクリプト）
zenithtune optimize --preset megatron \
    --args n_gpus=8,gbs=64 \
    --n-trials 50 \
    -- python <aibooster-examples>/intelligence/zenith_tune/nocode/fake_megatron_train.py

# Megatron-LM（HuggingFaceモデル設定を利用、GQA・VLM制約が自動有効化）
zenithtune optimize --preset megatron \
    --args model=Qwen/Qwen3-Omni-30B-A3B-Instruct,n_gpus=8,gbs=64 \
    --n-trials 50 --timeout-dynamic \
    -- python pretrain_gpt.py --num-layers 30 --hidden-size 4096 --train-iters 3

# ms-swift
zenithtune optimize --preset megatron \
    --args n_gpus=8,gbs=64 \
    --n-trials 50 --timeout-dynamic \
    -- swift sft --model Qwen/Qwen3-30B-A3B --max_steps 3

# 結果の分析
zenithtune analyze outputs/study_YYYYMMDD_HHMMSS/study.db

# 最適パラメータで実行（--argsも必要）
zenithtune apply \
    --db-path outputs/study_YYYYMMDD_HHMMSS/study.db \
    --preset megatron --args n_gpus=8,gbs=64 \
    -- python pretrain_gpt.py --num-layers 30 --hidden-size 4096 --train-iters 100

チューニング時のスクリプト設定​

チューニングでは1試行あたりの実行時間を短くすることが重要です。トレーニングスクリプトの設定を以下のように調整してください。

  • イテレーション数を小さくする — スループット測定に必要な最小限の値に設定してください。ただし、イテレーション数が小さすぎるとウォームアップの影響を受けるため、適切な値を指定してください
  • バリデーションをオフにする — バリデーションはスループット測定に不要で、試行時間を大幅に増加させます
  • チェックポイント保存をオフにする — チューニング中は不要です

これらの設定はZenithTuneではなく、トレーニングコマンドのオプションや設定ファイルで行ってください。

トラブルシューティング​

問題対処
すべての試行が失敗するスクリプトを直接実行し throughput per GPU (TFLOP/s/GPU): <値> が出力されるか確認
「n_gpus must be a power of 2」--strategy megatron-staged-blackbox を使用
探索が途中で打ち切られる--n-trials を増やす
試行の進みが遅いバリデーションがオンになっていないか、イテレーション数が大きすぎないか確認