Четыре видеокарты: DDP и FSDP на практике

Зачем это нужно

Теорию распределённого обучения вы прошли в уроке 5.4.3: параллелизм по данным, all-reduce, ZeRO и FSDP, пузыри конвейера, DualPipe. Там же, в уроке 5.6.1 и уроке 5.6.2 вы уже запускали DDP «по образцу». Этот урок — инженерный: один скрипт, который честно меряет, что дают 1, 2 и 4 карты вашей машины, и при этом доказуемо правильный.

Было в 5.4.3 и 5.6	Добавляем здесь
DDP и FSDP2 в виде отличий от train.py	полный скрипт: DDP и FSDP2, накопление, смешанная точность, перевычисление активаций
«свой seed на каждой карте»	глобальный батч, одинаковый при любом числе карт, и тест: кривые loss совпадают
torch.save на ранге 0	распределённые чекпоинты FSDP, продолжение обучения без потери шага, экспорт полных весов
формула времени all-reduce	модель «счёт против обмена», трасса профайлера с перекрытием, таблица и графики масштабирования
—	топология PCIe, типичные зависания и их отладка

В конце урока у вас будет готовый инструмент, чтобы ускорить претрейн из модуля 5.6 и дообучение на ступени 6.

Как torchrun запускает обучение

torchruntorchrunEN torchrunУтилита PyTorch для запуска распределённого обучения: стартует нужное число процессов и передаёт им RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR и MASTER_PORT через переменные окружения. Пример: torchrun --standalone --nproc_per_node=4 train.py. — это запускалка. Команда torchrun --standalone --nproc_per_node=4 train_dist.py стартует четыре независимых процесса Python с одним и тем же скриптом. Различаются они только переменными окружения:

Переменная	Значение у процессов	Зачем
RANK	0, 1, 2, 3	рангранг процессаEN rankНомер процесса в распределённом запуске, от 0 до world_size − 1. LOCAL_RANK — номер на своей машине, по нему выбирают GPU. Ранг 0 обычно печатает логи и сохраняет чекпоинты. — номер среди всех процессов (на всех машинах)
LOCAL_RANK	0, 1, 2, 3	номер на этой машине — по нему выбирают карту cuda:LOCAL_RANK
WORLD_SIZE	4	сколько процессов всего
MASTER_ADDR, MASTER_PORT	один адрес на всех	где процессы находят друг друга

--standalone означает «одна машина, свободный порт выбери сам». Без него при двух одновременных запусках второй упадёт с Address already in use. Чтобы занять две карты из четырёх, добавьте CUDA_VISIBLE_DEVICES=0,1 и --nproc_per_node=2.

Внутри скрипта dist.init_process_group("nccl") делает три вещи: процессы встречаются по MASTER_ADDR:MASTER_PORT, договариваются о составе группы, и библиотека NCCLNCCLEN NCCL (NVIDIA Collective Communications Library)Библиотека NVIDIA для коллективных операций между GPU (all-reduce, all-gather, reduce-scatter): сама выбирает путь — NVLink, PCIe напрямую (P2P), через память CPU или по сети. Бэкенд torch.distributed для видеокарт. строит маршруты обмена между картами — кольца и деревья для all-reduceall-reduce (суммирование со всех на все)EN all-reduceКоллективная операция — каждый GPU отдаёт свой тензор и получает сумму (или среднее) тензоров всех GPU. Так DDP усредняет градиенты., all-gatherall-gatherEN all-gatherКоллективная операция: каждый GPU отдаёт свой кусок тензора и получает склейку кусков всех GPU. FSDP собирает так веса слоя перед вычислением. и reduce-scatterreduce-scatterEN reduce-scatterКоллективная операция: тензоры всех GPU суммируются, и каждый GPU получает только свой кусок суммы. Вместе с all-gather даёт all-reduce; на этой паре построены ZeRO и FSDP.. Каждый процесс обязан до этого вызвать torch.cuda.set_device(LOCAL_RANK). Иначе все четыре процесса создадут свои тензоры на cuda:0: нулевая карта переполнится, остальные будут простаивать.

flowchart TB
TR["torchrun --nproc_per_node=4"]
TR --> P0["процесс RANK=0<br/>cuda:0"]
TR --> P1["процесс RANK=1<br/>cuda:1"]
TR --> P2["процесс RANK=2<br/>cuda:2"]
TR --> P3["процесс RANK=3<br/>cuda:3"]
P0 -.->|"NCCL"| P1
P1 -.->|"NCCL"| P2
P2 -.->|"NCCL"| P3
P3 -.->|"NCCL"| P0

Четыре процесса — по одному на карту. NCCL связывает их в кольцо для all-reduce

Ваша топология: PCIe без NVLink

Команда nvidia-smi topo -m печатает матрицу «как карта $i$ связана с картой $j$». Расшифровка обозначений из её легенды:

Обозначение	Путь между картами	Скорость
NV#	# линий NVLink напрямую	сотни ГБ/с — у вас не встретится
PIX, PXB	через один или несколько мостов PCIe, мимо процессора	быстрее всего для PCIe
PHB	через корневой мост PCIe процессора
NODE	между корневыми мостами внутри одного NUMA-узла процессора
SYS	между NUMA-узлами (сокетами) процессора	медленнее всего

На вашей машине ждите PHB, NODE или SYS. Вторая команда, nvidia-smi topo -p2p r, показывает, разрешён ли прямой обмен P2P (карта читает память соседней по PCIe напрямую). У GeForce он обычно запрещён драйвером (NS — not supported), и тогда NCCL гоняет данные через оперативную память процессора: карта → RAM → карта. Теоретический предел PCIe 4.0 x16 — около 32 ГБ/с в одну сторону («Ваша лаборатория»). Реальную шинную скорость all-reduce вы измеряли в уроке 5.4.3; дальше она обозначена $W$, а в примерах подставлено условное 10 ГБ/с.

