torchrun (эластичный запуск)

Использование

torchrun
    --standalone
    --nnodes=1
    --nproc-per-node=$NUM_TRAINERS
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

torchrun
    --rdzv-backend=c10d
    --rdzv-endpoint=localhost:0
    --nnodes=1
    --nproc-per-node=$NUM_TRAINERS
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

torchrun
    --nnodes=$NUM_NODES
    --nproc-per-node=$NUM_TRAINERS
    --max-restarts=3
    --rdzv-id=$JOB_ID
    --rdzv-backend=c10d
    --rdzv-endpoint=$HOST_NODE_ADDR
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

torchrun
    --nnodes=1:4
    --nproc-per-node=$NUM_TRAINERS
    --max-restarts=3
    --rdzv-id=$JOB_ID
    --rdzv-backend=c10d
    --rdzv-endpoint=$HOST_NODE_ADDR
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

Автодополнение командной оболочки

torchrun может создать скрипт автодополнения для bash, zsh или tcsh. Скрипт создаётся на основе анализатора аргументов, поэтому его параметры всегда соответствуют приведённым выше. Для этого требуется необязательный пакет shtab (pip install shtab); torchrun импортирует его только при использовании соответствующего флага.

# zsh
torchrun --print-completion zsh > ~/.zsh/completions/_torchrun

# bash
torchrun --print-completion bash > ~/.local/share/bash-completion/completions/torchrun

Путь к каталогу, из которого оболочка загружает скрипты автодополнения, см. в документации вашей оболочки; приведённые выше пути используются по умолчанию в большинстве случаев.

Примечание о серверной части rendezvous

Для обучения на нескольких узлах необходимо указать:

1. --rdzv-id: уникальный идентификатор задания (одинаковый для всех узлов, участвующих в задании)
2. --rdzv-backend: реализацию torch.distributed.elastic.rendezvous.RendezvousHandler
3. --rdzv-endpoint: конечную точку, где работает серверная часть rendezvous; обычно в формате host:port.

В настоящее время из коробки поддерживаются серверные части rendezvous c10d (рекомендуется), etcd-v2 и etcd (устаревшая). Чтобы использовать etcd-v2 или etcd, настройте сервер etcd с включённым API v2 (например, --enable-v2).

Определения

1. Node — физический экземпляр или контейнер; соответствует единице, с которой работает менеджер заданий.
2. Worker — рабочий процесс в контексте распределённого обучения.
3. WorkerGroup — набор рабочих процессов, выполняющих одну и ту же функцию (например, обучение).
4. LocalWorkerGroup — подмножество рабочих процессов группы, запущенных на одном узле.
5. RANK — ранг рабочего процесса в группе рабочих процессов.
6. WORLD_SIZE — общее количество рабочих процессов в группе рабочих процессов.
7. LOCAL_RANK — ранг рабочего процесса в локальной группе рабочих процессов.
8. LOCAL_WORLD_SIZE — размер локальной группы рабочих процессов.
9. rdzv_id — определяемый пользователем идентификатор, уникально задающий группу рабочих процессов для задания. Этот идентификатор используется каждым узлом для присоединения к определённой группе рабочих процессов.

9. rdzv_backend — серверная часть rendezvous (например, c10d). Обычно это строго согласованное хранилище «ключ-значение».
10. rdzv_endpoint — конечная точка серверной части rendezvous; обычно в формате <host>:<port>.

На Node запускаются LOCAL_WORLD_SIZE рабочих процессов, образующих LocalWorkerGroup. Объединение всех LocalWorkerGroups на узлах задания образует WorkerGroup.

Переменные среды

В вашем скрипте доступны следующие переменные среды:

1. LOCAL_RANK — локальный ранг.
2. RANK — глобальный ранг.
3. GROUP_RANK — ранг группы рабочих процессов. Число от 0 до max_nnodes. При запуске одной группы рабочих процессов на узел это ранг узла.
4. ROLE_RANK — ранг рабочего процесса среди всех рабочих процессов с одинаковой ролью. Роль рабочего процесса задаётся в WorkerSpec.
5. LOCAL_WORLD_SIZE — размер локального мира (например, число рабочих процессов, запущенных локально); равен значению --nproc-per-node, заданному в torchrun.
6. WORLD_SIZE — размер мира (общее количество рабочих процессов в задании).
7. ROLE_WORLD_SIZE — общее количество рабочих процессов, запущенных с ролью, заданной в WorkerSpec.
8. MASTER_ADDR — полное доменное имя (FQDN) хоста, на котором запущен рабочий процесс с рангом 0; используется для инициализации серверной части Torch Distributed.
9. MASTER_PORT — порт на MASTER_ADDR, который можно использовать для размещения хранилища C10d TCP.
10. TORCHELASTIC_RESTART_COUNT — количество перезапусков группы рабочих процессов на данный момент.
11. TORCHELASTIC_MAX_RESTARTS — заданное максимальное количество перезапусков.
12. TORCHELASTIC_RUN_ID — совпадает с run_id rendezvous (например, уникальным идентификатором задания).
13. PYTHON_EXEC — переопределение исполняемого файла системы. Если задано, пользовательский скрипт Python будет использовать значение PYTHON_EXEC в качестве исполняемого файла. По умолчанию используется sys.executable.

Ведение журнала

По умолчанию потоки stdout/stderr каждого рабочего процесса выводятся в консоль без изменений, поэтому выводы всех рангов перемешиваются и их невозможно различить.

--redirects записывает потоки в файлы журнала в каталоге --log-dir вместо вывода в консоль; --tee записывает их в файлы журнала и выводит в консоль. Для обоих параметров используется одинаковый формат: одно значение применяется ко всем рабочим процессам (3 для обоих потоков, 1 для stdout, 2 для stderr) либо задаётся сопоставление для каждого локального ранга, например 0:1,1:2. Например, --tee 3 перенаправляет в консоль оба потока для каждого рабочего процесса.

Строки, выводимые в консоль с помощью tee, получают префикс [${role_name}${local_rank}]: (например, [default3]: foobar). Чтобы изменить его, используйте --log-line-prefix-template; макросы ${role_name}, ${local_rank}, ${rank} и ${hostname} подставляются отдельно для каждого рабочего процесса. ${hostname} — имя узла, на котором запущен рабочий процесс; оно позволяет определить проблемный хост в задании на нескольких узлах:

torchrun --nnodes 2 --nproc-per-node 8 --tee 3              --log-line-prefix-template "${hostname}:${rank}: " train.py

