torch.distributed.fsdp.fully_shard

Группировка и планирование коммуникаций

Каждый вызов fully_shard создаёт одну группу коммуникаций, содержащую все параметры модуля, которые ещё не были назначены группе в результате более раннего вызова для подмодуля. Параметры каждой группы совместно собираются с помощью одной коллективной операции перед прямым проходом, а их градиенты совместно распределяются с помощью одной коллективной операции после обратного прохода. В отличие от DDP, в FSDP2 нет параметра bucket_cap_mb — границы коммуникаций полностью определяются тем, к каким модулям вы применяете fully_shard.

Рассмотрим модель с четырьмя подмодулями, где a, b, c и d обозначают количество параметров в каждом из них:

model[ m1[a] -> m2[b] -> m3[c] -> m4[d] ]

Если вызвать только fully_shard(model) (только для корневого модуля), все параметры будут входить в одну группу. Это означает, что весь прямой и обратный проход будут выглядеть так:

all-gather(a+b+c+d) -> forward(m1,m2,m3,m4) -> backward(m4,m3,m2,m1) -> reduce-scatter(a+b+c+d)

Вся коммуникация выполняется в виде двух крупных блокирующих операций без перекрытия с вычислениями. Так делать почти никогда не следует.

Если применять fully_shard для каждого подмодуля — например, вызвать fully_shard(m2), fully_shard(m3), а затем fully_shard(model), — оставшиеся параметры (a и d) образуют корневую группу, а m2 и m3 получат собственные группы.

Во время прямого прохода операции all-gather выполняются в отдельном потоке CUDA, поэтому all-gather следующего модуля может перекрываться с вычислениями прямого прохода текущего модуля. Хук pre-forward каждого модуля запускает собственный all-gather и ожидает его завершения перед запуском модуля. Поскольку процессор обычно опережает графический процессор, all-gather следующего модуля запускается в потоке AG, пока прямой проход текущего модуля ещё выполняется в потоке вычислений:

              time ──────────────────────────────────────────────►

compute:      [wait] [ fwd(m1)   | fwd(m2)    | fwd(m3,m4)     ]
AG stream:    [AG(a,d)]  [AG(b)  |    AG(c)   ]

Пока fwd(m1) выполняется в потоке вычислений, процессор запускает хук pre-forward для m2, который запускает AG(b) в потоке AG. Чтобы сделать такое перекрытие надёжнее (например, если накладные расходы на стороне процессора сокращают временной запас), используйте set_modules_to_forward_prefetch, чтобы запустить следующий all-gather раньше — в хуке pre-forward текущего модуля, не дожидаясь срабатывания хука следующего модуля.

Во время обратного прохода FSDP2 также явно предварительно загружает all-gather следующего модуля и запускает reduce-scatter в отдельном потоке CUDA — и всё это без дополнительной настройки:

              time ──────────────────────────────────────────────►

compute:      [ bwd(m4,m3)     | bwd(m2)        | bwd(m1)       ]
AG stream:    [AG(c)] [ AG(b)  |   AG(a,d)      ]
RS stream:                     |[RS(c)]  [ RS(b)|     RS(a,d)   ]

Пока bwd(m4,m3) выполняется в потоке вычислений, all-gather для b (необходимый для m2) предварительно загружается в потоке AG. Пока выполняется bwd(m2), и AG(a,d), и RS(c) перекрываются с вычислениями. Именно благодаря такому конвейерному выполнению рекомендуется применять fully_shard снизу вверх к каждому слою, прежде чем применять его к корневому модулю.

Чтобы управлять размером каждой группы коммуникаций, выберите модули для обёртывания: обёртывание более мелких модулей приводит к созданию меньших групп, которые легче перекрывать с вычислениями (аналогично уменьшению корзин DDP), а обёртывание меньшего числа модулей — к созданию более крупных групп. Автоматической разбивки на корзины нет: группировка задаётся явно и определяется структурой модулей.

По сравнению с PyTorch FSDP1 (FullyShardedDataParallel):

