Nemo Mbridge Resiliency

nvidia/skills/skills/nemo-mbridge-resiliency

作者 nvidiacf5224d14250Apache-2.03.5K 個星標收錄於 2026年10月8日更新於 2026年10月8日儲存庫今天更新

Resiliency features in Megatron Bridge including fault tolerance, straggler detection, in-process restart, preemption, and re-run state machine.

僅含說明DevOps & Cloud
AI 產生的概覽

說明 Megatron Bridge 的容錯、落後者偵測、搶占、重啟與檢查點等韌性功能。

功能
此技能是一份參考指南,用於啟用與排查 Megatron Bridge 訓練中的韌性功能。內容涵蓋容錯設定與啟動器用法、NVRx 落後者偵測、搶占訊號處理、實驗性重跑狀態機、行程內重啟,以及非同步或本機檢查點。它提供設定片段、參數表、程式碼位置、常見陷阱與驗證指令。
適用情境
適用於在 Megatron Bridge 中啟用韌性功能,或排查訓練停滯、落後者偵測失敗、重啟中斷與搶占問題。也適合在確認某項韌性功能需要哪些設定參數與環境變數時使用。
執行需求
不隨附指令碼,僅為說明文件。依指南操作需要 Megatron Bridge 訓練環境,包含 PyTorch、NCCL、NeMo-Run 與 nvidia-resiliency-ext 套件;容錯與搶占外掛還需要 Slurm;行程內重啟需要特定 PyTorch/NCCL 版本與環境變數。

Resiliency

Stable docs: @docs/training/resiliency.md, @docs/training/checkpointing.md Card: @skills/nemo-mbridge-resiliency/card.yaml

Enablement

Fault tolerance (Slurm only)

Option 1: NeMo Run plugin (recommended)
python
from megatron.bridge.recipes.run_plugins import FaultTolerancePluginimport 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)
Plugin parameterDefaultDescription
num_in_job_restarts3Max restarts within same job
num_job_retries_on_failure2Max new job launches on failure
initial_rank_heartbeat_timeout1800First heartbeat timeout (seconds)
rank_heartbeat_timeout300Subsequent heartbeat timeout (seconds)
Option 2: Direct config + ft_launcher
python
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",)

Launch with ft_launcher (not torchrun):

bash
export GROUP_RANK=0  # required for non-Slurmft_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
Config parameterDefaultDescription
enable_ft_packageFalseEnable fault tolerance
calc_ft_timeoutsFalseAuto-compute optimal timeouts
simulate_faultFalseEnable fault simulation for testing
simulated_fault_type"random""rank_hung", "rank_killed", or "random"
simulated_fault_rankNoneSpecific rank to fault (random if None)
simulated_fault_base_delay0Base delay before simulating fault

Section-based timeout monitoring covers setup, training steps, checkpointing, and out-of-section time independently. Timeouts are saved to ft_state.json for subsequent runs when calc_ft_timeouts=True.

NVRx straggler detection

python
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,)
ParameterDefaultDescription
enabledFalseEnable straggler detection
report_time_interval300.0Seconds between straggler checks
calc_relative_gpu_perfTrueCompare ranks against each other
calc_individual_gpu_perfTrueTrack per-rank degradation over time
gpu_relative_perf_threshold0.7Threshold for relative performance (0-1)
gpu_individual_perf_threshold0.7Threshold for individual performance (0-1)
stop_if_detectedFalseTerminate training on straggler
num_gpu_perf_scores_to_print5Number of best/worst scores to print
profiling_interval1Profiling interval for detector

Preemption

Plugin (Slurm)
python
from megatron.bridge.recipes.run_plugins import PreemptionPlugin
plugins = [    PreemptionPlugin(        preempt_time=60,        enable_exit_handler=True,        enable_exit_handler_for_data_loader=False,    )]
Plugin parameterDefaultDescription
preempt_time60Seconds before job limit to send signal
enable_exit_handlerTrueEnable signal handler in training
enable_exit_handler_for_data_loaderFalseEnable for dataloader workers
Direct config
python
import signalcfg.train.exit_signal_handler = Truecfg.train.exit_signal = signal.SIGTERMcfg.train.exit_signal_handler_for_dataloader = False

Re-run state machine (experimental)

python
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,)
ParameterDefaultDescription
rerun_mode"disabled""disabled", "validate_results", "report_determinism_stats"
check_for_nan_in_lossTrueCheck for NaN in loss
check_for_spiky_lossFalseCheck for unexpectedly large loss
spiky_loss_factor10.0Loss flagged if > factor * max observed (increase for large models)

Exit codes: 16 = resume to disambiguate, 17 = failed validation.

In-process restart (experimental)

python
from megatron.bridge.training.config import InProcessRestartConfig
cfg.inprocess_restart = InProcessRestartConfig(    enabled=True,    granularity="node",    soft_timeout=60.0,    hard_timeout=90.0,)
ParameterDefaultDescription
enabledFalseEnable in-process restart
active_world_sizeNoneRanks executing workload (rest are warm reserves)
granularity"node""node" or "rank" restart granularity
max_iterationsNoneMax restart attempts (None = unlimited)
soft_timeout60.0Detect GIL-released hangs (seconds)
hard_timeout90.0Force-terminate hung ranks (seconds)
heartbeat_interval30.0Heartbeat interval (seconds)
heartbeat_timeout60.0Missing heartbeat timeout (seconds)
barrier_timeout120.0Distributed barrier timeout (seconds)
completion_timeout120.0Completion barrier timeout (seconds)
empty_cuda_cacheTrueClear CUDA cache during restart
max_rank_faultsNoneMax rank faults before terminating
monitor_process_logdirNoneDirectory for monitor logs