Данные: каждому рангу — своё

Это место, где чаще всего ошибаются. В уроке 5.6.2 каждая карта брала случайные окна со своим seed 42 + rank. Это правильно, но с двумя недостатками: результат нельзя сравнить с одной картой шаг в шаг, а забыть + rank очень легко — и тогда все четыре карты учатся на одних и тех же окнах. Токенов в секунду вчетверо больше, а пользы — как от одной карты.

Наш вариант: глобальный батч шага генерируется одинаково на всех рангах (генератор зависит только от номера шага), и каждый ранг берёт свой кусок:

g = torch.Generator().manual_seed(SEED * 100_003 + step)        # один и тот же на всех рангах
ix = torch.randint(len(data) - T - 1, (GLOBAL,), generator=g)    # GLOBAL = MICRO · ACCUM · WORLD окон
part = (micro * world + rank) * MICRO                            # кусок ранга rank на микрошаге micro
ix = ix[part : part + MICRO]

При 4 картах без накопления ранг 0 берёт окна 0–15 глобального батча, ранг 1 — 16–31 и так далее. На одной карте с ACCUM=4 те же 64 окна набираются за четыре микрошага. Градиенты в обоих случаях — среднее по одним и тем же 64 окнам, поэтому кривые loss совпадают (урок 4.6.3: накопление эквивалентно большому батчу, урок 5.4.3: all-reduce тоже). А в начале обучения скрипт собирает индексы со всех рангов через all_gather и падает, если они одинаковые, — защита от ошибки с seed.

Если данные идут через DataLoader

Для датасета в виде Dataset (урок 3.3.2) то же делает DistributedSampler(dataset, num_replicas=world, rank=rank, shuffle=True): он раздаёт рангам непересекающиеся части. Не забудьте sampler.set_epoch(epoch) в начале каждой эпохи, иначе порядок перемешивания будет одним и тем же во всех эпохах.

Скрипт train_dist.py

Один файл на все случаи: одна карта, DDP или FSDP2 на 2 и 4 картах. Положите его рядом с model.py и input.txt из урока 4.6.2. Размер модели выбирается переменной SIZE: baby (10.8M), gpt2 (около 85M, как GPT-2 small) и big (около 1.2B — для FSDP).

"""train_dist.py — GPT из урока 4.6 на 1, 2 или 4 GPU: DDP или FSDP2, скорость, память, чекпоинты.

python train_dist.py                                               # одна карта (или CPU)
torchrun --standalone --nproc_per_node=4 train_dist.py             # DDP на 4 картах
MODE=fsdp torchrun --standalone --nproc_per_node=4 train_dist.py   # FSDP2 на 4 картах
Переменные окружения: MODE=ddp|fsdp, SIZE=tiny|baby|gpt2|big, STEPS, MICRO (окон на карту за микрошаг),
ACCUM (шагов накопления), ACT_CKPT=1, BUCKET_MB, SAVE_EVERY, RESUME=1, TAG.  Рядом — model.py и input.txt.
"""
import contextlib
import json
import os
import time

import torch
import torch.distributed as dist

from model import GPT, GPTConfig

# ── 1. Процессы и устройства ────────────────────────────────────────────────
distributed = "RANK" in os.environ                 # torchrun выставляет RANK, LOCAL_RANK, WORLD_SIZE
use_cuda = torch.cuda.is_available()
if distributed:
    dist.init_process_group("nccl" if use_cuda else "gloo")
    rank, world = dist.get_rank(), dist.get_world_size()
    local_rank = int(os.environ["LOCAL_RANK"])
else:
    rank, world, local_rank = 0, 1, 0
device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
if use_cuda:
    torch.cuda.set_device(device)                  # без этого все процессы полезут на cuda:0
master = rank == 0


def log(*args):
    if master:                                     # печатает только ранг 0, иначе вывод ×4
        print(*args, flush=True)


# ── 2. Настройки ────────────────────────────────────────────────────────────
MODE = os.environ.get("MODE", "ddp")
SIZE = os.environ.get("SIZE", "baby" if use_cuda else "tiny")
STEPS = int(os.environ.get("STEPS", 300))
MICRO = int(os.environ.get("MICRO", 16))           # окон на одну карту за один микрошаг
ACCUM = int(os.environ.get("ACCUM", 1))            # микрошагов на один шаг оптимизатора
SAVE_EVERY = int(os.environ.get("SAVE_EVERY", 0))  # 0 — не сохранять чекпоинты
SEED = 1337
TAG = os.environ.get("TAG", f"{MODE}_{SIZE}_w{world}")
sizes = {
    "tiny": dict(n_layer=2, n_head=2, n_embd=64, block_size=64),
    "baby": dict(n_layer=6, n_head=6, n_embd=384, block_size=256),     # 10.8M, урок 4.6.2
    "gpt2": dict(n_layer=12, n_head=12, n_embd=768, block_size=256),   # ≈ 85M
    "big": dict(n_layer=24, n_head=16, n_embd=2048, block_size=256),   # ≈ 1.2B — для FSDP
}

# ── 3. Данные: символы из input.txt (урок 4.6.2) ────────────────────────────
RU = "абвгдеёжзийклмнопрстуфхцчшщъыьэюя"
EN = "abcdefghijklmnopqrstuvwxyz"
ALLOWED = set(RU + RU.upper() + EN + EN.upper() + "0123456789" + " \n.,:;!?-—«»()")
text = "".join(ch for ch in open("input.txt", encoding="utf-8").read() if ch in ALLOWED)
chars = sorted(set(text))
stoi = {ch: i for i, ch in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text], dtype=torch.long).to(device)   # корпус целиком на карте
cfg = GPTConfig(vocab_size=len(chars), dropout=0.0, **sizes[SIZE])
T = cfg.block_size
GLOBAL = MICRO * ACCUM * world                     # эффективный батч в окнах (урок 4.6.3)