• FSDP2 использует шардирование каждого параметра по dim-0 на основе DTensor, что обеспечивает более простое представление шардирования по сравнению с шардированием плоских параметров в FSDP1 при сохранении сопоставимой производительности. Точнее, FSDP2 разбивает каждый параметр по dim-0 между рабочими процессами параллелизма данных (используя torch.chunk(dim=0)), тогда как FSDP1 выравнивает, объединяет и разбивает на части группу тензоров, из-за чего сложнее понять, какие данные есть на каждом рабочем процессе, и изменить шардирование для других видов параллелизма. Шардирование каждого параметра делает работу интуитивно понятнее, ослабляет ограничения для замороженных параметров и позволяет создавать словари состояний без коммуникаций (шардированные), для которых в FSDP1 в противном случае требуются операции all-gather.
• В FSDP2 реализован другой подход к управлению памятью для работы с несколькими потоками, позволяющий избежать torch.Tensor.record_stream. Это обеспечивает детерминированное и ожидаемое использование памяти и не требует блокировки процессора, как в limit_all_gathers=True в FSDP1.
• FSDP2 предоставляет API для ручного управления предварительной загрузкой и планированием коллективных операций, давая опытным пользователям больше возможностей для настройки. Подробности см. в описании методов FSDPModule ниже.
• В FSDP2 упрощена часть API: например, FSDP2 не поддерживает полные словари состояний напрямую. Вместо этого пользователи могут самостоятельно преобразовать шардированные словари состояний, содержащие DTensor, в полные словари состояний с помощью API DTensor, например DTensor.full_tensor(), или API распределённых словарей состояний более высокого уровня из PyTorch Distributed Checkpoint. Также удалены некоторые другие аргументы; подробности см. здесь.

Пользовательский API — это fully_shard, который можно вызвать для module:

torch.distributed.fsdp.fully_shard(module, *, mesh=None, reshard_after_forward=None, shard_placement_fn=None, mp_policy=MixedPrecisionPolicy(param_dtype=None, reduce_dtype=None, output_dtype=None, cast_forward_inputs=True), offload_policy=OffloadPolicy(), ignored_params=None, dp_mesh_dims=None) [исходный код]

Применяет полностью шардированный параллелизм данных (FSDP) к module: FSDP шардирует параметры, градиенты и состояния оптимизатора модуля между рабочими процессами параллелизма данных, экономя память за счёт дополнительных коммуникаций.

При инициализации FSDP шардирует параметры модуля между рабочими процессами параллелизма данных, заданными mesh. Перед прямым проходом FSDP собирает шардированные параметры между рабочими процессами параллелизма данных, чтобы получить нешардированные параметры для вычислений прямого прохода. Если reshard_after_forward равно True, FSDP освобождает нешардированные параметры после прямого прохода и повторно собирает их с помощью all-gather в обратном проходе перед вычислением градиентов. После вычисления градиентов FSDP освобождает нешардированные параметры и распределяет нешардированные градиенты между рабочими процессами параллелизма данных с помощью reduce-scatter.

В этой реализации шардированные параметры представлены как DTensor, шардированные по dim-0, а нешардированные параметры имеют тот же вид, что и исходные параметры на module (например, torch.Tensor, если изначально это был torch.Tensor). Хук pre-hook прямого прохода модуля на module собирает параметры с помощью all-gather, а хук прямого прохода модуля на module освобождает их (если необходимо). Аналогичные хуки обратного прохода собирают параметры с помощью all-gather, а затем освобождают параметры и распределяют градиенты с помощью reduce-scatter.

Поскольку объединение нескольких тензоров для одной коллективной операции критически важно для эффективности коммуникаций, в этой реализации такая группировка является одной из основных возможностей. Вызов fully_shard() для module создаёт одну группу, включающую параметры module.parameters(), кроме тех, которые уже были назначены группе в результате более раннего вызова для подмодуля. Поэтому fully_shard() следует вызывать для модели снизу вверх. Параметры каждой группы собираются с помощью одной коллективной операции all-gather, а градиенты распределяются с помощью одной коллективной операции reduce-scatter. Разбиение модели на несколько групп («по слоям») позволяет максимально сократить использование памяти и перекрывать коммуникации с вычислениями. Как правило, пользователям не следует вызывать fully_shard() только для самого верхнего корневого модуля.

При вызове со списком (fully_shard([a, b, ...])) прямой проход модели может выполняться только для подмножества сгруппированных модулей, а остальные модули вызываться позднее в той же итерации. Типичный пример — обучение с разбиением функции потерь на фрагменты с помощью fully_shard([norm, head]): основной прямой проход выполняет только norm, а затем head вызывается для каждого фрагмента. Важные замечания:

• Каждый отдельный вызов для фрагмента регистрирует собственный узел autograd post_backward, поэтому N вызовов для фрагментов приводят к N операциям reduce-scatter для этой группы.
• mp_policy.cast_forward_inputs и mp_policy.output_dtype применяются к каждому модулю в группе: каждый вызов (включая отдельные вызовы для каждого фрагмента) преобразует входные данные к param_dtype, а выходные — к output_dtype.

Примечание

Если возникает исключение в forward() или backward(), состояние FSDP для текущей итерации (маркер корневого модуля прямого прохода, отслеживание запусков сгруппированных модулей, состояние выполняющихся коллективных операций, состояния обучения для каждой группы) остаётся неопределённым. Чтобы восстановить состояние и запустить следующую итерацию, вызовите FSDPModule.reset_iter_state() для корневого модуля FSDP. Градиенты неудачной итерации отбрасываются, включая состояние накопления частичных reduce-операций no_sync / HSDP.

Параметры:

• module (Union[nn.Module, List[nn.Module]]) – Модуль или модули, параметры которых нужно шардировать с помощью FSDP и объединить в группы для коммуникаций.
• mesh (Optional[DeviceMesh]) – Эта сетка параллелизма данных задаёт шардирование и устройство. Если сетка одномерная, параметры полностью шардируются по одномерной сетке (FSDP) с размещением (Shard(0),). Если сетка двумерная, параметры шардируются по первому измерению и реплицируются по нулевому измерению (HSDP) с размещением (Replicate(), Shard(0)). Тип устройства сетки задаёт тип устройства, используемого для коммуникаций; для CUDA или совместимых с CUDA устройств используется текущее устройство.

• reshard_after_forward (Optional[Union[bool, int]]) –

  Этот параметр управляет поведением параметров после прямого прохода и позволяет выбирать компромисс между использованием памяти и объёмом коммуникаций:

  • Если True, параметры повторно шардируются после прямого прохода и снова собираются с помощью all-gather в обратном проходе.
  • Если False, нешардированные параметры остаются в памяти после прямого прохода, что позволяет избежать all-gather в обратном проходе. Для максимальной производительности обычно задают False для корневого модуля, поскольку он, как правило, требуется сразу после начала обратного прохода.
  • Если None, для некорневых модулей устанавливается значение True, а для корневых — False.
  • Если задано целое число int, оно обозначает размер мира, до которого нужно повторно шардировать параметры после прямого прохода. Оно должно быть нетривиальным делителем размера измерения шардирования mesh (то есть не равняться 1 или самому размеру измерения). Например, можно выбрать размер внутриузловой группы (например, torch.cuda.device_count()). Это позволяет выполнять all-gather в обратном проходе для меньшего размера мира, но требует больше памяти, чем при значении True.
  • После прямого прохода зарегистрированные в модуле параметры зависят от этого значения: это шардированные параметры при True; нешардированные параметры при False; в противном случае — параметры, повторно шардированные по меньшей сетке. Чтобы изменить параметры между прямым и обратным проходами, зарегистрированные параметры должны быть шардированными. Для False или целого числа int это можно сделать, выполнив повторное шардирование вручную с помощью reshard().

• shard_placement_fn (Optional[Callable[[nn.Parameter], Optional[Shard | ShardPlacementResult]]]) –

  Эта вызываемая функция позволяет переопределить размещение шардов и/или сетку для параметра. Она может возвращать:

  • None: использовать шардирование по умолчанию (Shard(0)) в сетке, переданной в fully_shard.
  • Shard: шардировать параметр по указанному измерению, используя сетку, переданную в fully_shard.
  • ShardPlacementResult: задать и размещение шардов, и пользовательскую FSDPMeshInfo. Это позволяет шардировать разные параметры по разным группам процессов и поддерживает такие сценарии, как Mixture of Experts, где для параметров экспертов используется другая сетка, чем для обычных параметров.

  При шардировании по ненулевому измерению в настоящее время требуется равномерное шардирование: размер измерения тензора должен делиться на размер сетки шардирования FSDP без остатка.

