mcpField
N
nemo-mbridge-resiliency

nemo-mbridge-resiliency

github.com/nvidiaクラウド・DevOps3.8k1.8k 回インストール今日更新
リポジトリ

nemo-mbridge-resiliency について

Megatron Bridge のトレーニングに回復性を設定する開発者向けスキルです。分散処理の障害に備え、フォールトトレランス、プリエンプション、遅延検出、再起動、チェックポイント保存を扱います。

主な機能

  1. 01フォールトトレランスとジョブ再試行
  2. 02NVRx による GPU 遅延検出
  3. 03プリエンプション信号処理
  4. 04プロセス内再起動と結果検証
  5. 05非同期・ローカルチェックポイント

活用シーン

  1. 01ジョブ障害後のチェックポイントからの再開
  2. 02性能の低い GPU の検出と報告
  3. 03ジョブ期限前のプリエンプション処理

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/nvidia/skills --skill nemo-mbridge-resiliency

インストール後、Claude Code で/nemo-mbridge-resiliencyを入力して読み込まれたか確認します。

同梱ファイル

SKILL.md と一緒にインストールされる参考資料やスクリプト。

BENCHMARK.md
card.yaml
skill-card.md
skill.oms.sig
evals/
evals.json

詳細

レジリエンス

安定版ドキュメント：@docs/training/resiliency.md、@docs/training/checkpointing.md カード：@skills/nemo-mbridge-resiliency/card.yaml

有効化

フォールトトレランス（Slurm のみ）
オプション 1：NeMo Run プラグイン（推奨）
from megatron.bridge.recipes.run_plugins import FaultTolerancePlugin
import nemo_run as run

task = run.Script(...)
run_plugins = [
    FaultTolerancePlugin(
        enable_ft_package=True,
        calc_ft_timeouts=True,
        num_in_job_restarts=3,
        num_job_retries_on_failure=2,
        initial_rank_heartbeat_timeout=1800,
        rank_heartbeat_timeout=300,
    )
]
run.run(task, plugins=run_plugins, executor=executor)
プラグインパラメーターデフォルト説明
num_in_job_restarts3同じジョブ内での最大再起動回数
num_job_retries_on_failure2失敗時に新たに起動するジョブの最大数
initial_rank_heartbeat_timeout1800最初のハートビートのタイムアウト（秒）
rank_heartbeat_timeout300以降のハートビートのタイムアウト（秒）
オプション 2：直接設定 + ft_launcher
from megatron.bridge.training.config import FaultToleranceConfig

cfg.ft = FaultToleranceConfig(
    enable_ft_package=True,
    calc_ft_timeouts=True,
    simulate_fault=False,
    simulated_fault_type="random",
)

ft_launcher（torchrun ではありません）で起動します：

export GROUP_RANK=0  # required for non-Slurm
ft_launcher \
    --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
    --nnodes=${NUM_NODES} --nproc-per-node=${NUM_GPUS_PER_NODE} \
    --ft-rank_section_timeouts=setup:600,step:180,checkpointing:420 \
    --ft-rank_out_of_section_timeout=300 \
    your_training_script.py
設定パラメーターデフォルト説明
enable_ft_packageFalseフォールトトレランスを有効にします
calc_ft_timeoutsFalse最適なタイムアウトを自動計算します
simulate_faultFalseテスト用の障害シミュレーションを有効にします
simulated_fault_type"random""rank_hung"、"rank_killed"、または "random"
simulated_fault_rankNone障害を発生させる特定のランク（None の場合はランダム）
simulated_fault_base_delay0障害をシミュレートするまでの基本遅延

セクション単位のタイムアウト監視により、セットアップ、トレーニングステップ、チェックポイント保存、およびセクション外の時間を個別に監視します。calc_ft_timeouts=True の場合、タイムアウトは以降の実行用に ft_state.json に保存されます。

NVRx のストラグラー検出
from megatron.bridge.training.config import NVRxStragglerDetectionConfig