def get_batch(step, micro):
    """Глобальный батч шага одинаков на всех рангах (общий seed), каждый ранг берёт СВОЙ кусок."""
    g = torch.Generator().manual_seed(SEED * 100_003 + step)
    ix = torch.randint(len(data) - T - 1, (GLOBAL,), generator=g)
    part = (micro * world + rank) * MICRO           # кусок ранга rank на микрошаге micro
    ix = ix[part : part + MICRO].to(device)
    chunk = data[ix[:, None] + torch.arange(T + 1, device=device)]   # (MICRO, T + 1) одним обращением
    return chunk[:, :-1].contiguous(), chunk[:, 1:].contiguous(), ix   # x и y — окно со сдвигом на 1


# ── 4. Модель: DDP или FSDP2 ────────────────────────────────────────────────
torch.manual_seed(SEED)                            # одинаковые веса на всех рангах
model = GPT(cfg).to(device)
n_params = sum(p.numel() for p in model.parameters())
if os.environ.get("ACT_CKPT") == "1":              # перевычисление активаций (урок 5.4.3)
    from torch.utils.checkpoint import checkpoint
    for blk in model.blocks:
        blk.forward = (lambda f: lambda x: checkpoint(f, x, use_reentrant=False))(blk.forward)
amp = torch.autocast(device_type="cuda", dtype=torch.bfloat16) if use_cuda else contextlib.nullcontext()
if distributed and MODE == "ddp":
    from torch.nn.parallel import DistributedDataParallel as DDP
    model = DDP(model, device_ids=[local_rank] if use_cuda else None,
                bucket_cap_mb=int(os.environ.get("BUCKET_MB", 25)),   # размер «корзины» градиентов
                gradient_as_bucket_view=True)       # градиенты прямо в буферах корзин: меньше копий
elif distributed and MODE == "fsdp":
    from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
    mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) if use_cuda else None
    kw = {"mp_policy": mp} if mp else {}
    for blk in model.blocks:
        fully_shard(blk, **kw)                     # каждый блок — своя единица шардирования
    fully_shard(model, **kw)                       # остальное: эмбеддинги, ln_f, lm_head
    amp = contextlib.nullcontext()                 # точность задаёт mp_policy, autocast не нужен
raw = model.module if MODE == "ddp" and distributed else model
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.1)
log(f"{TAG}: {n_params / 1e6:.1f}M параметров, {world} проц., глобальный батч {GLOBAL} × {T} токенов")

# ── 5. Чекпоинты ────────────────────────────────────────────────────────────
ckpt_dir = f"ckpt_{TAG}"


def save(step):
    if MODE == "fsdp" and distributed:             # у каждого ранга только свой шард
        import torch.distributed.checkpoint as dcp
        from torch.distributed.checkpoint.state_dict import get_state_dict
        msd, osd = get_state_dict(model, opt)
        dcp.save({"model": msd, "opt": osd}, checkpoint_id=f"{ckpt_dir}/step{step}")   # пишут все ранги
    elif master:                                   # DDP: копии одинаковые, пишет один ранг
        os.makedirs(ckpt_dir, exist_ok=True)
        torch.save({"model": raw.state_dict(), "opt": opt.state_dict()}, f"{ckpt_dir}/step{step}.pt")
    if distributed:
        dist.barrier()                             # никто не идёт дальше, пока чекпоинт не записан


def load(step):
    if MODE == "fsdp" and distributed:
        import torch.distributed.checkpoint as dcp
        from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
        msd, osd = get_state_dict(model, opt)
        state = {"model": msd, "opt": osd}
        dcp.load(state, checkpoint_id=f"{ckpt_dir}/step{step}")
        set_state_dict(model, opt, model_state_dict=state["model"], optim_state_dict=state["opt"])
    else:
        sd = torch.load(f"{ckpt_dir}/step{step}.pt", map_location=device)
        raw.load_state_dict(sd["model"])
        opt.load_state_dict(sd["opt"])


start = 0
if os.environ.get("RESUME") == "1":
    start = max(int(f.split("step")[1].split(".")[0]) for f in os.listdir(ckpt_dir))
    load(start)
    log(f"продолжаем с шага {start}")

# ── 6. Обучение ─────────────────────────────────────────────────────────────
prof = None
if os.environ.get("PROFILE") == "1":               # трасса шагов 8–10 каждого ранга (урок 5.7.1)
    from torch.profiler import ProfilerActivity, profile, schedule
    acts = [ProfilerActivity.CPU] + ([ProfilerActivity.CUDA] if use_cuda else [])
    prof = profile(activities=acts, schedule=schedule(wait=6, warmup=2, active=3))
    prof.start()
losses, t_start = [], None
for step in range(start, STEPS):
    if step == start + 5:                          # первые шаги — прогрев: компиляция, выделение памяти
        if use_cuda:
            torch.cuda.synchronize()
            torch.cuda.reset_peak_memory_stats()
        t_start = time.perf_counter()
    loss_sum = torch.zeros((), device=device)
    for micro in range(ACCUM):
        last = micro == ACCUM - 1
        if MODE == "ddp" and distributed:          # all-reduce только на последнем микрошаге
            sync = contextlib.nullcontext() if last else model.no_sync()
        else:
            sync = contextlib.nullcontext()
            if MODE == "fsdp" and distributed:
                model.set_requires_gradient_sync(last)
        xb, yb, ix = get_batch(step, micro)
        if step == start and micro == 0 and distributed:   # проверка: у рангов разные окна?
            all_ix = [torch.zeros_like(ix) for _ in range(world)]
            dist.all_gather(all_ix, ix)
            assert len(set(torch.cat(all_ix).tolist())) > MICRO, "все ранги получили одни и те же данные!"
        with sync:                                 # и прямой, и обратный проход — внутри
            with amp:
                _, loss = model(xb, yb)
            (loss / ACCUM).backward()
        loss_sum += loss.detach() / ACCUM
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    opt.zero_grad(set_to_none=True)
    if distributed:                                # средний loss по всем рангам — для лога
        dist.all_reduce(loss_sum)
        loss_sum /= world
    losses.append(loss_sum)                        # .item() на каждом шаге заставил бы CPU ждать GPU
    if step % 50 == 0 or step == STEPS - 1:
        log(f"шаг {step:4d}  loss {loss_sum.item():.4f}")
    if SAVE_EVERY and (step + 1) % SAVE_EVERY == 0:
        save(step + 1)
    if prof is not None:
        prof.step()