• mp_policy (MixedPrecisionPolicy) – Управляет политикой смешанной точности, которая задаёт смешанную точность параметров и редукции для этого модуля. Подробности см. в разделе MixedPrecisionPolicy.
• offload_policy (OffloadPolicy) – Управляет политикой выгрузки, которая позволяет выгружать параметры, градиенты и состояния оптимизатора. Подробности см. в разделе OffloadPolicy и описаниях его подклассов.
• ignored_params (set[nn.Parameter] | None) – Необязательный набор параметров (Set[nn.Parameter]), которые FSDP должен игнорировать. Они не будут шардироваться, перемещаться на устройство при инициализации или участвовать в редукции градиентов в обратном проходе.
• dp_mesh_dims (Optional[DataParallelMeshDims]) – Если задан, mesh рассматривается как полная SPMD-сетка, а параметры должны быть DTensor в этой сетке с Replicate() по всем измерениям параллелизма данных. Поле shard указывает измерения, по которым FSDP выполняет шардирование (несколько измерений объединяются). Поле replicate указывает измерения репликации HSDP (несколько измерений объединяются).

Возвращает:

Модуль, к которому применён FSDP (на месте).

Тип возвращаемого значения:

FSDPModule

class torch.distributed.fsdp.FSDPModule(*args, **kwargs)

reset_iter_state() [источник]

Сбрасывает состояние FSDP для текущей итерации после того, как исключение прервало выполнение прямого или обратного прохода. Поддерживаемый порядок восстановления:

1. Перехватите исключение из forward() или backward().
2. Вызовите reset_iter_state() для корневого модуля FSDP.
3. Выполните следующую итерацию обычным образом.

Сброс ожидает завершения всех выполняющихся операций all-gather/reduce-scatter, повторно разделяет на шарды каждую группу параметров и очищает отслеживание итерации (iter_forward_root, _modules_to_run_forward, порядок после прямого прохода, состояния обучения для каждой группы). Все выполняющиеся операции редукции градиентов отменяются: градиенты неудачной итерации теряются, включая состояние частичного накопления редукции HSDP и состояние накопления градиентов no_sync. При накоплении градиентов последовательность микропакетов следует считать недействительной и начать заново.

Метод необходимо вызывать для корневого модуля FSDP — то есть модуля, к которому применён fully_shard верхнего уровня; другими словами, модуля, для которого первым был выполнен прямой проход. Вызов для некорневого модуля приводит к RuntimeError.

reshard() [источник]

Повторно разделяет параметры модуля на шарды, освобождая неразделённые параметры, если они выделены, и регистрируя в модуле параметры, разделённые на шарды. Этот метод не является рекурсивным.

set_all_reduce_hook(hook, *, stream=None) [источник]

Параметры:

• hook (Callable[[torch.Tensor], None]) – Пользовательский хук all-reduce с ожидаемой сигнатурой hook(reduce_output: torch.Tensor) -> None, где reduce_output — результат reduce-scatter, если используется только FSDP, или результат all-reduce, если используется встроенный HSDP.
• stream (Optional[torch.cuda.Stream]) – Поток для выполнения хука all-reduce. Указывать его следует только при использовании не встроенного HSDP. При использовании встроенного HSDP хук выполняется во внутреннем потоке all-reduce, используемом встроенным HSDP.

set_allocate_memory_from_process_group_for_comm(enable) [источник]

Указывает, следует ли выделять временные буферы staging, используемые для отправки и получения данных при коллективном обмене данными, с помощью предоставленного самой ProcessGroup оптимизированного пользовательского аллокатора (если он есть). Это может повысить эффективность ProcessGroup. Например, при использовании NCCL это позволяет задействовать передачи без копирования через SHARP (для NVLink и/или InfiniBand).

Этот метод нельзя использовать совместно с set_custom_all_gather() или set_custom_reduce_scatter(), поскольку эти API позволяют задавать более детальные параметры для каждого обмена данными, а этот метод не может определить стратегию выделения буферов staging для них.

Параметры:

enable (bool) – Включать ли выделение памяти ProcessGroup.

set_custom_all_gather(comm) [источник]

Переопределяет стандартное поведение обмена данными all_gather, чтобы обеспечить более точный контроль обмена данными и использования памяти. Подробности см. в Comm и ReduceScatter.

Параметры:

comm (AllGather) – Пользовательская операция all-gather.

set_custom_reduce_scatter(comm) [источник]

Переопределяет стандартное поведение обмена данными reduce_scatter, чтобы обеспечить более точный контроль обмена данными и использования памяти. Подробности см. в Comm и ReduceScatter.

