Распределённое предобучение LLM с TorchTitan

Обеспечивает нативное для PyTorch распределённое предобучение LLM с использованием torchtitan и 4D-параллелизмом (FSDP2, TP, PP, CP). Используйте при предобучении Llama 3.1, DeepSeek V3 или пользовательских моделей в масштабе от 8 до 512+ GPU с Float8, torch.compile и распределённой контрольной точкой.

Метаданные навыка


Источник	Опционально — установка через `vibeos skills install official/mlops/torchtitan`
Путь	`optional-skills/mlops/torchtitan`
Версия	`1.0.0`
Автор	Orchestra Research
Лицензия	MIT
Зависимости	`torch>=2.6.0`, `torchtitan>=0.2.0`, `torchao>=0.5.0`
Платформы	linux, macos
Теги	`Model Architecture`, `Distributed Training`, `TorchTitan`, `FSDP2`, `Tensor Parallel`, `Pipeline Parallel`, `Context Parallel`, `Float8`, `Llama`, `Pretraining`

Справочник: полный SKILL.md

к сведению

Ниже приведено полное определение навыка, которое VibeOS загружает при его активации. Это те инструкции, которые видит агент, когда навык активен.

TorchTitan — нативное распределённое предобучение LLM на PyTorch

Быстрый старт

TorchTitan — официальная платформа PyTorch для крупномасштабного предобучения LLM с композируемым 4D-параллелизмом (FSDP2, TP, PP, CP), обеспечивающая ускорение на 65%+ по сравнению с базовыми подходами на GPU H100.

Установка:

# Из PyPI (стабильная версия)
pip install torchtitan

# Из исходников (новейшие функции, требуется PyTorch nightly)
git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt

Загрузка токенизатора:

# Получите HF-токен на https://huggingface.co/settings/tokens
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...

Запуск обучения на 8 GPU:

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh

Типовые рабочие процессы

Рабочий процесс 1: Предобучение Llama 3.1 8B на одном узле

Скопируйте этот чек-лист:

Предобучение на одном узле:
- [ ] Шаг 1: Загрузить токенизатор
- [ ] Шаг 2: Настроить обучение
- [ ] Шаг 3: Запустить обучение
- [ ] Шаг 4: Мониторинг и контрольные точки

Шаг 1: Загрузка токенизатора

python scripts/download_hf_assets.py \
  --repo_id meta-llama/Llama-3.1-8B \
  --assets tokenizer \
  --hf_token=ВАШ_HF_ТОКЕН

Шаг 2: Настройка обучения

Отредактируйте или создайте TOML-файл конфигурации:

# llama3_8b_custom.toml
[job]
dump_folder = "./outputs"
description = "Обучение Llama 3.1 8B"

[model]
name = "llama3"
flavor = "8B"
hf_assets_path = "./assets/hf/Llama-3.1-8B"

[optimizer]
name = "AdamW"
lr = 3e-4

[lr_scheduler]
warmup_steps = 200

[training]
local_batch_size = 2
seq_len = 8192
max_norm = 1.0
steps = 1000
dataset = "c4"

[parallelism]
data_parallel_shard_degree = -1  # Использовать все GPU для FSDP

[activation_checkpoint]
mode = "selective"
selective_ac_option = "op"

[checkpoint]
enable = true
folder = "checkpoint"
interval = 500

Шаг 3: Запуск обучения

# 8 GPU на одном узле
CONFIG_FILE="./llama3_8b_custom.toml" ./run_train.sh

# Или явно через torchrun
torchrun --nproc_per_node=8 \
  -m torchtitan.train \
  --job.config_file ./llama3_8b_custom.toml

Шаг 4: Мониторинг и контрольные точки

Логи TensorBoard сохраняются в ./outputs/tb/:

tensorboard --logdir ./outputs/tb

Рабочий процесс 2: Многоузловое обучение с SLURM

