Add low-overhead GPU progress heartbeat monitoring

[lzy-edu]

PR Category

Train

PR Types

New Features

PR Description

This PR adds a standalone, low-overhead GPU heartbeat monitoring capability for distributed training.

The main changes include:

• Record per-rank GPU progress heartbeats from existing Megatron fault-tolerance integration points.
• Run heartbeat recording on the CPU side without CUDA synchronization.
• Collect low-frequency, node-local GPU health information through nvidia-smi.
• Correlate each training rank with its assigned GPU.
• Integrate heartbeat lifecycle management into the training runner and SSH launcher.
• Write heartbeat status and findings as machine-readable JSONL logs.
• Add unit tests and two-worker integration scenarios.

The heartbeat distinguishes between:

• Rank process liveness.
• GPU training progress.
• Low-frequency GPU hardware health status.

Usage

Enable heartbeat under experiment.runner in the training configuration:

experiment:
  runner:
    no_shared_fs: false
    heartbeat:
      enabled: true
      publish_interval_s: 5
      process_timeout_s: 30
      progress_timeout_s: 300
      checkpoint_timeout_s: 1800
      hardware_health:
        enabled: true
        interval_s: 60

After enabling the heartbeat, start the FlagScale training task with the existing launch command. No additional build step or separate monitoring command is required.

The runner automatically:

• Starts a lightweight heartbeat publisher for each training rank.
• Starts an optional GPU health collector on each node.
• Starts the heartbeat monitor on node 0.
• Stops the related monitoring processes when training finishes.

Results

Heartbeat monitoring results are written to:

<log_dir>/heartbeat/<run_id>/
├── status.json
├── findings.jsonl
├── rank_<rank>*pid*<pid>.heartbeat.jsonl
└── gpu_health_node_<node_rank>.json

The expected behavior is:

• During normal training, each rank is reported as alive, and its GPU training progress is reported as progressing. No timeout finding is generated.
• After training completes normally, each rank is reported as exited, and its GPU progress is reported as completed.
• If a rank process or its heartbeat thread becomes unresponsive, a rank_process_heartbeat_timeout finding is recorded.
• If a rank process remains alive but its training progress stops advancing, a gpu_progress_timeout finding is recorded.
• When hardware health collection is enabled, the health status of the GPU assigned to the affected rank is included as supporting evidence.

For example, when a rank remains alive but its training progress stalls, findings.jsonl contains a record similar to:

{
  "finding_type": "gpu_progress_timeout",
  "rank": 1,
  "process_liveness": "alive",
  "gpu_progress": "stalled",
  "reason": "rank_alive_but_gpu_progress_not_advancing"
}

These results make it possible to distinguish a rank process that has become unresponsive from a live rank whose GPU training progress has stalled, while also correlating the issue with optional GPU hardware health information.