cfg.nvrx_straggler = NVRxStragglerDetectionConfig(
    enabled=True,
    report_time_interval=300.0,
    calc_relative_gpu_perf=True,
    calc_individual_gpu_perf=True,
    num_gpu_perf_scores_to_print=5,
    gpu_relative_perf_threshold=0.7,
    gpu_individual_perf_threshold=0.7,
    stop_if_detected=False,
    enable_logging=True,
)
パラメーターデフォルト説明
enabledFalseストラグラー検出を有効にします
report_time_interval300.0ストラグラーをチェックする間隔（秒）
calc_relative_gpu_perfTrue各ランクを相互に比較します
calc_individual_gpu_perfTrueランクごとの時間経過に伴う性能低下を追跡します
gpu_relative_perf_threshold0.7相対性能のしきい値（0～1）
gpu_individual_perf_threshold0.7個別性能のしきい値（0～1）
stop_if_detectedFalseストラグラー検出時にトレーニングを終了します
num_gpu_perf_scores_to_print5表示する最良／最悪スコアの数
profiling_interval1検出器のプロファイリング間隔
プリエンプション
プラグイン（Slurm）
from megatron.bridge.recipes.run_plugins import PreemptionPlugin

plugins = [
    PreemptionPlugin(
        preempt_time=60,
        enable_exit_handler=True,
        enable_exit_handler_for_data_loader=False,
    )
]
プラグインパラメーターデフォルト説明
preempt_time60ジョブの制限時間前にシグナルを送信するまでの秒数
enable_exit_handlerTrueトレーニングでシグナルハンドラーを有効にします
enable_exit_handler_for_data_loaderFalseデータローダーワーカーで有効にします
直接設定
import signal
cfg.train.exit_signal_handler = True
cfg.train.exit_signal = signal.SIGTERM
cfg.train.exit_signal_handler_for_dataloader = False
再実行ステートマシン（実験的）
from megatron.bridge.training.config import RerunStateMachineConfig

cfg.rerun_state_machine = RerunStateMachineConfig(
    rerun_mode="validate_results",
    check_for_nan_in_loss=True,
    check_for_spiky_loss=False,
    spiky_loss_factor=10.0,
)
パラメーターデフォルト説明
rerun_mode"disabled""disabled"、"validate_results"、"report_determinism_stats"
check_for_nan_in_lossTrue損失に NaN が含まれていないか確認します
check_for_spiky_lossFalse予想外に大きな損失がないか確認します
spiky_loss_factor10.0損失が観測済みの最大値 × 係数を超えた場合にフラグを付けます（大規模モデルでは値を大きくしてください）

終了コード：16 = 判別のために再開、17 = 検証失敗。

プロセス内再起動（実験的）
from megatron.bridge.training.config import InProcessRestartConfig

cfg.inprocess_restart = InProcessRestartConfig(
    enabled=True,
    granularity="node",
    soft_timeout=60.0,
    hard_timeout=90.0,
)
パラメーターデフォルト説明
enabledFalseプロセス内再起動を有効にします
active_world_sizeNoneワークロードを実行するランク（残りは待機予備になります）
granularity"node"再起動の粒度（"node" または "rank"）
max_iterationsNone再起動の最大試行回数（None = 無制限）
soft_timeout60.0GIL が解放された状態でのハングを検出します（秒）
hard_timeout90.0ハングしたランクを強制終了します（秒）
heartbeat_interval30.0ハートビートの間隔（秒）
heartbeat_timeout60.0ハートビートが届かない場合のタイムアウト（秒）
barrier_timeout120.0分散バリアのタイムアウト（秒）
completion_timeout120.0完了バリアのタイムアウト（秒）
empty_cuda_cacheTrue再起動中に CUDA キャッシュをクリアします
max_rank_faultsNone終了するまでのランク障害の最大回数
monitor_process_logdirNoneモニターのログ用ディレクトリ

必須の環境変数：

export TORCH_CPP_LOG_LEVEL=error
export TORCH_NCCL_RETHROW_CUDA_ERRORS=0
export NCCL_NVLS_ENABLE=0

PyTorch NCCL ウォッチドッグのタイムアウトは、hard_timeout より大きくする必要があります。NeMo-Run の Slurm Executor はサポートされていません。srun --kill-on-bad-exit=0 で直接起動してください。

非同期チェックポイント保存
cfg.checkpoint.async_save = True
cfg.checkpoint.ckpt_format = "torch_dist"
ローカルチェックポイント（NVRx）
cfg.checkpoint.non_persistent_local_ckpt_dir = "/local/scratch/ckpt"
cfg.checkpoint.non_persistent_local_ckpt_algo = "fully_parallel"

コードアンカー

耐障害性
  • 設定：src/megatron/bridge/training/config.py — FaultToleranceConfig
  • ランタイム：src/megatron/bridge/training/fault_tolerance.py
  • プラグイン：src/megatron/bridge/recipes/run_plugins.py — FaultTolerancePlugin
  • パフォーマンスプラグイン：scripts/performance/nemo-mbridge-resiliency_plugins.py
  • テスト：tests/unit_tests/training/test_fault_tolerance.py
  • 例：examples/training_features/nemo-mbridge-resiliency/fault_tolerance/
