Generated banner for Nemo Mbridge Resiliency
Generated logo for Nemo Mbridge Resiliency

Nemo Mbridge Resiliency

NVIDIAOfficialPublished by NVIDIA

  • Skill
  • Tool integration
  • Software development
  • Apache-2.0
  • Medium risk

From the author

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

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides configuration and implementation guidance for resiliency features within the Megatron Bridge framework. It covers mechanisms to maintain training stability in large-scale distributed environments, including automated fault tolerance, straggler detection, preemption handling, and experimental in-process restart capabilities. The documentation outlines how to integrate these features via NeMo Run plugins or direct configuration, along with specific requirements for environment variables and launcher utilities.

What it does

  • Enables fault tolerance with automated heartbeat monitoring and job restarts.
  • Detects performance degradation across GPUs using NVRx straggler detection.
  • Configures preemption handling for Slurm-based environments.
  • Supports experimental in-process restart and re-run state machine logic.
  • Optimizes checkpointing through asynchronous saving and local storage options.

When to use it

  • When setting up training jobs that require high fault tolerance in distributed clusters.
  • When investigating training hangs or performance bottlenecks caused by GPU stragglers.
  • When configuring training to handle preemption signals gracefully.
  • When implementing advanced recovery strategies like in-process restarts for long-running jobs.

Advantages

  • Reduces manual intervention for common training failures.
  • Provides granular control over performance monitoring and timeout thresholds.
  • Supports both plugin-based and direct configuration workflows.
  • Includes built-in fault simulation for testing resiliency configurations.

Usage examples

“How do I enable fault tolerance in my NeMo Run training script?”
You can import the FaultTolerancePlugin from megatron.bridge.recipes.run_plugins and add it to your run.run() call, specifying parameters like num_in_job_restarts and heartbeat timeouts.
“How can I detect slow GPUs during training?”
Configure the NVRxStragglerDetectionConfig in your training setup, setting enabled=True and adjusting the performance thresholds to monitor and report GPU degradation.

Install

  1. Download the skill as a ZIP:

    https://autogineer.com/agent-skills/nvidia/nemo-mbridge-resiliency/download
  2. Unzip it into the cross-client skills folder of your project:

    mkdir -p .agents/skills && unzip nemo-mbridge-resiliency.zip -d .agents/skills

Documentation by NVIDIA

Resiliency

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

Enablement

Fault tolerance (Slurm only)
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)
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
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):

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
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
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)
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
import signal
cfg.train.exit_signal_handler = True
cfg.train.exit_signal = signal.SIGTERM
cfg.train.exit_signal_handler_for_dataloader = False
Re-run state machine (experimental)
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)
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:

export TORCH_CPP_LOG_LEVEL=error
export TORCH_NCCL_RETHROW_CUDA_ERRORS=0
export 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
cfg.checkpoint.async_save = True
cfg.checkpoint.ckpt_format = "torch_dist"
Local checkpointing (NVRx)
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
./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
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
pytest tests/functional_tests/training/test_inprocess_restart.py -v

Requires compatible PyTorch/NCCL versions.

Files

card.yaml

title: resiliency
validated_on: "2026-03-16"
validation_method: file_existence_only
summary: >
  Megatron Bridge integrates nvidia-resiliency-ext for fault tolerance (hang
  detection + auto restart), NVRx straggler detection, in-process restart
  (experimental), preemption (graceful shutdown), and re-run state machine
  (experimental NaN attribution). Async checkpoint save and local checkpointing
  support faster recovery.
validation_status:
  fault_tolerance_config:
    - file_exists
  fault_tolerance_plugin:
    - file_exists
  fault_tolerance_unit_tests:
    - file_exists
  fault_tolerance_example:
    - file_exists
  nvrx_straggler_config:
    - file_exists
  nvrx_straggler_unit_tests:
    - file_exists
  nvrx_straggler_functional_test:
    - file_exists
  nvrx_straggler_example:
    - file_exists
  inprocess_restart_config:
    - file_exists
  inprocess_restart_unit_tests:
    - file_exists
  inprocess_restart_functional_test:
    - file_exists
  preemption_plugin:
    - file_exists
  preemption_unit_tests:
    - file_exists
  rerun_state_machine_config:
    - file_exists
  rerun_state_machine_runtime:
    - unclear
  async_checkpoint_save:
    - file_exists
  local_checkpointing:
    - file_exists
  local_checkpointing_functional_test:
    - file_exists