if prof is not None:
    prof.stop()
    prof.export_chrome_trace(f"trace_{TAG}_rank{rank}.json")

# ── 7. Итоги: скорость, память, кривая loss ─────────────────────────────────
if use_cuda:
    torch.cuda.synchronize()
n_timed = STEPS - start - 5
if t_start is not None and n_timed > 0:
    dt = (time.perf_counter() - t_start) / n_timed
    mem = torch.cuda.max_memory_allocated() / 2**30 if use_cuda else 0.0
    if distributed:                                # память — максимум по рангам
        m = torch.tensor([mem], device=device)
        dist.all_reduce(m, op=dist.ReduceOp.MAX)
        mem = m.item()
    tok_s = round(GLOBAL * T / dt)
    res = dict(tag=TAG, mode=MODE, size=SIZE, world=world, accum=ACCUM, micro=MICRO, step_ms=round(dt * 1000, 2),
               tok_s=tok_s, mem_gib=round(mem, 3), losses=[round(x.item(), 6) for x in losses])
    log(f"{TAG}: шаг {dt * 1000:.1f} мс, {tok_s} токенов/с, пик памяти {mem:.2f} ГиБ")
    if master:
        with open("results.jsonl", "a") as f:
            f.write(json.dumps(res) + "\n")
# полные веса одним файлом — для генерации обычной GPT на одной карте (урок 4.6.2)
if MODE == "fsdp" and distributed:
    from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
    full = get_model_state_dict(model, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
else:
    full = raw.state_dict()
if master:                                         # все ранги участвуют в сборке, пишет один
    torch.save({"model": full, "cfg": cfg, "chars": chars}, f"gpt_{TAG}.pt")
if distributed:
    dist.destroy_process_group()

Разберём то, чего не было в прошлых уроках.

DDP. Модель оборачивается в DistributedDataParallel; при создании он рассылает веса ранга 0 всем остальным. Градиенты собираются в «корзины» по bucket_cap_mb мегабайт, и all-reduce готовой корзины стартует, пока обратный проход ещё считает следующие слои (урок 5.4.3). gradient_as_bucket_view=True кладёт градиенты прямо в буферы корзин — минус одна копия и немного памяти.

FSDP2. fully_shard применяется к каждому блоку, затем к корню: перед вычислением блока его веса собираются со всех карт (all-gather) и после освобождаются, градиенты расходятся по владельцам (reduce-scatter). MixedPrecisionPolicy задаёт точность: собранные веса — в BF16 (вдвое меньше байт по PCIe и считаем на тензорных ядрах), усреднение градиентов — во FP32, сами шарды весов и состояние AdamW — во FP32. Поэтому autocast здесь не нужен.

Накопление. На всех микрошагах, кроме последнего, обмен выключен: у DDP — контекст no_sync() (он должен охватывать и прямой проход: решение о синхронизации DDP принимает в forward), у FSDP2 — set_requires_gradient_sync(False). Один all-reduce на шаг оптимизатора вместо ACCUM штук.

Остальное. clip_grad_norm_ работает и с DDP (градиенты уже одинаковые), и с FSDP2 (параметры — распределённые тензоры, норма считается по всем шардам). Loss для лога усредняется по рангам одним all-reduce, а .item() вызывается только при печати: на каждом шаге он заставил бы процессор ждать видеокарту. ACT_CKPT=1 включает перевычисление активацийперевычисление активаций (activation checkpointing)EN activation checkpointing (recomputation)Не хранить часть активаций прямого прохода, а пересчитать их при обратном. Экономит память ценой дополнительных вычислений. для каждого блока. Обратите внимание на .contiguous() в get_batch: срез chunk[:, 1:] не лежит в памяти подряд, и view внутри model.py упал бы с ошибкой «view size is not compatible».

Чекпоинты. У DDP копии одинаковые, и файл пишет только ранг 0, а dist.barrier() не даёт остальным уйти вперёд. У FSDP полной модели нет ни у кого, поэтому сохраняет распределённый чекпоинтраспределённый чекпоинт (DCP)EN distributed checkpoint (torch.distributed.checkpoint)Сохранение состояния модели и оптимизатора, при котором каждый ранг пишет свой шард параллельно. Позволяет загрузить чекпоинт и на другом числе GPU. Нужен для FSDP, где ни у одного ранга нет полной модели.: каждый ранг параллельно пишет свой шард (файлы __0_0.distcp, __1_0.distcp, …) и общий файл метаданных. Такой чекпоинт можно загрузить и на другом числе карт. RESUME=1 продолжает с последнего сохранения, а в конце скрипт собирает полные веса одним файлом gpt_<тег>.pt — его загрузит обычная GPT из урока 4.6.2 для генерации.

Проверка: кривые loss совпадают

Прежде чем мерить скорость, убедимся, что распределённое обучение считает то же самое. Запустим на процессоре (бэкенд gloo вместо nccl, крошечная модель SIZE=tiny) пять вариантов с одинаковым глобальным батчем в 64 окна и сравним кривые отчётом report.py (код — в блоке «На вашей машине» ниже). На машине с видеокартами добавьте к командам CUDA_VISIBLE_DEVICES="", чтобы скрипт ушёл на процессор:

STEPS=30 ACCUM=4 TAG=ref_acc4 python train_dist.py                                     # 1 процесс, накопление 4
STEPS=30 TAG=ddp_tiny_w4 torchrun --standalone --nproc_per_node=4 train_dist.py       # DDP, 4 процесса
STEPS=30 MODE=fsdp TAG=fsdp_tiny_w4 torchrun --standalone --nproc_per_node=4 train_dist.py
STEPS=30 ACCUM=2 MODE=fsdp ACT_CKPT=1 TAG=fsdp_ckpt_w2 torchrun --standalone --nproc_per_node=2 train_dist.py
python report.py

ddp_tiny_w4 против ref_acc4 (батч 64): max |Δloss| = 1.0e-06
fsdp_ckpt_w2 против ref_acc4 (батч 64): max |Δloss| = 1.0e-06
fsdp_tiny_w4 против ref_acc4 (батч 64): max |Δloss| = 1.0e-06

Расхождение $10^{-6}$ — это округление в логе: все варианты идут по одной кривой. Так же проверено продолжение после чекпоинта: 20 шагов с SAVE_EVERY=10, затем RESUME=1 до 30-го шага — и шаги 20–29 совпали с непрерывным прогоном с той же точностью $10^{-6}$ и у DDP, и у FSDP (данные зависят только от номера шага, а состояние AdamW сохранено целиком).

На видеокартах в BF16 совпадение будет не до шестого знака: другой порядок суммирования, недетерминированные ядра и разная точность у DDP с autocast и FSDP с MixedPrecisionPolicy. Расхождение порядка $10^{-3}$–$10^{-2}$, медленно растущее со временем, — норма. Расхождение на 0.1 и больше уже на первых шагах — ошибка: одинаковые данные на рангах, забытое деление на ACCUM, не тот глобальный батч.

Масштабирование и PCIe

Когда обмен перестаёт прятаться? Для DDP с одинаковым батчем на карту за шаг оптимизатора:

• счёт одной карты: $6 N \cdot \text{токенов} / (P \cdot \mathrm{MFU})$ — по правилу $6N$ FLOPs на токен (урок 5.1.1);
• обмен: all-reduce градиентов, $2 \cdot \frac{n-1}{n} \cdot 4N / W$ секунд при $n$ картах и градиентах во FP32 (урок 5.4.3);
• обмен прячется за обратным проходом последнего микрошага (около $2/3$ счёта), на виду остаётся только избыток.

Поделим одно на другое — $N$ сокращается:

$$\frac{t_{\text{обмен}}}{t_{\text{счёт}}} \approx \frac{1.5 \cdot 4 / W}{6 \cdot \text{токенов} / (P \cdot \mathrm{MFU})} = \frac{P \cdot \mathrm{MFU}}{W \cdot \text{токенов}}.$$

Читаем: «доля обмена зависит не от размера модели, а от того, сколько токенов карта обрабатывает между двумя all-reduce». При $P \cdot \mathrm{MFU} = 0.4 \cdot 97.5 = 39$ TFLOPS и $W = 10$ ГБ/с отношение ≈ $3900 / \text{токенов}$. При батче 16 окон по 256 токенов (4096 токенов) обмен почти равен счёту, и после перекрытия эффективность масштабированияэффективность масштабированияEN scaling efficiencyСкорость на N GPU, делённая на N скоростей одной GPU. 100% — идеальное ускорение в N раз; потери дают коммуникации, не спрятанные за вычислениями, и неравномерная загрузка. около 78%. С накоплением ACCUM=4 (16 384 токена на all-reduce) — около 93%. Эти числа вы получите в практике ниже.

Для FSDP есть приятная деталь. За шаг он делает два all-gather весов (в прямом и обратном проходе) и один reduce-scatter градиентов — в 1.5 раза больше пересылок, чем DDP (урок 5.4.3). Но с весами в BF16 и градиентами во FP32 байт выходит $\frac{n-1}{n} N (2 \cdot 2 + 4) = 2 \cdot \frac{n-1}{n} \cdot 4N$ — ровно столько же, сколько у DDP во FP32. Платим только задержками: обмен идёт мелкими порциями на каждом блоке.

Как увидеть перекрытие. Запустите PROFILE=1 — каждый ранг сохранит трассу шагов 8–10. В Perfetto (урок 5.7.1) под строкой вычислительного потока GPU появится вторая строка с ядрами NCCL (ncclDevKernel_AllReduce... у DDP, AllGather и ReduceScatter у FSDP). Если они идут параллельно ядрам обратного прохода — обмен спрятан. Если после последнего ядра backward тянется хвост из NCCL, а ядро шага оптимизатора ждёт его — это и есть открытый обмен, который ест эффективность.

Практика

Перенесите эти рассуждения в код: разбиение глобального батча, DDP-градиент с накоплением (тест сверит его с градиентом по всему батчу), время all-reduce и шага с перекрытием, эффективность и трафик FSDP. Подставьте потом свою шинную скорость.

Модель DDP: данные, градиенты и время шагаnumpy · нужен интернет7 тестов

import numpy as np


# ── Готовые помощники ────────────────────────────────────────────────────────
def global_batch(step, global_size, n_data, seed=1337):
    """Индексы глобального батча шага step. Генератор зависит только от (seed, step),
    поэтому на всех рангах получается один и тот же список — как в train_dist.py."""
    rng = np.random.default_rng([seed, step])
    return rng.integers(0, n_data, size=global_size)


def mse_grad(w, X, y):
    """Градиент средней квадратичной ошибки линейной модели X @ w по w (урок 2.2.1)."""
    return 2 * X.T @ (X @ w - y) / len(y)


# ── Ваша часть ───────────────────────────────────────────────────────────────
def rank_slice(ix, rank, world, micro_idx, micro):
    """Кусок глобального батча для ранга rank на микрошаге micro_idx.
    Раскладка как в train_dist.py: ix[(micro_idx · world + rank) · micro : … + micro]."""
    # TODO
    ...


def ddp_grad(w, X, y, ix, world, accum, micro):
    """Градиент, который получит каждый ранг после DDP-шага с накоплением.

    Каждый ранг на каждом микрошаге считает mse_grad на своём куске и делит на accum
    (как loss / ACCUM), складывает по микрошагам; затем all-reduce: среднее по рангам.
    """
    # TODO: два цикла — по рангам и по микрошагам
    ...


def ring_allreduce_time(nbytes, world, bus_bw):
    """Время кольцевого all-reduce: 2·(world − 1)/world · nbytes / bus_bw (урок 5.4.3). На одной карте — 0."""
    # TODO
    ...


def ddp_step_time(t_fwd, t_bwd, grad_bytes, world, bus_bw, accum=1, overlap=True):
    """Время шага оптимизатора в DDP.

    Вычисления: accum · (t_fwd + t_bwd). Обмен — один all-reduce за шаг (no_sync на остальных микрошагах).
    С перекрытием all-reduce прячется за обратным проходом ПОСЛЕДНЕГО микрошага:
    на виду остаётся max(0, t_comm − t_bwd). Без перекрытия — весь t_comm.
    """
    # TODO
    ...


def scaling_efficiency(t_one, t_many):
    """Эффективность при одинаковом батче на карту: токены/с на N картах против N × (токены/с на одной).
    Каждая карта обрабатывает столько же токенов, поэтому это просто t_one / t_many."""
    # TODO
    ...


def comm_bytes(n_params, world, mode, param_bytes=2, grad_bytes=4):
    """Байт, которые уходят с одной карты за шаг.
    DDP: all-reduce градиентов — 2·(N−1)/N · Ψ · grad_bytes.
    FSDP: all-gather весов в прямом и в обратном проходе + reduce-scatter градиентов —
          (N−1)/N · Ψ · (2·param_bytes + grad_bytes)."""
    # TODO
    ...


if __name__ == "__main__":
    ix = global_batch(0, 8, 100)
    print("глобальный батч:", ix)
    print("ранг 1 из 2, микрошаг 0, по 2 окна:", rank_slice(ix, 1, 2, 0, 2))

ЗапуститьПроверитьПодсказкаСбросить

🖥 На вашей машине (4×GPU)

1. Топология и связь.

nvidia-smi topo -m          # матрица связей: PHB / NODE / SYS
nvidia-smi topo -p2p r      # разрешён ли прямой обмен P2P между картами

Шинную скорость all-reduce возьмите из урока 5.4.3 (allreduce_sweep.py, плато на больших сообщениях) — это ваше $W$.

2. Отчёт. Сохраните рядом со скриптом report.py — он строит таблицу скорости и памяти, сверяет кривые loss и рисует графики:

"""report.py — таблица скорости и памяти по results.jsonl, сверка кривых loss, графики.  python report.py"""
import json

import matplotlib.pyplot as plt

last = {}
for line in open("results.jsonl"):
    r = json.loads(line)
    last[r["tag"]] = r                             # берём последний запуск с каждым тегом
runs = sorted(last.values(), key=lambda r: (r["size"], r["mode"], r["world"]))

# 1. Скорость и память. Эффективность — токены/с на N картах против N × (токены/с на одной)
one = {(r["size"], r["micro"], r["accum"]): r["tok_s"] for r in runs if r["world"] == 1}
print(f"{'тег':22s} {'GPU':>3s} {'токенов/с':>10s} {'эффект.':>8s} {'память, ГиБ':>12s}")
for r in runs:
    b = one.get((r["size"], r["micro"], r["accum"]))
    eff = f"{r['tok_s'] / (r['world'] * b):.0%}" if b and r["world"] > 1 else "—"
    print(f"{r['tag']:22s} {r['world']:3d} {r['tok_s']:10d} {eff:>8s} {r['mem_gib']:12.2f}")

# 2. Сверка loss: запуски с одинаковым глобальным батчем должны идти по одной кривой
groups = {}
for r in runs:
    groups.setdefault((r["size"], r["micro"] * r["accum"] * r["world"]), []).append(r)
for (size, glob), rs in groups.items():
    ref = next((r for r in rs if r["world"] == 1), None)
    for r in rs:
        if ref is not None and r is not ref:
            n = min(len(r["losses"]), len(ref["losses"]))
            d = max(abs(a - b) for a, b in zip(r["losses"][:n], ref["losses"][:n]))
            print(f"{r['tag']} против {ref['tag']} (батч {glob}): max |Δloss| = {d:.1e}")

# 3. Графики: токены/с и пиковая память от числа карт
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4))
for key in sorted({(r["size"], r["mode"], r["micro"], r["accum"]) for r in runs}):
    rs = [r for r in runs if (r["size"], r["mode"], r["micro"], r["accum"]) == key]
    ws = [r["world"] for r in rs]
    ax1.plot(ws, [r["tok_s"] for r in rs], "o-", label=f"{key[0]} {key[1]}")
    ax2.plot(ws, [r["mem_gib"] for r in rs], "o-", label=f"{key[0]} {key[1]}")
    b = one.get((key[0], key[2], key[3]))
    if b:
        ax1.plot([1, 2, 4], [b, 2 * b, 4 * b], ":", color="gray")   # идеальное масштабирование
