Upgrade to CUDA 13 base container and update torch indexes
- 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
Assessment
This issue has not been assessed yet.