Potential core dump during checkpoint saving when log_energy=False
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Describe the bug**
I encountered a core dump issue during checkpoint saving in Megatron-LM when log_energyis not enabled. The problem appears to be related to missing energy logging initialization during the save process.
The root cause appears to be:
1. When log_energy=False, the NVML device handler is not initialized (handler=None). See [this](https://github.com/NVIDIA/Megatron-LM/blob/bcdd591fd6001fc47601886597ab1ee9074e0360/megatron/training/training.py#L2074).
2. During checkpoint saving, the code attempts to call the function `nvmlDeviceGetTotalEnergyConsumption` with uninitialized parameter handler(=None). See [code1](https://github.com/NVIDIA/Megatron-LM/blob/bcdd591fd6001fc47601886597ab1ee9074e0360/megatron/training/training.py#L1720) and [code2](https://github.com/NVIDIA/Megatron-LM/blob/bcdd591fd6001fc47601886597ab1ee9074e0360/megatron/core/energy_monitor.py#L62)
3. This leads to a segmentation fault/core dump
**Steps/Code to reproduce bug**
1. Configure Megatron-LM training without --log-energyflag
2. Run training until checkpoint saving is triggered
3. Observe core dump during checkpoint saving process
**Additional context**
I've confirmed that passing None to nvmlDeviceGetTotalEnergyConsumption will definitely cause a core dump.
```python3
>>> import pynvml
>>> from pynvml import nvmlDeviceGetTotalEnergyConsumption
>>> pynvml.nvmlInit()
>>> nvmlDeviceGetTotalEnergyConsumption(None)
Segmentation fault (core dumped)
```
My test environment is as follows:
pynvml version: 12.0.0
CUDA Version: 12.9
Driver Version: 570.86.15
Contributor guide
Research direction
Start in megatron/training/training.py at the checkpoint-saving and NVML initialization locations linked in the report, then read megatron/core/energy_monitor.py around line 62. Reproduce training with --log-energy omitted and a checkpoint triggered. Done means checkpoint saving no longer reaches the NVML call with an uninitialized handler and the no-energy-logging path completes without a core dump.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100