ax1.set(xlabel="число GPU", ylabel="токенов/с", title="Скорость (пунктир — идеал)", xticks=[1, 2, 4])
ax2.set(xlabel="число GPU", ylabel="ГиБ на карту", title="Пиковая память", xticks=[1, 2, 4])
ax1.legend()
fig.tight_layout()
fig.savefig("scaling.png", dpi=120)
print("графики: scaling.png")

3. Замеры. Сначала baby GPT, одинаковый батч на карту (16 окон по 256 токенов), затем сверка кривых:

CUDA_VISIBLE_DEVICES=0 python train_dist.py
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 train_dist.py
torchrun --standalone --nproc_per_node=4 train_dist.py
CUDA_VISIBLE_DEVICES=0,1 MODE=fsdp torchrun --standalone --nproc_per_node=2 train_dist.py
MODE=fsdp torchrun --standalone --nproc_per_node=4 train_dist.py
CUDA_VISIBLE_DEVICES=0 ACCUM=4 TAG=ref_acc4 python train_dist.py      # тот же глобальный батч, что у 4 карт
python report.py

Повторите для SIZE=gpt2 и для SIZE=gpt2 ACCUM=4. Получите две картинки в scaling.png: токены в секунду от числа карт (пунктир — идеал) и пиковая память на карту.

