[ENHANCEMENT] Support Distributed Checkpoint Saving Fault-Tolerant (Weak Dependency)
- 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
Assessment
This issue has not been assessed yet.