ストラグラー検出
  • 設定：src/megatron/bridge/training/config.py — NVRxStragglerDetectionConfig
  • ランタイム：src/megatron/bridge/training/nvrx_straggler.py
  • トレーニングループ：src/megatron/bridge/training/train.py — check_nvrx_straggler_detection
  • テスト：tests/unit_tests/training/test_nvrx_straggler.py、tests/functional_tests/training/test_nvrx_straggler.py
  • 例：examples/training_features/nemo-mbridge-resiliency/straggler_detection/
プロセス内再起動
  • 設定：src/megatron/bridge/training/config.py — InProcessRestartConfig
  • ランタイム：src/megatron/bridge/training/inprocess_restart.py
  • エントリーポイント：src/megatron/bridge/training/pretrain.py — maybe_wrap_for_inprocess_restart
  • テスト：tests/unit_tests/training/test_inprocess_restart.py、tests/functional_tests/training/test_inprocess_restart.py
プリエンプション
  • プラグイン：src/megatron/bridge/recipes/run_plugins.py — PreemptionPlugin
  • シグナルハンドラー：src/megatron/bridge/training/utils/sig_utils.py
  • テスト：tests/unit_tests/recipes/test_run_plugins.py
再実行ステートマシン
  • 設定：src/megatron/bridge/training/config.py — RerunStateMachineConfig
  • 初期化：src/megatron/bridge/training/initialize.py — init_rerun_state
チェックポイント
  • 非同期保存：src/megatron/bridge/training/checkpointing.py — schedule_async_save
  • ローカルチェックポイント：src/megatron/bridge/training/checkpointing.py — LocalCheckpointManager
  • テスト：tests/functional_tests/training/test_local_checkpointing.py

落とし穴

  1. torchrun ではなく ft_launcher：FaultToleranceConfig を直接使用するには、 ft_launcher が必要です。torchrun を使用すると、FT が暗黙的に無効になります。Slurm を使用しない場合は、 GROUP_RANK=0 を設定してください。

  2. 非同期保存には torch_dist が必要です：async_save=True が機能するのは、 ckpt_format="torch_dist" の場合のみです。他の形式では、暗黙的に失敗するかエラーになります。

  3. IPR と NeMo-Run：インプロセス再起動は NeMo-Run や Slurm のプリエンプションプラグインと互換性がありません。特定の PyTorch／NCCL バージョン と環境変数が必要です。

  4. NVRx とレガシーのストラグラー検出：検出器が 2 種類あります。NVRx (nvrx_straggler) を使用してください。両方を有効にしないでください。

  5. stop_if_detected のデフォルト値：NVRx はデフォルトではログを記録しますが、トレーニングを停止しません。 自動終了するには stop_if_detected=True を設定してください。

  6. NCCL watchdog と hard_timeout：IPR では、NCCL watchdog のタイムアウトを hard_timeout より長くする必要があります。そうしないと、復旧前に PyTorch がプロセスを終了します。

  7. Rerun のステートマシンはアルファ版です：NaN の検出には check_for_nan_in_loss=True を使用してください。ただし、 まだ rerun ワークフロー全体には依存しないでください。

検証

フォールトトレランス
./examples/training_features/nemo-mbridge-resiliency/fault_tolerance/run_fault_tolerance.sh
./examples/training_features/nemo-mbridge-resiliency/fault_tolerance/run_fault_tolerance.sh --simulate-fault

セクションのタイムアウトを含む [FaultTolerance]／[RankMonitorServer] のログ行を確認してください。 シミュレートした障害が発生すると、チェックポイントから再起動されるはずです。

ストラグラー検出
uv run python -m torch.distributed.run --nproc_per_node=2 \
    examples/training_features/nemo-mbridge-resiliency/straggler_detection/straggler_detection_example.py

ランクごとのスコアを含む GPU relative performance および GPU individual performance のレポートを確認してください。

非同期チェックポイント

ログで Scheduling async checkpoint save を確認してください。チェックポイントファイルの書き込み中も、 トレーニングのイテレーションが続行されるはずです。

インプロセス再起動
pytest tests/functional_tests/training/test_inprocess_restart.py -v

互換性のある PyTorch／NCCL バージョンが必要です。

コメント

Markdown 対応