NVIDIA / NVIDIA/Megatron-LM

Potential core dump during checkpoint saving when log_energy=False

Open
#1,870 0 comments 1 reaction 0 assignees View on GitHub
bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.