NVIDIA CUDA 12.4: что нового для ML-инженеров

NVIDIA выпустила CUDA 12.4 в начале 2024 года, и с тех пор это стало стандартным окружением для большинства ML-команд. Если вы ещё на 11.x или ранних версиях 12.x — самое время разобраться, что изменилось и стоит ли переходить прямо сейчас.

cuDNN 9.x, идущий вместе с CUDA 12.4, получил переработанный backend для трансформерных архитектур. Основные изменения:

• Нативная поддержка Flash Attention 2 на уровне библиотеки — больше не нужно тащить отдельный пакет с расширениями на C++
• Улучшенный fused MHA (Multi-Head Attention) kernel для Hopper (H100) и Ampere (A100)
• Оптимизации для grouped query attention (GQA) — архитектура, которую используют Llama, Mistral и большинство современных LLM

До cuDNN 9.x разработчики использовали flash-attn как отдельный пакет, который нужно было компилировать под конкретную версию CUDA и PyTorch. Теперь Flash Attention доступна через стандартный torch.nn.functional.scaled_dot_product_attention при наличии cuDNN 9+.

import torch
import torch.nn.functional as F

# Автоматически использует Flash Attention при cuDNN 9+
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
    output = F.scaled_dot_product_attention(q, k, v, is_causal=True)

Рекомендую устанавливать через официальный репозиторий NVIDIA, а не через apt из стандартных реп:

# Добавляем репозиторий NVIDIA
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
dpkg -i cuda-keyring_1.1-1_all.deb
apt update

# Устанавливаем CUDA 12.4
apt install cuda-12-4

# Добавляем в PATH
echo 'export PATH=/usr/local/cuda-12.4/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

nvcc --version
nvidia-smi

PyTorch 2.2+ полностью поддерживает CUDA 12.4. Важный момент: версии PyTorch и CUDA должны совпадать. Используйте официальный матрикс совместимости.

# PyTorch 2.2 с CUDA 12.4
pip install torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cu124

# Проверка
python -c "import torch; print(torch.version.cuda); print(torch.cuda.is_available())"

torch.compile() в PyTorch 2.x работает значительно лучше с CUDA 12.4 за счёт улучшенного Triton backend.

Тест на A100 80GB, модель Llama-2−13B, batch size 8, sequence length 2048:

Прирост ~15−20% по сравнению с CUDA 11.8 достигается без каких-либо изменений в коде модели.

MIG позволяет разделить один GPU на несколько изолированных инстансов. Полезно, когда одна задача не утилизирует карту целиком.

# Включаем MIG mode (требует перезапуска GPU)
nvidia-smi -i 0 -mig 1

# Смотрим доступные профили
nvidia-smi mig -lgip

# Создаём два инстанса 3g.40gb (половина A100)
nvidia-smi mig -cgi 9,9 -C

# Проверяем
nvidia-smi -L

С CUDA 12.4 MIG работает стабильнее: исправлены проблемы с изоляцией памяти и контекстами CUDA при одновременном запуске нескольких процессов.

1. Не обновляйте на production без тестирования. Соберите тестовый стенд с такой же моделью и прогоните бенчмарки.

2. Версия драйвера важна. CUDA 12.4 требует драйвер >= 550.54. Проверьте: nvidia-smi | grep Driver.

3. Docker-образы. Используйте официальные образы nvidia/cuda:12.4.0-cudnn9-devel-ubuntu22.04 вместо самостоятельной установки — меньше проблем с зависимостями.

4. NCCL для multi-GPU. Обновите NCCL до 2.20+ вместе с CUDA 12.4, иначе distributed training может работать медленнее.

В целом, переход с CUDA 11.x на 12.4 — это одно из самых безболезненных обновлений, которое приносит ощутимый прирост производительности без необходимости переписывать код.