feature_meaning:
  fault_tolerance: >
    Hang detection and automatic job restart via ft_launcher and
    nvidia-resiliency-ext RankMonitorClient. Slurm-only.
  nvrx_straggler_detection: >
    GPU performance monitoring that identifies slow ranks and optionally
    terminates training. Uses nvidia-resiliency-ext attribution module.
  inprocess_restart: >
    Experimental restart within the same process using
    nvidia-resiliency-ext inprocess module. Does not require a new job.
  preemption: >
    Graceful shutdown on SIGTERM with checkpoint save before exit.
    Slurm preemption support via PreemptionPlugin.
  rerun_state_machine: >
    Experimental NaN and spiky loss detection with automatic rerun
    attribution. Alpha-level feature.
  async_checkpoint_save: >
    Non-blocking checkpoint writes using persistent workers. Overlaps
    save I/O with training compute.
  local_checkpointing: >
    Fast local checkpoint save/load using nvidia-resiliency-ext local
    checkpointing with replication strategies.
recommended_path:
  fault_tolerance:
    ft.enable_ft_package: true
    ft.calc_ft_timeouts: true
    plugin: FaultTolerancePlugin
    launcher: ft_launcher
  straggler_detection:
    nvrx_straggler.enabled: true
    nvrx_straggler.report_time_interval: 300.0
    nvrx_straggler.stop_if_detected: false
  async_checkpoint:
    checkpoint.async_save: true
    checkpoint.ckpt_format: torch_dist
known_constraints:
  - Fault tolerance requires Slurm and ft_launcher (not torchrun).
  - Async save requires ckpt_format=torch_dist.
  - In-process restart requires PyTorch >= 2.5.1 and NCCL >= 2.26.2.
  - In-process restart is incompatible with NeMo-Run and Slurm preemption.
  - NVRx straggler and legacy StragglerDetector should not both be enabled.
  - NCCL watchdog timeout must exceed in-process restart hard_timeout.
  - nvidia-resiliency-ext (~0.5.0) is required for FT, straggler, IPR, and local ckpt.
known_limitations:
  - No torchrun-based fault tolerance path exists.
  - Re-run state machine integration is alpha-level.
  - In-process restart functional test is excluded from default CI.
  - Not all recipes enable resiliency features by default.
  - Preemption plugin is Slurm-specific.
evidence:
  - docs/training/resiliency.md
  - docs/training/checkpointing.md
  - src/megatron/bridge/training/fault_tolerance.py
  - src/megatron/bridge/training/nvrx_straggler.py
  - src/megatron/bridge/training/inprocess_restart.py
  - src/megatron/bridge/training/checkpointing.py
  - src/megatron/bridge/training/config.py
  - src/megatron/bridge/training/utils/sig_utils.py
  - src/megatron/bridge/recipes/run_plugins.py
  - scripts/performance/nemo-mbridge-resiliency_plugins.py
  - tests/unit_tests/training/test_fault_tolerance.py
  - tests/unit_tests/training/test_nvrx_straggler.py
  - tests/unit_tests/training/test_inprocess_restart.py
  - tests/unit_tests/recipes/test_run_plugins.py
  - tests/functional_tests/training/test_nvrx_straggler.py
  - tests/functional_tests/training/test_inprocess_restart.py
  - tests/functional_tests/training/test_local_checkpointing.py
  - examples/training_features/nemo-mbridge-resiliency/fault_tolerance/
  - examples/training_features/nemo-mbridge-resiliency/straggler_detection/
follow_up_validation:
  - Add a Slurm-based FT end-to-end test with actual hang recovery.
  - Add a checked-in recipe that enables FT + straggler detection by default.
  - Validate in-process restart on current container and NCCL versions.
  - Promote re-run state machine from alpha once runtime integration is complete.
  - Add benchmark for async save overhead vs sync save.

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Needs Python, Needs Node.js

Security

What our scanner noticed· Medium risk

  • Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Published by
NVIDIA
Publisher
NVIDIA (Official)
Source
NVIDIA/skills · skills/nemo-mbridge-resiliency · commit fce3cca
Licence
Apache License 2.0 (redistribution allowed)
  • LICENSE (added by AutoGineer from LICENSE-APACHE)
  • ATTRIBUTION.md (written by AutoGineer)
Synced

AutoGineer is not affiliated with or endorsed by NVIDIA or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download nemo-mbridge-resiliency.zip