Чего ждать (оценки по модели из урока, а не замеры: проверьте у себя). У baby GPT шаг короткий (миллисекунды), и на 4 картах заметную долю съедают и обмен, и задержки запуска: эффективность может оказаться ниже 80%. У gpt2 без накопления — около 75–80% при $W \approx 10$ ГБ/с, с ACCUM=4 — за 90%. Память DDP от числа карт не зависит; у FSDP она падает, но у маленьких моделей почти незаметно — основную память занимают активации, а не веса.

4. Большая модель: где DDP упирается в память. SIZE=big — около 1.2B параметров. В DDP только веса, градиенты и состояние AdamW занимают $16 \cdot 1.2 \approx 19$ ГБ на каждой карте, плюс активации и BF16-копии весов под autocast — впритык к 32 ГБ, первая команда может упасть с OOM. В FSDP на 4 картах — около 5 ГБ плюс активации:

SIZE=big MICRO=4 STEPS=60 torchrun --standalone --nproc_per_node=4 train_dist.py
SIZE=big MICRO=4 STEPS=60 MODE=fsdp torchrun --standalone --nproc_per_node=4 train_dist.py
SIZE=big MICRO=4 STEPS=60 MODE=fsdp ACT_CKPT=1 torchrun --standalone --nproc_per_node=4 train_dist.py
SIZE=big MICRO=16 STEPS=60 MODE=fsdp ACT_CKPT=1 torchrun --standalone --nproc_per_node=4 train_dist.py

