nemo-mbridge-resiliency

韧性

稳定文档：@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	指定发生故障的 rank（若为 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	比较各 rank 之间的性能
calc_individual_gpu_perf	True	跟踪各 rank 随时间的性能下降情况
gpu_relative_perf_threshold	0.7	相对性能阈值（0-1）
gpu_individual_perf_threshold	0.7	单个 GPU 的性能阈值（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	若损失 > 因子 * 已观测到的最大值，则标记该损失（大型模型应调高此值）

退出代码：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	执行工作负载的 Rank（其余为备用 Rank）
granularity	"node"	"node" 或 "rank" 重启粒度
max_iterations	None	最大重启尝试次数（None = 不限）
soft_timeout	60.0	检测释放 GIL 的挂起（秒）
hard_timeout	90.0	强制终止挂起的 Rank（秒）
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	终止前允许的最大 Rank 故障数
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.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. 使用 ft_launcher，而非 torchrun：直接使用 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 与旧版 straggler 检测器：存在两种检测器。请使用 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 状态机仍处于 alpha 阶段：使用 check_for_nan_in_loss=True 检测 NaN，但暂时不要依赖完整的 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] 日志行。模拟故障应触发从检查点重启。

Straggler 检测

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

查找包含各 Rank 分数的 GPU relative performance 和 GPU individual performance 报告。

异步检查点

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

进程内重启

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

需要兼容的 PyTorch/NCCL 版本。