NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Support Distributed Checkpoint Saving Fault-Tolerant (Weak Dependency)

Open
#1,435 2 comments 0 reactions 0 assignees View on GitHub
community-request enhancement waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

### Description

In large-scale model training, checkpoint saving is a critical but potentially fragile step. Currently, if a checkpoint save fails (e.g., due to temporary storage issues, network hiccups, or quota limits), Megatron-LM will raise an exception and terminate the training process. This strict dependency on successful checkpointing can cause significant loss of training progress and waste of computational resources, especially for long-running or large-scale jobs.

---

### Motivation

- In real-world scenarios (cloud, shared storage, large clusters), checkpoint failures are often transient and not always critical.
- Strict checkpoint dependency means a single failure can waste days or weeks of compute.
- Many users would prefer to continue training and handle checkpoint issues separately, rather than lose all progress.

---

### Expected Behavior

- When a checkpoint save (including async finalize) fails, training should **not be interrupted** by default.
- The failure should be **logged clearly** (with warning or error level), and the training loop should continue.
- Optionally, provide configuration to control whether checkpoint failures are fatal or non-fatal (strict vs. weak mode).
- Optionally, support limited retries or fallback strategies.

---

### Suggested Implementation

- Wrap all checkpoint save and finalize logic in try/except blocks.
- On exception, log the error and continue training.
- Optionally, expose a config flag (e.g., `--checkpoint-strictness=weak/strict`) to control this behavior.

**Example Pseudocode:**
```python
try:
save_checkpoint(...)
except Exception as e:
print_rank_0(f\"[WARNING] Checkpoint save failed: {e}\")
# Optionally: retry, skip, or escalate
# Training continues
```
And for async finalize:
```python
try:
for finalize_fn in async_request.finalize_fns:
finalize_fn()
except Exception as e:
print_rank_0(f\"[WARNING] Checkpoint finalize failed: {e}\")
# Training continues
```

---

### Benefits

- Greatly improves robustness for long-running, large-scale, or cloud-based training jobs.
- Reduces risk of losing significant progress due to non-critical, transient checkpoint failures.
- Makes Megatron-LM more production-friendly and resource-efficient.

---

### Additional Context

- In distributed training, it is important that all ranks handle checkpoint failures in a synchronized and consistent manner to avoid resource deadlock or hanging at barriers.
- Related code: `save_checkpoint`, `maybe_finalize_async_save`, `AsyncCallsQueue`, etc.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.