Required environment variables:

bash
export TORCH_CPP_LOG_LEVEL=errorexport TORCH_NCCL_RETHROW_CUDA_ERRORS=0export NCCL_NVLS_ENABLE=0

The PyTorch NCCL watchdog timeout must exceed hard_timeout. NeMo-Run's Slurm Executor is not supported; launch directly with srun --kill-on-bad-exit=0.

Async checkpoint save

python
cfg.checkpoint.async_save = Truecfg.checkpoint.ckpt_format = "torch_dist"

Local checkpointing (NVRx)

python
cfg.checkpoint.non_persistent_local_ckpt_dir = "/local/scratch/ckpt"cfg.checkpoint.non_persistent_local_ckpt_algo = "fully_parallel"

Code Anchors

Fault tolerance

  • Config: src/megatron/bridge/training/config.py — FaultToleranceConfig
  • Runtime: src/megatron/bridge/training/fault_tolerance.py
  • Plugin: src/megatron/bridge/recipes/run_plugins.py — FaultTolerancePlugin
  • Perf plugin: scripts/performance/nemo-mbridge-resiliency_plugins.py
  • Tests: tests/unit_tests/training/test_fault_tolerance.py
  • Example: examples/training_features/nemo-mbridge-resiliency/fault_tolerance/

Straggler detection

  • Config: src/megatron/bridge/training/config.py — NVRxStragglerDetectionConfig
  • Runtime: src/megatron/bridge/training/nvrx_straggler.py
  • Train loop: src/megatron/bridge/training/train.py — check_nvrx_straggler_detection
  • Tests: tests/unit_tests/training/test_nvrx_straggler.py, tests/functional_tests/training/test_nvrx_straggler.py
  • Example: examples/training_features/nemo-mbridge-resiliency/straggler_detection/

In-process restart

  • Config: src/megatron/bridge/training/config.py — InProcessRestartConfig
  • Runtime: src/megatron/bridge/training/inprocess_restart.py
  • Entry point: src/megatron/bridge/training/pretrain.py — maybe_wrap_for_inprocess_restart
  • Tests: tests/unit_tests/training/test_inprocess_restart.py, tests/functional_tests/training/test_inprocess_restart.py

Preemption

  • Plugin: src/megatron/bridge/recipes/run_plugins.py — PreemptionPlugin
  • Signal handler: src/megatron/bridge/training/utils/sig_utils.py
  • Tests: tests/unit_tests/recipes/test_run_plugins.py

Re-run state machine

  • Config: src/megatron/bridge/training/config.py — RerunStateMachineConfig
  • Init: src/megatron/bridge/training/initialize.py — init_rerun_state

Checkpointing

  • Async save: src/megatron/bridge/training/checkpointing.py — schedule_async_save
  • Local ckpt: src/megatron/bridge/training/checkpointing.py — LocalCheckpointManager
  • Tests: tests/functional_tests/training/test_local_checkpointing.py

Pitfalls

  1. ft_launcher, not torchrun: Direct FaultToleranceConfig requires ft_launcher. Using torchrun silently disables FT. For non-Slurm, set GROUP_RANK=0.

  2. Async save requires torch_dist: async_save=True only works with ckpt_format="torch_dist". Other formats silently fail or error.

  3. IPR + NeMo-Run: In-process restart is not compatible with NeMo-Run or Slurm preemption plugins. Requires specific PyTorch/NCCL versions and env vars.

  4. NVRx vs legacy straggler: Two detectors exist. Use NVRx (nvrx_straggler); do not enable both.

  5. stop_if_detected default: NVRx logs but does not stop training by default. Set stop_if_detected=True for automatic termination.

  6. NCCL watchdog vs hard_timeout: For IPR, NCCL watchdog timeout must exceed hard_timeout or PyTorch kills the process before recovery.

  7. Rerun state machine is alpha: Use check_for_nan_in_loss=True for NaN detection, but don't rely on full rerun workflows yet.

Verification

Fault tolerance

bash
./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

Look for [FaultTolerance] / [RankMonitorServer] log lines with section timeouts. Simulated fault should trigger restart from checkpoint.

Straggler detection

bash
uv run python -m torch.distributed.run --nproc_per_node=2 \    examples/training_features/nemo-mbridge-resiliency/straggler_detection/straggler_detection_example.py

Look for GPU relative performance and GPU individual performance reports with per-rank scores.

Async checkpoint

Look for Scheduling async checkpoint save in logs. Training iterations should continue while checkpoint files are being written.

In-process restart

bash
pytest tests/functional_tests/training/test_inprocess_restart.py -v

Requires compatible PyTorch/NCCL versions.

來源與署名

來源:nvidia/skills位於skills/nemo-mbridge-resiliency提交cf5224d

授權條款: Apache-2.0

內容歸原作者所有。SourceWeft 從公開儲存庫中收錄這些內容。

檢舉或申請下架