Параметры:

comm (ReduceScatter) – Пользовательская операция reduce_scatter.

set_force_sum_reduction_for_comms(enable) [источник]

Указывает, следует ли требовать, чтобы низкоуровневые примитивы коллективного обмена данными использовали исключительно редукции типа «sum», даже если для этого придётся выполнять дополнительные операции предварительного или последующего масштабирования. Это необходимо, например, потому что в настоящее время NCCL поддерживает передачи без копирования только для такого типа коллективных операций.

Примечание: для устройств MTIA этот режим всегда включён неявно.

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

Параметры:

enable (bool) – Использовать ли для обмена данными только ReduceOp.SUM.

set_gradient_divide_factor(factor) [источник]

Задаёт пользовательский коэффициент деления для редукции градиентов. Для этого может использоваться пользовательская операция редукции с PreMulSum в NCCL, которая позволяет умножить данные на коэффициент перед редукцией.

Параметры:

factor (float) – Пользовательский коэффициент деления.

set_is_last_backward(is_last_backward) [источник]

Указывает, будет ли следующий обратный проход последним. Во время последнего обратного прохода FSDP ожидает завершения ожидающей редукции градиентов и очищает внутренние структуры данных для предварительной выборки в обратном проходе. Это может быть полезно при обработке микропакетов.

set_modules_to_backward_prefetch(modules) [источник]

Задаёт модули FSDP, для которых этот модуль FSDP должен явно выполнять предварительную выборку операций all-gather в обратном проходе. Это переопределяет стандартную реализацию предварительной выборки в обратном проходе, которая выбирает следующий модуль FSDP на основе обратного порядка после прямого прохода.

Передача списка из одного элемента с предыдущим модулем FSDP обеспечивает такое же перекрытие all-gather, как и стандартное поведение перекрытия. Для более интенсивного перекрытия необходимо передать список длиной не менее двух элементов; при этом будет использоваться больше зарезервированной памяти.

Параметры:

modules (List[FSDPModule]) – Модули FSDP для предварительной выборки.

set_modules_to_forward_prefetch(modules) [источник]

Задаёт модули FSDP, для которых этот модуль FSDP должен явно выполнять предварительную выборку операций all-gather в прямом проходе. Предварительная выборка выполняется после копирования результатов all-gather этого модуля.

Передача списка из одного элемента со следующим модулем FSDP обеспечивает такое же перекрытие all-gather, как и стандартное поведение перекрытия, однако предварительно выбранная операция all-gather отправляется с CPU раньше. Для более интенсивного перекрытия необходимо передать список длиной не менее двух элементов; при этом будет использоваться больше зарезервированной памяти.

Параметры:

modules (List[FSDPModule]) – Модули FSDP для предварительной выборки.

set_post_optim_event(event) [источник]

Задаёт событие после шага оптимизатора, которого корневой модуль FSDP будет ожидать в потоках all-gather.

По умолчанию корневой модуль FSDP ожидает в текущем потоке завершения операций all-gather, чтобы шаг оптимизатора завершился до начала all-gather. Однако это может приводить к ложным зависимостям, если после шага оптимизатора выполняются несвязанные вычисления. Этот API позволяет пользователю задать собственное событие для ожидания. После того как корневой модуль ожидает событие, оно удаляется, поэтому при каждом шаге этот API следует вызывать с новым событием.

Параметры:

event (torch.Event) – Событие, зарегистрированное после шага оптимизатора, которого будут ожидать потоки all-gather.

set_reduce_scatter_divide_factor(factor) [источник]

Вместо этого используйте set_gradient_divide_factor()

set_reduce_scatter_max_input_buffers(max_input_buffers, *, recurse=True) [источник]

Задаёт максимальное количество входных буферов градиентов reduce-scatter, которые могут одновременно находиться в обработке, — лимит буферов копирования входных данных (chunk_cat), cap-K (экспериментальная возможность).

По умолчанию FSDP одновременно обрабатывает 1 такой буфер, поэтому поток вычислений должен ожидать завершения предыдущего reduce-scatter, прежде чем следующий этап копирования входных данных сможет повторно использовать этот буфер. Когда операция reduce-scatter становится узким местом (обмен данными идёт медленнее, чем скрывающие его вычисления обратного прохода), ожидание повторного использования буфера приостанавливает поток вычислений. Увеличение лимита позволяет следующему копированию входных данных записывать данные в новый буфер, не дожидаясь завершения, что устраняет остановку, но требует дополнительной пиковой памяти для хранения буферов. Копирование входных данных остаётся в потоке вычислений; дополнительный поток и record_stream не используются. Это помогает только тогда, когда reduce-scatter становится узким местом.