Многоузловое обучение:
- [ ] Шаг 1: Настроить параллелизм для масштабирования
- [ ] Шаг 2: Создать SLURM-скрипт
- [ ] Шаг 3: Отправить задачу
- [ ] Шаг 4: Возобновить с контрольной точки

Шаг 1: Настройка параллелизма для масштабирования

Для модели 70B на 256 GPU (32 узла):

[parallelism]
data_parallel_shard_degree = 32  # FSDP на 32 рангах
tensor_parallel_degree = 8        # TP внутри узла
pipeline_parallel_degree = 1      # Без PP для 70B
context_parallel_degree = 1       # Увеличить для длинных последовательностей

Шаг 2: Создание SLURM-скрипта

#!/bin/bash
#SBATCH --job-name=llama70b
#SBATCH --nodes=32
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8

srun torchrun \
  --nnodes=32 \
  --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  -m torchtitan.train \
  --job.config_file ./llama3_70b.toml

Шаг 3: Отправка задачи

sbatch multinode_trainer.slurm

Шаг 4: Возобновление с контрольной точки

Обучение автоматически возобновляется, если контрольная точка существует в настроенной папке.

Рабочий процесс 3: Включение Float8-обучения для H100

Float8 обеспечивает ускорение на 30-50% на GPU H100.

Float8-обучение:
- [ ] Шаг 1: Установить torchao
- [ ] Шаг 2: Настроить Float8
- [ ] Шаг 3: Запустить с compile

Шаг 1: Установка torchao

USE_CPP=0 pip install git+https://github.com/pytorch/ao.git

Шаг 2: Настройка Float8

Добавьте в ваш TOML-конфиг:

[model]
converters = ["quantize.linear.float8"]

[quantize.linear.float8]
enable_fsdp_float8_all_gather = true
precompute_float8_dynamic_scale_for_fsdp = true
filter_fqns = ["output"]  # Исключить выходной слой

[compile]
enable = true
components = ["model", "loss"]

Шаг 3: Запуск с compile

CONFIG_FILE="./llama3_8b.toml" ./run_train.sh \
  --model.converters="quantize.linear.float8" \
  --quantize.linear.float8.enable_fsdp_float8_all_gather \
  --compile.enable

Рабочий процесс 4: 4D-параллелизм для моделей 405B

4D-параллелизм (FSDP + TP + PP + CP):
- [ ] Шаг 1: Создать начальную контрольную точку
- [ ] Шаг 2: Настроить 4D-параллелизм
- [ ] Шаг 3: Запустить на 512 GPU

Шаг 1: Создание начальной контрольной точки

Необходимо для единообразной инициализации на всех стадиях PP:

NGPU=1 CONFIG_FILE=./llama3_405b.toml ./run_train.sh \
  --checkpoint.enable \
  --checkpoint.create_seed_checkpoint \
  --parallelism.data_parallel_shard_degree 1 \
  --parallelism.tensor_parallel_degree 1 \
  --parallelism.pipeline_parallel_degree 1

Шаг 2: Настройка 4D-параллелизма

[parallelism]
data_parallel_shard_degree = 8   # FSDP
tensor_parallel_degree = 8       # TP внутри узла
pipeline_parallel_degree = 8     # PP между узлами
context_parallel_degree = 1      # CP для длинных последовательностей

[training]
local_batch_size = 32
seq_len = 8192

Шаг 3: Запуск на 512 GPU

# 64 узла x 8 GPU = 512 GPU
srun torchrun --nnodes=64 --nproc_per_node=8 \
  -m torchtitan.train \
  --job.config_file ./llama3_405b.toml

Когда использовать vs альтернативы

Используйте TorchTitan когда:

Предобучаете LLM с нуля (от 8B до 405B+)
Нужно нативное для PyTorch решение без сторонних зависимостей
Требуется композируемый 4D-параллелизм (FSDP2, TP, PP, CP)
Обучаете на H100 с поддержкой Float8
Нужны совместимые контрольные точки с torchtune/HuggingFace

