Lightning-AI / Lightning-AI/pytorch-lightning

Loading checkpoint from CLI using SLURM doesn't use GPU even though it says it does

Open
#20,689 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug environment: slurm ver: 2.5.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

When I load my checkpoint, it says LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

But I check my slurm jobstats and it's like

GPU utilization per node
stellar-m01g3 (GPU 0): 0% <--- GPU was not used

GPU memory usage per node - maximum used/total
stellar-m01g3 (GPU 0): 12.2GB/40.0GB (30.5%)

I even made sure to later put accelerator: gpu in the trainer section of the yaml file

### What version are you seeing the problem on?

v2.5

### How to reproduce the bug

```python
cli = ModelCLI(
subclass_mode_model=True,
subclass_mode_data=True,
parser_kwargs={"parser_mode": "omegaconf"},
save_config_callback=None,
)

trainer:
max_epochs: 10
accelerator: gpu
enable_progress_bar: False

#SBATCH --gres=gpu:1
#SBATCH --time=04:00:00

module purge
source ...

export SLURM_JOB_ID=$SLURM_JOB_ID

srun python -m specseg.models.train \
fit \
--config config/model/config_label.yaml \
--ckpt_path /path
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```

### More info

_No response_

cc @lantiga

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the ModelCLI invocation in specseg.models.train and the provided config/model/config_label.yaml, then inspect how the trainer's accelerator setting is passed under the srun command. Reproduce with the v2.5 setup and compare the reported LOCAL_RANK/CUDA_VISIBLE_DEVICES values with SLURM GPU utilization; done means the checkpoint-loading CLI job demonstrably uses the allocated GPU.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
cli, distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.