Параметры:

• max_input_buffers (int) – Максимальное количество входных буферов reduce-scatter, сохраняемых в обработке (параметр баланса между памятью и перекрытием); должно быть >= 1. 1 соответствует поведению FSDP по умолчанию (один буфер; остановка при узком месте reduce-scatter). Небольшое значение (например, 2) ограничивает пиковое использование памяти и не вызывает остановок, пока оно >= глубины конвейера reduce-scatter (иначе узкое место reduce-scatter снова приведёт к остановке в конце); большее значение сохраняет больше буферов, обеспечивая более глубокое перекрытие ценой увеличения пиковой памяти.
• recurse (bool) – Задавать ли значение для всех вложенных модулей FSDP или только для переданного модуля.

set_reduce_scatter_unused_params(reduce_scatter_unused_params, *, recurse=True) [источник]

Указывает, следует ли включать нулевые градиенты для параметров, не получивших градиент в обратном проходе. Это необходимо, когда на разных рангах используются разные параметры из-за условного управления потоком (например, в мультимодальных моделях или моделях со смесью экспертов), что приводит к несовпадению коллективных операций reduce-scatter. Аналогично find_unused_parameters в DDP.

Параметры:

• reduce_scatter_unused_params (bool) – Включать ли нулевые градиенты для неиспользуемых параметров в редукцию градиентов.
• recurse (bool) – Задавать ли значение для всех вложенных модулей FSDP или только для переданного модуля.

set_requires_all_reduce(requires_all_reduce, *, recurse=True) [источник]

Указывает, следует ли выполнять all-reduce градиентов для модуля. Это можно использовать для реализации накопления градиентов с выполнением только reduce-scatter, но без all-reduce, в HSDP.

set_requires_gradient_sync(requires_gradient_sync, *, recurse=True) [источник]

Указывает, следует ли синхронизировать градиенты модуля. Это можно использовать для реализации накопления градиентов без обмена данными. Для HSDP этот параметр одновременно управляет reduce-scatter и all-reduce. Эквивалент no_sync в FSDP1.

Параметры:

• requires_gradient_sync (bool) – Выполнять ли редукцию градиентов для параметров модуля.
• recurse (bool) – Задавать ли значение для всех вложенных модулей FSDP или только для переданного модуля.

set_reshard_after_backward(reshard_after_backward, *, recurse=True) [источник]

Указывает, следует ли повторно разделять параметры модуля на шарды после обратного прохода. Это можно использовать при накоплении градиентов, обменивая больший расход памяти на сокращение обмена данными, поскольку перед следующим прямым проходом не потребуется повторно выполнять all-gather неразделённых параметров.

Параметры:

• reshard_after_backward (bool) – Повторно разделять ли параметры на шарды после обратного прохода.
• recurse (bool) – Задавать ли значение для всех вложенных модулей FSDP или только для переданного модуля.

set_reshard_after_forward(reshard_after_forward, recurse=True) [источник]

Указывает, следует ли повторно разделять параметры модуля на шарды после прямого прохода. Это позволяет во время выполнения изменить аргумент FSDP reshard_after_forward. Например, можно задать для корневого модуля FSDP значение True (поскольку иначе для него специально устанавливается False) или установить для модуля FSDP значение False при запуске оценки, а затем вернуть значение True для обучения.

Параметры:

• reshard_after_forward (bool) – Повторно разделять ли параметры на шарды после прямого прохода.
• recurse (bool) – Задавать ли значение для всех вложенных модулей FSDP или только для переданного модуля.

set_separate_reduce_scatter_group(enable=True, *, recurse=True) [источник]

Включает (или отключает) выполнение reduce-scatter градиентов в отдельной группе процессов, чтобы эта операция могла перекрываться с all-gather в обратном проходе (экспериментальная возможность).