Сравните пиковую память и токены в секунду. Перевычисление активаций стоит примерно ещё одного прямого прохода (около +30% ко времени), зато позволяет вчетверо поднять MICRO — а по формуле выше больше токенов на all-reduce значит выше эффективность. Обмен здесь тяжёлый: 4.8 ГБ градиентов во FP32 — это сотни миллисекунд на шаг по PCIe. Найдите конфигурацию с наибольшими токенами в секунду.

5. Трасса перекрытия. PROFILE=1 STEPS=12 SIZE=gpt2 torchrun --standalone --nproc_per_node=4 train_dist.py, затем откройте trace_ddp_gpt2_w4_rank0.json в Perfetto. Найдите ядра NCCL и проверьте, идут ли они параллельно обратному проходу. Повторите с BUCKET_MB=100 и BUCKET_MB=5: как меняется число и длина all-reduce?

6. Чекпоинты. SAVE_EVERY=100 STEPS=200 MODE=fsdp torchrun ..., прервите по Ctrl+C после сотого шага и продолжите с RESUME=1. Затем загрузите gpt_<тег>.pt в обычную GPT и сгенерируйте текст, как в уроке 4.6.2 (torch.load(path, weights_only=False): в файле лежит ещё и конфиг).

Настройка карт и проверка связи — на странице «Ваша лаборатория (4×GPU)».

Типичные ошибки

Симптом	Причина	Что делать
Все процессы на cuda:0, OOM на нулевой карте	нет torch.cuda.set_device(LOCAL_RANK)	вызывать до создания модели и тензоров
4 карты, токенов/с вчетверо больше, а loss падает как на одной	одинаковые данные на всех рангах	общий генератор + свой кусок или seed с + rank; проверка через all_gather
Зависание без ошибки, через 10 минут — таймаут NCCL	коллективную операцию вызвали не все ранги	не вызывать FSDP-модель только на ранге 0; одинаковое число шагов на всех рангах; печать и запись — под if rank == 0, а forward — у всех
Address already in use при старте	порт занят другим запуском	--standalone или свой --master_port
Чекпоинт битый или пустой	файл пишут все ранги разом / ранг 0 вышел раньше записи	писать на ранге 0 (DDP) или через DCP (FSDP), затем dist.barrier()
Чекпоинт не грузится в обычную модель	префиксы module. (DDP) и _orig_mod. (torch.compile)	сохранять model.module.state_dict(); для FSDP — get_model_state_dict(..., full_state_dict=True)
Loss на 4 картах хуже, чем на одной, при той же длине обучения	эффективный батч вырос в 4 раза, а число шагов и lr прежние	сравнивать при одинаковом глобальном батче или пересчитать lr и число шагов

Для отладки: NCCL_DEBUG=INFO (как NCCL построил связь, урок 5.4.3), TORCH_DISTRIBUTED_DEBUG=DETAIL (проверка, что ранги вызывают одинаковые коллективные операции с одинаковыми размерами), dist.init_process_group(..., timeout=timedelta(minutes=2)) (падать быстрее, а не ждать 10 минут) и py-spy dump --pid <PID> — увидеть, на какой строке висит каждый процесс.

Где это пригодится