Используйте альтернативы когда:

Megatron-LM: Максимальная производительность для развёртываний только на NVIDIA
DeepSpeed: Более широкая экосистема оптимизации ZeRO, поддержка инференса
Axolotl/TRL: Тонкая настройка, а не предобучение
LitGPT: Образовательные цели, обучение меньшего масштаба

Частые проблемы

Проблема: Нехватка памяти на больших моделях

Включите активационную контрольную точку и уменьшите размер батча:

[activation_checkpoint]
mode = "full"  # Вместо "selective"

[training]
local_batch_size = 1

Или используйте накопление градиентов:

[training]
local_batch_size = 1
global_batch_size = 32  # Накопление градиентов

Проблема: TP вызывает высокое потребление памяти с асинхронными коллективами

Установите переменную окружения:

export TORCH_NCCL_AVOID_RECORD_STREAMS=1

Проблема: Float8-обучение не быстрее

Float8 выгоден только для больших GEMM. Отфильтруйте маленькие слои:

[quantize.linear.float8]
filter_fqns = ["attention.wk", "attention.wv", "output", "auto_filter_small_kn"]

Проблема: Загрузка контрольной точки не удаётся после изменения параллелизма

Используйте возможность перешардирования DCP:

# Преобразование шардированной контрольной точки в один файл
python -m torch.distributed.checkpoint.format_utils \
  dcp_to_torch checkpoint/step-1000 checkpoint.pt

Проблема: Инициализация конвейерного параллелизма

Сначала создайте начальную контрольную точку (см. Рабочий процесс 4, Шаг 1).

Поддерживаемые модели

Модель	Размеры	Статус
Llama 3.1	8B, 70B, 405B	Продуктовая
Llama 4	Различные	Экспериментальная
DeepSeek V3	16B, 236B, 671B (MoE)	Экспериментальная
GPT-OSS	20B, 120B (MoE)	Экспериментальная
Qwen 3	Различные	Экспериментальная
Flux	Diffusion	Экспериментальная

Бенчмарки производительности (H100)

Модель	GPU	Параллелизм	TPS/GPU	Техники
Llama 8B	8	FSDP	5 762	Базовый
Llama 8B	8	FSDP+compile+FP8	8 532	+48%
Llama 70B	256	FSDP+TP+AsyncTP	876	2D-параллелизм
Llama 405B	512	FSDP+TP+PP	128	3D-параллелизм

Продвинутые темы

Конфигурация FSDP2: См. references/fsdp.md для подробного сравнения FSDP2 vs FSDP1 и эквивалентов ZeRO.

Float8-обучение: См. references/float8.md для рецептов тензорного vs построчного масштабирования.

Контрольные точки: См. references/checkpoint.md для конвертации в HuggingFace и асинхронных контрольных точек.

Добавление пользовательских моделей: См. references/custom-models.md для протокола TrainSpec.

Ресурсы

GitHub: https://github.com/pytorch/torchtitan
Статья: https://arxiv.org/abs/2410.06511
ICLR 2025: https://iclr.cc/virtual/2025/poster/29620
Форум PyTorch: https://discuss.pytorch.org/c/distributed/torchtitan/44

Метаданные навыка​

Справочник: полный SKILL.md​

TorchTitan — нативное распределённое предобучение LLM на PyTorch

Быстрый старт​

Типовые рабочие процессы​

Рабочий процесс 1: Предобучение Llama 3.1 8B на одном узле​

Рабочий процесс 2: Многоузловое обучение с SLURM​

Рабочий процесс 3: Включение Float8-обучения для H100​

Рабочий процесс 4: 4D-параллелизм для моделей 405B​

Когда использовать vs альтернативы​

Частые проблемы​

Поддерживаемые модели​

Бенчмарки производительности (H100)​

Продвинутые темы​

Ресурсы​