По умолчанию FSDP выполняет all-gather и reduce-scatter в отдельных потоках CUDA, но через одну и ту же группу процессов — один коммуникатор NCCL, который обрабатывает по одной коллективной операции за раз и поэтому сериализует их при передаче по сети. Если включить этот параметр, FSDP создаст отдельную группу процессов для рангов шардов (dist.new_group(..., use_local_synchronization=True)) — по одной на каждое уникальное множество рангов шардов, обычно с одним коммуникатором, — чтобы обе коллективные операции могли выполняться параллельно, если пропускная способность сети это позволяет. Настройка является коллективной для каждого множества рангов шардов: как и при настройке других операций обмена данными FSDP, вызывайте её согласованно на рангах, использующих эту сетку FSDP.

Параметры:

• enable (bool) – True (по умолчанию) выделяет reduce-scatter отдельную группу процессов; False возвращает общую группу процессов для шардов/all-gather.
• recurse (bool) – Задавать ли значение для всех вложенных модулей FSDP или только для переданного модуля.

set_symm_mem_for_comm(backend='NCCL') [источник]

Задаёт бэкенд симметричной памяти (symm_mem) для выделения буферов staging, используемых в коллективных операциях all-gather. Это позволяет NCCL использовать оптимизированные реализации all-gather с помощью симметричной памяти. Такая оптимизация может зависеть от топологии системы. Для одного узла может использоваться Copy Engine All-Gather. Для нескольких узлов может использоваться Symmetric Kernel All-Gather.