r12i0n8:3: python: src/psm2_nccl_net.c:756: Assertion `r->used' failed.
r12i0n8:3: Fatal Python error: Segmentation fault

Шаблон также можно задать с помощью переменной среды TORCHELASTIC_LOG_LINE_PREFIX_TEMPLATE; параметр командной строки имеет более высокий приоритет. --local-ranks-filter ограничивает ранги, выводимые в консоль, не влияя на файлы журнала, создаваемые с помощью --redirects/--tee.

Развёртывание

1. (Не требуется для серверной части C10d) Запустите сервер серверной части rendezvous и получите конечную точку (её нужно передать как --rdzv-endpoint в torchrun)
2. Несколько рабочих процессов на одном узле: запустите torchrun на хосте, чтобы запустить процесс агента, который создаёт и отслеживает локальную группу рабочих процессов.
3. Несколько рабочих процессов на нескольких узлах: запустите torchrun с одинаковыми аргументами на всех узлах, участвующих в обучении.

При использовании менеджера заданий/кластера командой запуска задания на нескольких узлах должна быть torchrun.

Сценарии сбоев

1. Сбой рабочего процесса: если в задании обучения используется n рабочих процессов и сбой происходит у k<=n рабочих процессов, все рабочие процессы останавливаются и перезапускаются — не более max_restarts раз.
2. Сбой агента: сбой агента приводит к сбою локальной группы рабочих процессов. Менеджер заданий может завершить всё задание (семантика одновременного запуска группы) или попытаться заменить узел. Агент поддерживает оба варианта.
3. Сбой узла: обрабатывается так же, как сбой агента.

Изменения состава

1. Отключение узла (уменьшение масштаба): агент получает уведомление об отключении, все существующие рабочие процессы останавливаются, формируется новый WorkerGroup, после чего все рабочие процессы запускаются с новыми значениями RANK и WORLD_SIZE.
2. Подключение узла (увеличение масштаба): новый узел добавляется в задание, все существующие рабочие процессы останавливаются, формируется новый WorkerGroup, после чего все рабочие процессы запускаются с новыми значениями RANK и WORLD_SIZE.

Привязка NUMA

В многопроцессорных системах с несколькими GPU и архитектурой NUMA (неоднородный доступ к памяти) производительность можно повысить, привязав рабочие процессы к CPU рядом с назначенными им GPU. Используйте флаг --numa-binding:

torchrun --numa-binding=node --nproc-per-node=8 train.py

Подробнее см. в разделе Привязка NUMA.

Важные замечания

1. В настоящее время эта утилита и многопроцессное распределённое обучение на GPU (на одном или нескольких узлах) обеспечивают наилучшую производительность только при использовании распределённой серверной части NCCL. Поэтому для обучения на GPU рекомендуется использовать серверную часть NCCL.
2. Эта утилита предоставляет необходимые для инициализации группы процессов Torch переменные среды, поэтому передавать RANK вручную не нужно. Чтобы инициализировать группу процессов в скрипте обучения, достаточно выполнить:

>>> import torch.distributed as dist
>>> dist.init_process_group(backend="gloo|nccl")

3. В программе обучения можно использовать обычные распределённые функции или модуль torch.nn.parallel.DistributedDataParallel(). Если программа обучения использует GPU и вы хотите задействовать модуль torch.nn.parallel.DistributedDataParallel(), настройте его следующим образом.

local_rank = int(os.environ["LOCAL_RANK"])
model = torch.nn.parallel.DistributedDataParallel(
    model, device_ids=[local_rank], output_device=local_rank
)

Убедитесь, что аргумент device_ids задан как единственный идентификатор GPU, с которым будет работать ваш код. Обычно это локальный ранг процесса. Иными словами, для использования этой утилиты значение device_ids должно быть [int(os.environ("LOCAL_RANK"))], а значение output_device — int(os.environ("LOCAL_RANK")).

4. При сбоях или изменениях состава ВСЕ оставшиеся рабочие процессы немедленно завершаются. Сохраняйте контрольные точки прогресса. Частота создания контрольных точек должна зависеть от того, какой объём выполненной работы допустимо потерять.
5. Этот модуль поддерживает только однородные LOCAL_WORLD_SIZE. Предполагается, что на всех узлах запускается одинаковое количество локальных рабочих процессов (для каждой роли).
6. RANK НЕ является стабильным. После перезапусков локальным рабочим процессам на узле может быть назначен диапазон рангов, отличающийся от прежнего. НИКОГДА не закладывайте в код предположения о стабильности рангов или о какой-либо взаимосвязи между RANK и LOCAL_RANK.
7. При использовании эластичного режима (min_size!=max_size) НЕ закладывайте в код предположения о WORLD_SIZE, поскольку размер мира может меняться при подключении и отключении узлов.
8. Рекомендуется организовать скрипт следующим образом:

def main():
    load_checkpoint(checkpoint_path)
    initialize()
    train()


def train():
    for batch in iter(dataset):
        train_step(batch)

        if should_checkpoint:
            save_checkpoint(checkpoint_path)

9. (Рекомендуется) При ошибках рабочих процессов эта утилита суммирует сведения об ошибке (например, время, ранг, хост, PID, трассировку стека и т. д.). На каждом узле первая по времени ошибка эвристически указывается как ошибка «первопричины». Чтобы включить трассировки стека в сводку, декорируйте главную функцию точки входа скрипта обучения, как показано в примере ниже. Если декоратор не используется, сводка не будет содержать трассировку стека исключения и будет включать только код завершения. Подробнее об обработке ошибок в torchelastic см. здесь: https://pytorch.org/docs/stable/elastic/errors.html

from torch.distributed.elastic.multiprocessing.errors import record


@record
def main():
    # do train
    pass


if __name__ == "__main__":
    main()