NVIDIA-NeMo / NVIDIA-NeMo/RL

Upgrade to CUDA 13 base container and update torch indexes

Open
#2,111 0 comments 0 reactions 1 assignee Claimed by @chtruong814 View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

We currently use CUDA 12 in our base container, but upstream Megatron-LM and Megatron-Bridge have moved to `core_cu13` for transformer-engine. We're working around this by overriding `core_cu13` → `core_cu12` in several places, which should be cleaned up when we upgrade.

## What needs to change

### transformer-engine CUDA extra
The following files override `core_cu13` to `core_cu12` and need to be reverted to `core_cu13`:
- `pyproject.toml` — `transformer-engine[pytorch,core_cu12]`
- `3rdparty/Megatron-LM-workspace/setup.py` — `transformer-engine[pytorch,core_cu12]`
- `3rdparty/Megatron-Bridge-workspace/setup.py` — `transformer-engine[pytorch,core_cu12]`

### Workspace setup.py consistency check normalization
Both workspace `setup.py` files have `_normalize_te_cuda()` functions that normalize `core_cu13` → `core_cu12` so the consistency check against the submodule pyproject.toml doesn't fail. These should be removed:
- `3rdparty/Megatron-LM-workspace/setup.py`
- `3rdparty/Megatron-Bridge-workspace/setup.py`

### Base container and torch indexes
- Update the base Docker container to CUDA 13
- Update PyTorch wheel indexes from `cu12` to `cu13` variants
- Update `nvidia-cudnn-cu12` pin in `pyproject.toml` to the `cu13` equivalent

## Context

This was introduced during the transformers v5 upgrade (`hemil/automodel-transformers-v5` branch). The upstream submodules (Megatron-LM, Megatron-Bridge) default to CUDA 13, but our container is still on CUDA 12. The mismatch also causes cuDNN version conflicts (system cuDNN 9.10.1 vs pip nvidia-cudnn-cu12 9.19.0) which prevents TE's FusedAttention backend from loading correctly.

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.