• Претрейн из модуля 5.6. Модель L из урока 5.6.2 (80M параметров) по масштабу как наш gpt2: с накоплением 4 микрошагов и no_sync четыре карты дадут ускорение около 3.7 раза (оценка по модели урока): прогон на одну ночь вместо почти четырёх.
• Ступень 6. LLaMA-Factory и TRL запускают SFT и DPO через тот же torchrun (FORCE_TORCHRUN=1), а ZeRO-3 и FSDP включаются строкой конфига. Теперь вы знаете, что происходит внутри, и сможете отличить «медленно, потому что PCIe» от «медленно, потому что ошибка».

Китайский угол: DeepSeek начинала на PCIe

До H800 и DualPipe (урок 5.4.3) у DeepSeek и её материнской компании High-Flyer был кластер Fire-Flyer 2, построенный в 2021 году. Его описывает статья Fire-Flyer AI-HPC (2024): 10 000 видеокарт A100 в исполнении PCIe, а не серверных SXM-модулей с NVLink, — то есть та же связь, что у вас, только в масштабе дата-центра. По статье, кластер приблизился к производительности систем DGX-A100 при вдвое меньшей стоимости и на 40% меньшем энергопотреблении.

Ключ — своё программное обеспечение. Библиотека HFReduce заменяет all-reduce из NCCL: градиенты копируются из видеопамяти в память процессора, складываются на CPU (сначала внутри узла, затем между узлами по InfiniBand) и возвращаются на GPU. Сложение на процессоре не отнимает у видеокарт ни вычислений, ни лишних передач по PCIe между ними (детали — в статье о Fire-Flyer, ссылка в источниках урока). Поверх неё — фреймворк HaiScale с перекрытием вычислений и обмена и распределённая файловая система 3FS для данных и чекпоинтов (позже DeepSeek открыла её код).

Урок тот же, что и с DualPipe: когда связь — самое узкое место, побеждает тот, кто считает байты и прячет обмен за счётом. Ваш стенд — маленькая копия этой ситуации, и приёмы из урока (больше токенов на all-reduce, no_sync, перекрытие, BF16 для пересылок) работают и у вас.

Где учиться по-китайски. Курс ZOMI AIInfra, раздел 04Train: DDP, ZeRO, FSDP и коммуникации.

Термины. Параллелизм по данным, FSDP («полностью шардированный параллелизм по данным») и накопление градиентов по-китайски называются дословно, как в английском. Коллективные операции — «коллективная связь», а чекпоинт — «точка проверки», как английское checkpoint.

Проверь себя

Проверь себя6 вопросов·

1 Скрипт запущен командой torchrun --nproc_per_node=4. Какой должна быть строка выбора карты в каждом процессе?

torch.cuda.set_device(0) — все считают на одной карте, остальные хранят копииtorch.cuda.set_device(int(os.environ["LOCAL_RANK"]))Ничего делать не нужно: torchrun сам раскладывает процессы по картамtorch.cuda.set_device(int(os.environ["WORLD_SIZE"]))

2 Глобальный батч 64 окна: MICRO = 8, 4 карты. Каким должен быть ACCUM?

3 На 4 картах токенов/с стало в 3.6 раза больше, а кривая loss по шагам такая же, как у одной карты с тем же MICRO и ACCUM = 1. Что вероятнее всего?

Всё хорошо: эффективность 90%Все ранги получают одни и те же окна, и реальный батч не выросМешает PCIeЗабыли clip_grad_norm_

4 По модели урока: 39 TFLOPS полезного счёта, шинная скорость 10 ГБ/с. Сколько токенов на карту за один all-reduce нужно, чтобы обмен занял не больше 25% времени счёта?

5 Какие утверждения о чекпоинтах верны? Выберите все.Можно выбрать несколько ответов

При DDP достаточно сохранить model.module.state_dict() на ранге 0При FSDP каждый ранг через DCP пишет свой шард параллельноПри FSDP можно вызвать model.state_dict() только на ранге 0 и сохранитьПосле сохранения полезен dist.barrier(), чтобы никто не ушёл вперёд до конца записи

6 Почему с весами в BF16 и градиентами во FP32 FSDP гонит по PCIe столько же байт, сколько DDP во FP32?

FSDP не пересылает градиентыДва all-gather по 2 байта на параметр и reduce-scatter по 4 байта дают $(N-1)/N \cdot 8$ байт на параметр — как all-reduce по 4 байта ($2 \cdot (N-1)/N \cdot 4$)FSDP сжимает данныеЭто неверно: FSDP всегда гонит в 3 раза больше

Проверить ответы

Итоги

• torchrun запускает процесс на каждую карту и передаёт RANK, LOCAL_RANK, WORLD_SIZE; процесс выбирает карту set_device(LOCAL_RANK), NCCL строит маршруты по PCIe (у GeForce — часто через память процессора).
• Глобальный батч, общий для всех рангов и разрезанный на куски, делает обучение на 1, 2 и 4 картах одной и той же кривой loss — это проверяемо тестом. Одинаковые данные на рангах — самая тихая и дорогая ошибка.
• DDP прячет all-reduce за обратным проходом, FSDP2 шардирует веса и состояние оптимизатора; накопление с no_sync / set_requires_gradient_sync оставляет один обмен на шаг. Чекпоинты FSDP — через DCP, с продолжением и экспортом полных весов.
• Доля обмена ≈ $P \cdot \mathrm{MFU} / (W \cdot \text{токенов})$ и от размера модели не зависит: на PCIe нужны 10–20 тысяч токенов на карту за один all-reduce. Это ваш домашний вариант задачи, которую DeepSeek решала на PCIe A100 и на H800.

Дальше: ступень 6 начинается с основ обучения с подкреплением — урок MDP, отдача и функции ценности. Четыре карты пригодятся там уже для SFT и GRPO.