Megatron桥接弹性Skill nemo-mbridge-resiliency

本技能介绍Megatron Bridge中的弹性(Resiliency)功能,涵盖故障容错、滞后者(straggler)检测、进程内重启、抢占和重新运行状态机,并提供完整的配置说明、插件使用、代码锚点、常见陷阱与验证方法。用于确保大规模分布式训练任务在遭遇节点故障、性能不均衡或外部抢占时能够自动恢复、稳定运行。关键词:Megatron Bridge, 弹性, 容错, 滞后者检测, 挂起检测, 自动重启, 进程内重启, 抢占, 重新运行状态机, NVRx, NeMo, 异步检查点, 分布式训练。

大模型训练框架 0 次安装 0 次浏览 更新于 9/7/2026
名称 nemo-mbridge-resiliency
描述 Megatron Bridge中的弹性特征,包括故障容错、滞后者检测、进程内重启、抢占以及重新运行状态机。
开源协议 Apache-2.0 when_to_use: 启用弹性特性,或调查导致训练挂起、滞后者检测失败或重启失败的提交;相关关键词:故障容错、滞后者检测、挂起检测、自动重启、进程内重启、抢占、nvidia-resiliency-ext。

弹性

稳定文档:@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_restarts 3 同一作业内的最大重启次数
num_job_retries_on_failure 2 失败时最多重新启动作业次数
initial_rank_heartbeat_timeout 1800 首个心跳超时(秒)
rank_heartbeat_timeout 300 后续心跳超时(秒)

选项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_package False 启用故障容错
calc_ft_timeouts False 自动计算最优超时时间
simulate_fault False 启用故障模拟以进行测试
simulated_fault_type "random" "rank_hung""rank_killed""random"
simulated_fault_rank None 指定故障排名(如果为None则随机)
simulated_fault_base_delay 0 模拟故障前的基础延迟

基于部分的超时监控独立覆盖设置、训练步骤、检查点和部分外时间。当 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,
)
参数 默认值 描述
enabled False 启用滞后者检测
report_time_interval 300.0 滞后者检查之间的秒数
calc_relative_gpu_perf True 比较各排名间的相对性能
calc_individual_gpu_perf True 追踪每个排名随时间的性能退化
gpu_relative_perf_threshold 0.7 相对性能阈值(0-1)
gpu_individual_perf_threshold 0.7 单独性能阈值(0-1)
stop_if_detected False 检测到滞后者时终止训练
num_gpu_perf_scores_to_print 5 打印最佳/最差分数数量
profiling_interval 1 检测器的性能分析间隔

抢占

插件(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_time 60 在作业限制前发送信号的秒数
enable_exit_handler True 在训练中启用信号处理程序
enable_exit_handler_for_data_loader False 为数据加载器工作进程启用

直接配置

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_loss True 检查损失中是否包含NaN
check_for_spiky_loss False 检查是否出现异常大的损失
spiky_loss_factor 10.0 如果损失超过 factor × 最大值,则标记为尖峰(大模型可调大)

退出码: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,
)
参数 默认值 描述
enabled False 启用进程内重启
active_world_size None 执行工作负载的排名(其余为热备)
granularity "node" "node""rank" 重启粒度
max_iterations None 最大重启尝试次数(None = 不限)
soft_timeout 60.0 检测GIL释放挂起(秒)
hard_timeout 90.0 强制终止挂起的排名(秒)
heartbeat_interval 30.0 心跳间隔(秒)
heartbeat_timeout 60.0 心跳超时时间(秒)
barrier_timeout 120.0 分布式屏障超时(秒)
completion_timeout 120.0 完成屏障超时(秒)
empty_cuda_cache True 重启期间清空CUDA缓存
max_rank_faults None 终止前允许的最大排名故障数
monitor_process_logdir None 监控进程日志目录

必需的环境变量:

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.pyFaultToleranceConfig
  • 运行时:src/megatron/bridge/training/fault_tolerance.py
  • 插件:src/megatron/bridge/recipes/run_plugins.pyFaultTolerancePlugin
  • 性能插件: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.pyNVRxStragglerDetectionConfig
  • 运行时:src/megatron/bridge/training/nvrx_straggler.py
  • 训练循环:src/megatron/bridge/training/train.pycheck_nvrx_straggler_detection
  • 测试:tests/unit_tests/training/test_nvrx_straggler.pytests/functional_tests/training/test_nvrx_straggler.py
  • 示例:examples/training_features/nemo-mbridge-resiliency/straggler_detection/

进程内重启

  • 配置:src/megatron/bridge/training/config.pyInProcessRestartConfig
  • 运行时:src/megatron/bridge/training/inprocess_restart.py
  • 入口点:src/megatron/bridge/training/pretrain.pymaybe_wrap_for_inprocess_restart
  • 测试:tests/unit_tests/training/test_inprocess_restart.pytests/functional_tests/training/test_inprocess_restart.py

抢占

  • 插件:src/megatron/bridge/recipes/run_plugins.pyPreemptionPlugin
  • 信号处理:src/megatron/bridge/training/utils/sig_utils.py
  • 测试:tests/unit_tests/recipes/test_run_plugins.py

重新运行状态机

  • 配置:src/megatron/bridge/training/config.pyRerunStateMachineConfig
  • 初始化:src/megatron/bridge/training/initialize.pyinit_rerun_state

检查点

  • 异步保存:src/megatron/bridge/training/checkpointing.pyschedule_async_save
  • 本地检查点:src/megatron/bridge/training/checkpointing.pyLocalCheckpointManager
  • 测试:tests/functional_tests/training/test_local_checkpointing.py

常见陷阱

  1. 使用ft_launcher而非torchrun:直接使用 FaultToleranceConfig 需要 ft_launcher。使用 torchrun 会静默禁用故障容错。对于非Slurm环境,请设置 GROUP_RANK=0

  2. 异步保存需要torch_distasync_save=True 仅适用于 ckpt_format="torch_dist"。其他格式会静默失败或报错。

  3. 进程内重启与NeMo-Run不兼容:进程内重启与NeMo-Run或Slurm抢占插件不兼容。需要特定的PyTorch/NCCL版本和环境变量。

  4. NVRx与旧版滞后者检测的区别:存在两个检测器。请使用NVRx(nvrx_straggler);不要同时启用两者。

  5. stop_if_detected 默认为False:NVRx默认记录但不停止训练。如需自动终止,请设置为 stop_if_detected=True

  6. NCCL看门狗与 hard_timeout 的关系:对于进程内重启,NCCL看门狗超时时间必须超过 hard_timeout,否则PyTorch会在恢复前杀掉进程。

  7. 重新运行状态机处于alpha阶段:可使用 check_for_nan_in_loss=True 进行NaN检测,但不要依赖完整的重新运行工作流。

验证

故障容错

./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 performanceGPU individual performance 报告。

异步检查点

在日志中查找 Scheduling async checkpoint save。在写入检查点文件时,训练迭代应继续进行。

进程内重启

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

需要兼容的PyTorch/NCCL版本。