Чтобы включить Copy Engine All-Gather, необходимо задать для группы процессов NCCL политику zero-CTA. `python opts = dist.ProcessGroupNCCL.Options() opts.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO dist.init_process_group(backend="nccl", pg_options=opts, device_id=device) ` В качестве альтернативы можно установить значение 2 для переменной среды NCCL_CTA_POLICY. `bash export NCCL_CTA_POLICY=2 ` Дополнительные сведения см. в разделе [Коллективные операции Copy Engine](https://docs.pytorch.org/docs/2.11/symmetric_memory.html#copy-engine-collectives).

Этот метод нельзя использовать совместно с set_custom_all_gather() или set_custom_reduce_scatter().

Параметры:

backend (str) – Используемый бэкенд симметричной памяти. По умолчанию используется "NCCL". В настоящее время поддерживается только "NCCL".

set_unshard_in_backward(unshard_in_backward) [источник]

Указывает, нужно ли восстанавливать параметры модуля FSDP из шардов в обратном проходе. Этот параметр можно использовать в особых случаях, когда пользователь знает, что параметры в группе параметров этого модуля FSDP не нужны для вычислений обратного прохода (например, для embedding).

unshard(async_op=False) [источник]

Восстанавливает параметры модуля из шардов, выделяя память и выполняя all-gather параметров. Этот метод не является рекурсивным. Восстановление из шардов выполняется в соответствии с MixedPrecisionPolicy, поэтому, если задан параметр param_dtype, операция all-gather будет выполнена с его учётом.

Параметры:

async_op (bool) – Если задано True, возвращает UnshardHandle с методом wait() для ожидания завершения операции восстановления из шардов. Если задано False, возвращает None и ожидает завершения операции через дескриптор внутри этой функции.

Тип возвращаемого значения:

UnshardHandle | None

Примечание

Если задано async_op=True, FSDP самостоятельно дождётся завершения ожидающей операции восстановления из шардов в предварительном прямом проходе модуля. Пользователю нужно явно вызвать wait(), только если ожидание должно произойти до предварительного прямого прохода.

class torch.distributed.fsdp.UnshardHandle

Дескриптор для ожидания завершения операции FSDPModule.unshard().

wait() [источник]

Ожидает завершения операции восстановления из шардов. Это гарантирует, что текущий поток сможет использовать восстановленные параметры, которые теперь зарегистрированы в модуле.

torch.distributed.fsdp.register_fsdp_forward_method(module, method_name) [источник]

Регистрирует метод module как метод прямого прохода для FSDP.

FSDP выполняет all-gather параметров перед прямым проходом и при необходимости освобождает параметры после него (в зависимости от reshard_after_forward). По умолчанию FSDP знает, как делать это только для nn.Module.forward(). Эта функция изменяет указанный пользователем метод так, чтобы перед ним и после него выполнялись соответствующие хуки предварительного и последующего прямого прохода. Если module не является экземпляром FSDPModule, функция ничего не делает.

Параметры:

• module (nn.Module) – Модуль, для которого регистрируется метод прямого прохода.
• method_name (str) – Имя метода прямого прохода.

class torch.distributed.fsdp.MixedPrecisionPolicy(param_dtype=None, reduce_dtype=None, output_dtype=None, cast_forward_inputs=True)

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

FSDP хорошо работает со смешанной точностью на уровне модулей, поскольку высокоточные шардированные параметры и так остаются в памяти. Другими словами, FSDP не требует дополнительной памяти для хранения высокоточной копии параметров на шаге оптимизатора.

Переменные:

• param_dtype (Optional[torch.dtype]) – Задаёт тип данных нешардированного параметра, а следовательно, тип данных для вычислений прямого и обратного проходов и для операции all-gather параметров. Если задано None, нешардированный параметр использует исходный тип данных. На шаге оптимизатора используется шардированный параметр с исходным типом данных. (По умолчанию: None)
• reduce_dtype (Optional[torch.dtype]) – Задаёт тип данных для редукции градиентов (то есть reduce-scatter или all-reduce). Если задано None, но param_dtype не равно None, для редукции используется тип данных вычислений. Это позволяет выполнять редукцию градиентов с полной точностью, используя низкую точность для вычислений. Если редукция градиентов также отключена с помощью set_requires_gradient_sync(), FSDP будет накапливать градиенты с использованием reduce_dtype. (По умолчанию: None)
• output_dtype (Optional[torch.dtype]) – Задаёт тип данных для преобразования выходных значений прямого прохода с плавающей точкой. Это можно использовать для реализации случаев, когда для разных модулей заданы разные политики смешанной точности. (По умолчанию: None)
• cast_forward_inputs (bool) – Указывает, следует ли FSDP преобразовывать входные тензоры прямого прохода с плавающей точкой к типу param_dtype. Для сгруппированных fully_shard([a, b, ...]) преобразование выполняется для каждого модуля перед его прямым проходом.

class torch.distributed.fsdp.OffloadPolicy

Этот базовый класс представляет политику без выгрузки и используется только в качестве значения по умолчанию для аргумента offload_policy.

class torch.distributed.fsdp.CPUOffloadPolicy(pin_memory=True)

Эта политика выгружает параметры, градиенты и состояния оптимизатора на CPU. Перед операцией all-gather шардированные параметры копируются с хоста на устройство. Параметры, собранные операцией all-gather, освобождаются в соответствии с reshard_after_forward. Во время обратного прохода шардированные градиенты копируются с устройства на хост, а шаг оптимизатора выполняется на CPU с состояниями оптимизатора на CPU.

Переменные:

pin_memory (bool) – Определяет, следует ли закреплять в памяти хоста память для шардированных параметров и градиентов. Закрепление памяти повышает эффективность копирования H2D/D2H и позволяет выполнять копирование одновременно с вычислениями. Однако закреплённая память недоступна другим процессам. Установите значение False, если памяти CPU недостаточно. (По умолчанию: True)

torch.distributed.fsdp.share_comm_ctx(modules) [исходный код]

Совместное использование потоков CUDA несколькими FSDPModules

Пример использования:

from torch.distributed.fsdp import share_comm_ctx share_comm_ctx([fsdp_model_1, fsdp_model_2, …])

При параллелизме конвейера (PP) каждый фрагмент модели является корневым элементом FSDP. Мы хотим совместно использовать потоки CUDA для операций all-gather, reduce-scatter и all-reduce. Это позволяет избежать фрагментации памяти между потоками.

Параметры:

modules (List[FSDPModule]) – модули, для которых нужно совместно использовать потоки CUDA

class torch.distributed.fsdp.DataParallelMeshDims(shard=None, replicate=None)

Указывает, какие измерения полной сетки SPMD DeviceMesh соответствуют параллелизму данных при использовании fully_shard(), если параметры уже являются DTensor в этой сетке.

Переменные:

• shard (Optional[Union[str, tuple[str, ...]]]) – Имя или имена измерений сетки, по которым FSDP разделяет параметры на шарды. Если указана кортежем несколько имён, соответствующие измерения объединяются в одно измерение шардинга. Должен быть задан как минимум один из параметров shard и replicate.
• replicate (Optional[Union[str, tuple[str, ...]]]) – Имя или имена измерений сетки для репликации HSDP или DDP. Если указана кортежем несколько имён, соответствующие измерения объединяются в одно измерение репликации.