Spec-Zone.ru › PyTorch 2.14

FullyShardedDataParallel

Создано: 2 февраля 2022 г. | Последнее обновление: 12 мая 2026 г.

class torch.distributed.fsdp.FullyShardedDataParallel(module, process_group=None, sharding_strategy=None, cpu_offload=None, auto_wrap_policy=None, backward_prefetch=BackwardPrefetch.BACKWARD_PRE, mixed_precision=None, ignored_modules=None, param_init_fn=None, device_id=None, sync_module_states=False, forward_prefetch=False, limit_all_gathers=True, use_orig_params=False, ignored_states=None, device_mesh=None) [исходный код]

Обёртка для распределения параметров модуля между рабочими процессами параллелизма по данным.

Эта реализация вдохновлена работой Xu и др., а также подходом ZeRO Stage 3 из DeepSpeed. FullyShardedDataParallel обычно сокращённо называют FSDP.

Пример:

>>> import torch
>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> torch.cuda.set_device(device_id)
>>> sharded_module = FSDP(my_module)
>>> optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
>>> x = sharded_module(x, y=3, z=torch.Tensor([1]))
>>> loss = x.sum()
>>> loss.backward()
>>> optim.step()

Для использования FSDP необходимо обернуть модуль, а затем инициализировать оптимизатор. Это необходимо, поскольку FSDP изменяет переменные параметров.

При настройке FSDP необходимо учитывать целевое устройство CUDA. Если у устройства есть идентификатор (dev_id), доступны три варианта:

  • Разместить модуль на этом устройстве
  • Задать устройство с помощью torch.cuda.set_device(dev_id)
  • Передать dev_id в аргумент конструктора device_id.

Это гарантирует, что вычислительным устройством экземпляра FSDP будет целевое устройство. В вариантах 1 и 3 инициализация FSDP всегда выполняется на GPU. В варианте 2 инициализация FSDP происходит на текущем устройстве модуля, которым может быть CPU.

Если используется флаг sync_module_states=True, необходимо убедиться, что модуль находится на GPU, либо использовать аргумент device_id, чтобы указать устройство CUDA, на которое FSDP переместит модуль в конструкторе FSDP. Это необходимо, поскольку sync_module_states=True требует обмена данными через GPU.

FSDP также перемещает входные тензоры, передаваемые методу forward, на вычислительное устройство GPU, поэтому вручную переносить их с CPU не нужно.

Для use_orig_params=True, ShardingStrategy.SHARD_GRAD_OP предоставляет нешардированные параметры после forward, а не шардированные параметры, в отличие от ShardingStrategy.FULL_SHARD. Чтобы проверить градиенты, можно использовать метод summon_full_params с with_grads=True.

При использовании limit_all_gathers=True перед выполнением forward в FSDP может наблюдаться пауза, во время которой поток CPU не запускает никаких ядер. Это сделано намеренно и показывает, что ограничитель частоты работает. Такая синхронизация потока CPU предотвращает избыточное выделение памяти для последующих операций all-gather и фактически не должна задерживать выполнение ядер GPU.

Во время вычислений forward и backward FSDP заменяет параметры управляемых модулей представлениями torch.Tensor по причинам, связанным с autograd. Если при выполнении forward модуль использует сохранённые ссылки на параметры вместо повторного получения этих ссылок на каждой итерации, он не увидит новые представления, созданные FSDP, и autograd будет работать некорректно.

Наконец, при использовании sharding_strategy=ShardingStrategy.HYBRID_SHARD, когда группа процессов шардирования работает внутри узла, а группа процессов репликации — между узлами, параметр NCCL_CROSS_NIC=1 может помочь ускорить all-reduce для некоторых кластерных конфигураций.

Ограничения

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

  • В настоящее время FSDP не поддерживает накопление градиентов за пределами no_sync() при выгрузке данных на CPU. Это связано с тем, что FSDP использует вновь редуцированный градиент вместо его накопления с уже имеющимся, что может привести к некорректным результатам.
  • FSDP не поддерживает выполнение прямого прохода подмодуля, входящего в экземпляр FSDP. Параметры подмодуля будут шардированы, но сам подмодуль не является экземпляром FSDP, поэтому во время прямого прохода не будет выполнен необходимый all-gather полных параметров.
  • FSDP несовместим с двойным обратным проходом из-за способа регистрации хуков backward.
  • При замораживании параметров FSDP накладывает некоторые ограничения. Для use_orig_params=False каждый экземпляр FSDP должен управлять только замороженными или только незамороженными параметрами. Для use_orig_params=True FSDP поддерживает сочетание замороженных и незамороженных параметров, но рекомендуется этого избегать, чтобы не допустить неожиданно высокого потребления памяти градиентами.
  • Начиная с PyTorch 1.12, FSDP обеспечивает ограниченную поддержку общих параметров. Если для вашего сценария использования требуется расширенная поддержка общих параметров, оставьте сообщение в этой задаче.
  • Не следует изменять параметры между прямым и обратным проходами без использования контекста summon_full_params, поскольку изменения могут не сохраниться.
Параметры:
  • module (nn.Module) – Модуль, который нужно обернуть с помощью FSDP.
  • process_group (Optional[Union[ProcessGroup, Tuple[ProcessGroup, ProcessGroup]]]) – Группа процессов, между которыми распределяется модель и которая используется для коллективного обмена данными FSDP с помощью all-gather и reduce-scatter. Если None, FSDP использует группу процессов по умолчанию. Для стратегий гибкого шардирования, таких как ShardingStrategy.HYBRID_SHARD, можно передать кортеж групп процессов, представляющих группы, между которыми выполняется шардирование и репликация соответственно. Если None, FSDP создаёт группы процессов для шардирования внутри узла и репликации между узлами. (По умолчанию: None)
  • sharding_strategy (Optional[ShardingStrategy]) – Задаёт стратегию шардирования, позволяющую выбирать компромисс между экономией памяти и накладными расходами на обмен данными. Подробнее см. в разделе ShardingStrategy. (По умолчанию: FULL_SHARD)
  • cpu_offload (Optional[CPUOffload]) – Настраивает выгрузку данных на CPU. Если задано значение None, выгрузка на CPU не выполняется. Подробнее см. в разделе CPUOffload. (По умолчанию: None)

  • auto_wrap_policy (Optional[Union[Callable[[nn.Module, bool, int], bool], ModuleWrapPolicy, CustomPolicy]]) –

    Задаёт политику применения FSDP к подмодулям module, необходимую для перекрытия обмена данными и вычислений и, следовательно, влияющую на производительность. Если None, FSDP применяется только к module, а пользователям следует вручную применять FSDP к родительским модулям, двигаясь снизу вверх. Для удобства можно напрямую передать ModuleWrapPolicy, чтобы указать классы модулей для обёртывания (например, блок трансформера). В противном случае это должна быть вызываемая функция, принимающая три аргумента: module: nn.Module, recurse: bool и nonwrapped_numel: int, и возвращающая bool, которое указывает, следует ли применять FSDP к переданному module, если recurse=False, или продолжить обход поддерева модуля, если recurse=True. В вызываемую функцию можно добавить дополнительные аргументы. В size_based_auto_wrap_policy в torch.distributed.fsdp.wrap.py приведён пример вызываемой функции, которая применяет FSDP к модулю, если число элементов параметров в его поддереве превышает 100 млн. Рекомендуется выводить модель после применения FSDP и при необходимости корректировать её.

    Пример:

    >>> def custom_auto_wrap_policy(
>>>     module: nn.Module,
>>>     recurse: bool,
>>>     nonwrapped_numel: int,
>>>     # Additional custom arguments
>>>     min_num_params: int = int(1e8),
>>> ) -> bool:
>>>     return nonwrapped_numel >= min_num_params
>>> # Configure a custom `min_num_params`
>>> my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5))
  • backward_prefetch (Optional[BackwardPrefetch]) – Настраивает явную предварительную выборку all-gather во время обратного прохода. Если None, FSDP не выполняет предварительную выборку для обратного прохода, и перекрытия обмена данными и вычислений во время обратного прохода не происходит. Подробнее см. в разделе BackwardPrefetch. (По умолчанию: BACKWARD_PRE)
  • mixed_precision (Optional[MixedPrecision]) – Настраивает встроенную поддержку смешанной точности в FSDP. Если задано значение None, смешанная точность не используется. В противном случае можно задать типы данных параметров, буферов и редукции градиентов. Подробнее см. в разделе MixedPrecision. (По умолчанию: None)
  • ignored_modules (Optional[Iterable[torch.nn.Module]]) – Модули, собственные параметры которых, параметры дочерних модулей и буферы игнорируются этим экземпляром. Ни один из модулей непосредственно в ignored_modules не должен быть экземпляром FullyShardedDataParallel, а уже созданные экземпляры FullyShardedDataParallel не будут игнорироваться, если они вложены в этот экземпляр. Этот аргумент можно использовать, чтобы не шардировать отдельные параметры на уровне модуля при использовании auto_wrap_policy или если шардированием параметров не управляет FSDP. (По умолчанию: None)

  • param_init_fn (Optional[Callable[[nn.Module], None]]) –

    Callable[torch.nn.Module] -> None, задающая способ переноса модулей, находящихся в данный момент на meta-устройстве, на настоящее устройство с их инициализацией. Начиная с версии v1.12, FSDP обнаруживает модули с параметрами или буферами на meta-устройстве с помощью is_meta и либо применяет param_init_fn, если он задан, либо вызывает nn.Module.reset_parameters(). В обоих случаях реализация должна инициализировать только параметры и буферы модуля, но не его подмодулей. Это необходимо, чтобы избежать повторной инициализации. Кроме того, FSDP поддерживает отложенную инициализацию через API deferred_init() torchdistX (pytorch/torchdistX), при которой отложенные модули инициализируются вызовом param_init_fn, если он задан, или стандартного materialize_module() из torchdistX. Если задан param_init_fn, он применяется ко всем модулям на meta-устройстве, поэтому, вероятно, должен учитывать тип модуля. FSDP вызывает функцию инициализации до сведения параметров и шардирования.

    Пример:

    >>> module = MyModule(device="meta")
>>> def my_init_fn(module: nn.Module):
>>>     # E.g. initialize depending on the module type
>>>     ...
>>> fsdp_model = FSDP(module, param_init_fn=my_init_fn, auto_wrap_policy=size_based_auto_wrap_policy)
>>> print(next(fsdp_model.parameters()).device) # current CUDA device
>>> # With torchdistX
>>> module = deferred_init.deferred_init(MyModule, device="cuda")
>>> # Will initialize via deferred_init.materialize_module().
>>> fsdp_model = FSDP(module, auto_wrap_policy=size_based_auto_wrap_policy)
  • device_id (Optional[Union[int, torch.device]]) – int или torch.device, задающий устройство CUDA, на котором выполняется инициализация FSDP, включая инициализацию модуля при необходимости и шардирование параметров. Чтобы ускорить инициализацию, этот параметр следует задать, если module находится на CPU. Если задано CUDA-устройство по умолчанию (например, с помощью torch.cuda.set_device), пользователь может передать сюда torch.cuda.current_device. (По умолчанию: None)
  • sync_module_states (bool) – Если True, каждый модуль FSDP передаёт параметры и буферы модуля с ранга 0, чтобы обеспечить их репликацию на всех рангах (это увеличивает накладные расходы на обмен данными при создании экземпляра). Это может помочь загружать контрольные точки state_dict с помощью load_state_dict, эффективно расходуя память. Пример см. в разделе FullStateDictConfig. (По умолчанию: False)
  • forward_prefetch (bool) – Если True, FSDP явно выполняет предварительную выборку следующей операции all-gather прямого прохода до начала текущих вычислений прямого прохода. Это полезно только для нагрузок, ограниченных производительностью CPU; в таком случае более ранний запуск следующей операции all-gather может улучшить перекрытие. Используйте этот параметр только для моделей со статичным графом, поскольку предварительная выборка следует порядку выполнения первой итерации. (По умолчанию: False)
  • limit_all_gathers (bool) – Если True, FSDP явно синхронизирует поток CPU, чтобы использование памяти GPU было ограничено двумя последовательными экземплярами FSDP (текущим экземпляром, выполняющим вычисления, и следующим экземпляром, для которого предварительно выбрана операция all-gather). Если False, FSDP позволяет потоку CPU запускать операции all-gather без дополнительной синхронизации. (По умолчанию: True) Эта возможность часто называется «ограничителем частоты». Устанавливайте этот флаг в False только для отдельных нагрузок, ограниченных производительностью CPU и не испытывающих дефицита памяти: в этом случае поток CPU может активно запускать все ядра, не учитывая использование памяти GPU.
  • use_orig_params (bool) – Если задать значение True, FSDP будет использовать исходные параметры module. FSDP предоставляет эти исходные параметры пользователю через nn.Module.named_parameters() вместо внутренних FlatParameter FSDP. Это означает, что шаг оптимизатора выполняется над исходными параметрами, что позволяет задавать гиперпараметры для каждого исходного параметра. FSDP сохраняет переменные исходных параметров и изменяет их данные между нешардированным и шардированным представлениями; при этом они всегда являются представлениями базового нешардированного или шардированного FlatParameter соответственно. В текущем алгоритме шардированное представление всегда одномерное, поэтому исходная структура тензора теряется. Для заданного ранга исходный параметр может содержать все свои данные, часть данных или не содержать их вовсе. В последнем случае его данные будут представлены пустым тензором нулевого размера. Не следует создавать программы, которые рассчитывают на определённые данные исходного параметра в его шардированном представлении. Для использования torch.compile() требуется True. Если задать значение False, FSDP предоставит пользователю свои внутренние FlatParameter через nn.Module.named_parameters(). (По умолчанию: False)
  • ignored_states (Optional[Iterable[torch.nn.Parameter]], Optional[Iterable[torch.nn.Module]]) – Игнорируемые параметры или модули, которыми не будет управлять этот экземпляр FSDP. Это означает, что параметры не будут шардироваться, а их градиенты не будут редуцироваться между рангами. Этот аргумент объединяет существующий аргумент ignored_modules; вскоре мы можем объявить ignored_modules устаревшим. Для обратной совместимости сохраняются оба аргумента — ignored_states и ignored_modules`, но FSDP разрешает задавать только один из них, отличный от None.
  • device_mesh (Optional[DeviceMesh]) – DeviceMesh можно использовать вместо process_group. Если передан device_mesh, FSDP использует базовые группы процессов для коллективного обмена данными с помощью all-gather и reduce-scatter. Поэтому эти два аргумента взаимоисключающие. Для стратегий гибкого шардирования, таких как ShardingStrategy.HYBRID_SHARD, можно передать двумерный DeviceMesh вместо кортежа групп процессов. Для совместного использования 2D FSDP и TP необходимо передать device_mesh вместо process_group. Дополнительные сведения о DeviceMesh см. здесь: https://pytorch.org/tutorials/recipes/distributed_device_mesh.html
apply(fn) [исходный код]

Рекурсивно применяет fn к каждому подмодулю (возвращаемому .children()), а также к самому модулю.

Обычно этот метод используется для инициализации параметров модели (см. также torch.nn.init).

В отличие от torch.nn.Module.apply, эта версия дополнительно собирает полные параметры перед применением fn. Её не следует вызывать внутри другого контекста summon_full_params.

Параметры:

fn (Module -> None) – функция, применяемая к каждому подмодулю

Возвращает:

self

Тип возвращаемого значения:
Module
check_is_root() [исходный код]

Проверяет, является ли этот экземпляр корневым модулем FSDP.

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

bool

clip_grad_norm_(max_norm, norm_type=2.0) [исходный код]

Ограничивает норму градиентов всех параметров.

Норма вычисляется для градиентов всех параметров, рассматриваемых как единый вектор; градиенты изменяются на месте.

Параметры:
  • max_norm (float или int) – максимальная норма градиентов
  • norm_type (float или int) – тип используемой p-нормы. Для бесконечной нормы можно указать 'inf'.
Возвращает:

Общую норму параметров, рассматриваемых как единый вектор.

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

Tensor

Если во всех экземплярах FSDP используется NO_SHARD, то есть градиенты не шардируются между рангами, можно напрямую использовать torch.nn.utils.clip_grad_norm_().

Если хотя бы в одном экземпляре FSDP используется стратегия шардирования (то есть любая стратегия, кроме NO_SHARD), следует использовать этот метод вместо torch.nn.utils.clip_grad_norm_(), поскольку он учитывает шардирование градиентов между рангами.

Возвращаемая общая норма будет иметь «наибольший» тип данных среди всех параметров и градиентов согласно правилам продвижения типов PyTorch. Например, если все параметры и градиенты используют тип данных низкой точности, тип данных возвращаемой нормы также будет низкой точности. Если же хотя бы один параметр или градиент использует FP32, тип данных возвращаемой нормы будет FP32.

Предупреждение

Этот метод необходимо вызывать на всех рангах, поскольку он использует коллективный обмен данными.

static flatten_sharded_optim_state_dict(sharded_optim_state_dict, model, optim) [исходный код]

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

API аналогичен shard_full_optim_state_dict(). Единственное отличие состоит в том, что входной sharded_optim_state_dict должен быть возвращён методом sharded_optim_state_dict(). Поэтому на каждом ранге будут выполняться вызовы all-gather для сбора ShardedTensor.

Параметры:
  • sharded_optim_state_dict (Dict[str, Any]) – Словарь состояния оптимизатора, соответствующий несведённым параметрам и содержащий шардированное состояние оптимизатора.
  • model (torch.nn.Module) – См. описание в shard_full_optim_state_dict().
  • optim (torch.optim.Optimizer) – Оптимизатор для параметров model.
Возвращает:

См. описание в shard_full_optim_state_dict().

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

dict[str, Any]

forward(*args, **kwargs) [исходный код]

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

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

Any

static fsdp_modules(module, root_only=False) [исходный код]

Возвращает все вложенные экземпляры FSDP.

Список может включать сам module; корневые модули FSDP включаются только при root_only=True.

Параметры:
  • module (torch.nn.Module) – Корневой модуль, который может быть модулем FSDP, но это не обязательно.
  • root_only (bool) – Определяет, возвращать ли только корневые модули FSDP. (По умолчанию: False)
Возвращает:

Модули FSDP, вложенные во входной module.

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

List[FullyShardedDataParallel]

static full_optim_state_dict(model, optim, optim_input=None, rank0_only=True, group=None) [источник]

Возвращает полный словарь состояния оптимизатора.

Собирает полное состояние оптимизатора на ранге 0 и возвращает его в виде dict в соответствии с соглашением torch.optim.Optimizer.state_dict(), то есть с ключами "state" и "param_groups". Развёрнутые параметры в модулях FSDP, содержащихся в model, сопоставляются с исходными неразвёрнутыми параметрами.

Метод необходимо вызывать на всех рангах, поскольку он использует коллективные операции обмена данными. Однако, если rank0_only=True, словарь состояния заполняется только на ранге 0, а все остальные ранги возвращают пустой dict.

В отличие от torch.optim.Optimizer.state_dict(), этот метод использует в качестве ключей полные имена параметров, а не идентификаторы параметров.

Как и в torch.optim.Optimizer.state_dict(), тензоры в словаре состояния оптимизатора не клонируются, поэтому возможны неожиданные эффекты из-за псевдонимов. Рекомендуется сразу сохранять возвращённый словарь состояния оптимизатора, например, с помощью torch.save().

Параметры:
  • model (torch.nn.Module) – Корневой модуль (который может быть экземпляром FullyShardedDataParallel или не быть им), параметры которого были переданы оптимизатору optim.
  • optim (torch.optim.Optimizer) – Оптимизатор для параметров model.
  • optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – Входные данные, переданные оптимизатору optim, представляющие собой либо list групп параметров, либо итерируемый объект параметров; если None, то этот метод предполагает, что входные данные были model.parameters(). Этот аргумент устарел и больше не требуется. (По умолчанию: None)
  • rank0_only (bool) – Если True, сохраняет заполненный dict только на ранге 0; если False, сохраняет его на всех рангах. (По умолчанию: True)
  • group (dist.ProcessGroup) – Группа процессов модели или None, если используется группа процессов по умолчанию. (По умолчанию: None)
Возвращает:

dict, содержащий состояние оптимизатора для исходных неразвёрнутых параметров model и включающий ключи «state» и «param_groups» в соответствии с соглашением torch.optim.Optimizer.state_dict(). Если rank0_only=True, ненулевые ранги возвращают пустой dict.

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

Dict[str, Any]

static get_state_dict_type(module) [источник]

Получает тип state_dict и соответствующие настройки для модулей FSDP, корнем которых является module.

Целевой модуль не обязательно должен быть модулем FSDP.

Возвращает:

Объект StateDictSettings, содержащий текущие тип state_dict и настройки state_dict / optim_state_dict.

Вызывает исключение:
  • AssertionError` если настройки StateDictSettings для разных –
  • подмодулей FSDP различаются. –
Тип возвращаемого значения:

StateDictSettings

property module: Module

Возвращает обёрнутый модуль.

named_buffers(*args, **kwargs) [источник]

Возвращает итератор по буферам модуля, выдающий имя буфера и сам буфер.

Перехватывает имена буферов и удаляет все вхождения префикса развёрнутого буфера, специфичного для FSDP, внутри контекстного менеджера summon_full_params().

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

Iterator[tuple[str, Tensor]]

named_parameters(*args, **kwargs) [источник]

Возвращает итератор по параметрам модуля, выдающий имя параметра и сам параметр.

Перехватывает имена параметров и удаляет все вхождения префикса развёрнутого параметра, специфичного для FSDP, внутри контекстного менеджера summon_full_params().

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

Iterator[tuple[str, Parameter]]

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

Отключает синхронизацию градиентов между экземплярами FSDP.

Внутри этого контекста градиенты накапливаются в переменных модуля и синхронизируются при первом проходе вперёд-назад после выхода из контекста. Этот метод следует использовать только для корневого экземпляра FSDP; действие рекурсивно распространяется на все дочерние экземпляры FSDP.

Примечание

Это, вероятно, приведёт к увеличению потребления памяти, поскольку FSDP будет накапливать полные градиенты модели (вместо фрагментов градиентов) до последующей синхронизации.

Примечание

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

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

Generator

static optim_state_dict(model, optim, optim_state_dict=None, group=None) [источник]

Преобразует state-dict оптимизатора, соответствующий шардированной модели.

Переданный state-dict можно преобразовать в один из трёх типов: 1) полный state_dict оптимизатора, 2) шардированный state_dict оптимизатора, 3) локальный state_dict оптимизатора.

Для полного state_dict оптимизатора все состояния разворачиваются и не шардируются. Чтобы избежать OOM, можно задать параметры «только ранг 0» и «только CPU» с помощью state_dict_type().

Для шардированного state_dict оптимизатора все состояния разворачиваются, но остаются шардированными. Для дополнительной экономии памяти можно задать параметр «только CPU» с помощью state_dict_type().

Для локального state_dict преобразование не выполняется. Однако состояние преобразуется из nn.Tensor в ShardedTensor, чтобы отразить его шардирование (пока не поддерживается).

Пример:

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> from torch.distributed.fsdp import StateDictType
>>> from torch.distributed.fsdp import FullStateDictConfig
>>> from torch.distributed.fsdp import FullOptimStateDictConfig
>>> # Save a checkpoint
>>> model, optim = ...
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.FULL_STATE_DICT,
>>>     FullStateDictConfig(rank0_only=False),
>>>     FullOptimStateDictConfig(rank0_only=False),
>>> )
>>> state_dict = model.state_dict()
>>> optim_state_dict = FSDP.optim_state_dict(model, optim)
>>> save_a_checkpoint(state_dict, optim_state_dict)
>>> # Load a checkpoint
>>> model, optim = ...
>>> state_dict, optim_state_dict = load_a_checkpoint()
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.FULL_STATE_DICT,
>>>     FullStateDictConfig(rank0_only=False),
>>>     FullOptimStateDictConfig(rank0_only=False),
>>> )
>>> model.load_state_dict(state_dict)
>>> optim_state_dict = FSDP.optim_state_dict_to_load(
>>>     model, optim, optim_state_dict
>>> )
>>> optim.load_state_dict(optim_state_dict)
Параметры:
  • model (torch.nn.Module) – Корневой модуль (который может быть экземпляром FullyShardedDataParallel или не быть им), параметры которого были переданы оптимизатору optim.
  • optim (torch.optim.Optimizer) – Оптимизатор для параметров model.
  • optim_state_dict (Dict[str, Any]) – целевой state_dict оптимизатора для преобразования. Если значение равно None, будет использован optim.state_dict(). (По умолчанию: None)
  • group (dist.ProcessGroup) – Группа процессов модели, между которыми шардируются параметры, или None, если используется группа процессов по умолчанию. (По умолчанию: None)
Возвращает:

dict, содержащий состояние оптимизатора для model. Шардирование состояния оптимизатора определяется на основе state_dict_type.

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

Dict[str, Any]

static optim_state_dict_to_load(model, optim, optim_state_dict, is_named_optimizer=False, load_directly=False, group=None) [источник]

Преобразует state-dict оптимизатора, чтобы его можно было загрузить в оптимизатор, связанный с моделью FSDP.

Переданный optim_state_dict, преобразованный с помощью optim_state_dict(), преобразуется в развёрнутый state_dict оптимизатора, который можно загрузить в optim — оптимизатор для model. model должен быть шардирован с помощью FullyShardedDataParallel.

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> from torch.distributed.fsdp import StateDictType
>>> from torch.distributed.fsdp import FullStateDictConfig
>>> from torch.distributed.fsdp import FullOptimStateDictConfig
>>> # Save a checkpoint
>>> model, optim = ...
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.FULL_STATE_DICT,
>>>     FullStateDictConfig(rank0_only=False),
>>>     FullOptimStateDictConfig(rank0_only=False),
>>> )
>>> state_dict = model.state_dict()
>>> original_osd = optim.state_dict()
>>> optim_state_dict = FSDP.optim_state_dict(
>>>     model,
>>>     optim,
>>>     optim_state_dict=original_osd
>>> )
>>> save_a_checkpoint(state_dict, optim_state_dict)
>>> # Load a checkpoint
>>> model, optim = ...
>>> state_dict, optim_state_dict = load_a_checkpoint()
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.FULL_STATE_DICT,
>>>     FullStateDictConfig(rank0_only=False),
>>>     FullOptimStateDictConfig(rank0_only=False),
>>> )
>>> model.load_state_dict(state_dict)
>>> optim_state_dict = FSDP.optim_state_dict_to_load(
>>>     model, optim, optim_state_dict
>>> )
>>> optim.load_state_dict(optim_state_dict)
Параметры:
  • model (torch.nn.Module) – Корневой модуль (который может быть экземпляром FullyShardedDataParallel или не быть им), параметры которого были переданы оптимизатору optim.
  • optim (torch.optim.Optimizer) – Оптимизатор для параметров model.
  • optim_state_dict (Dict[str, Any]) – Состояния оптимизатора для загрузки.
  • is_named_optimizer (bool) – Является ли этот оптимизатор NamedOptimizer или KeyedOptimizer. Устанавливайте значение True, только если optim — это KeyedOptimizer из TorchRec или NamedOptimizer из torch.distributed.
  • load_directly (bool) – Если установлено значение True, этот API также вызовет optim.load_state_dict(result) перед возвратом результата. В противном случае пользователи должны самостоятельно вызвать optim.load_state_dict() (По умолчанию: False)
  • group (dist.ProcessGroup) – Группа процессов модели, между которыми шардируются параметры, или None, если используется группа процессов по умолчанию. (По умолчанию: None)
Тип возвращаемого значения:

dict[str, Any]

register_comm_hook(state, hook) [источник]

Регистрирует хук коммуникации.

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

Предупреждение

Хук коммуникации FSDP следует зарегистрировать до первого прямого прохода и только один раз.

Параметры:

  • state (object) –

    Передаётся хуку для хранения информации о состоянии в процессе обучения. Например, обратная связь по ошибке при сжатии градиентов, участники для следующего обмена данными в GossipGrad и т. д. Состояние хранится локально на каждом рабочем процессе и совместно используется всеми тензорами градиентов этого рабочего процесса.

  • hook (Callable) – Вызываемый объект с одной из следующих сигнатур: 1) hook: Callable[torch.Tensor] -> None: эта функция принимает тензор Python, представляющий полный, развёрнутый и нешардированный градиент для всех переменных модели, обёрнутой этим модулем FSDP (не обёрнутых другими подмодулями FSDP). Затем она выполняет всю необходимую обработку и возвращает None; 2) hook: Callable[torch.Tensor, torch.Tensor] -> None: эта функция принимает два тензора Python. Первый представляет полный, развёрнутый и нешардированный градиент для всех переменных модели, обёрнутой этим модулем FSDP (не обёрнутых другими подмодулями FSDP). Второй представляет предварительно выделенный тензор для хранения фрагмента шардированного градиента после редукции. В обоих случаях вызываемый объект выполняет всю необходимую обработку и возвращает None. Вызываемые объекты с сигнатурой 1 должны обрабатывать обмен градиентами для случая NO_SHARD. Вызываемые объекты с сигнатурой 2 должны обрабатывать обмен градиентами для шардированных случаев.
static rekey_optim_state_dict(optim_state_dict, optim_state_key_type, model, optim_input=None, optim=None) [источник]

Повторно назначает ключи словаря состояния оптимизатора optim_state_dict, используя тип ключей optim_state_key_type.

Это позволяет обеспечить совместимость словарей состояния оптимизатора моделей с экземплярами FSDP и без них.

Чтобы повторно назначить ключи полного словаря состояния оптимизатора FSDP (то есть полученного из full_optim_state_dict()), используя идентификаторы параметров, чтобы его можно было загрузить в необёрнутую модель:

>>> wrapped_model, wrapped_optim = ...
>>> full_osd = FSDP.full_optim_state_dict(wrapped_model, wrapped_optim)
>>> nonwrapped_model, nonwrapped_optim = ...
>>> rekeyed_osd = FSDP.rekey_optim_state_dict(full_osd, OptimStateKeyType.PARAM_ID, nonwrapped_model)
>>> nonwrapped_optim.load_state_dict(rekeyed_osd)

Чтобы повторно назначить ключи обычного словаря состояния оптимизатора необёрнутой модели, чтобы его можно было загрузить в обёрнутую модель:

>>> nonwrapped_model, nonwrapped_optim = ...
>>> osd = nonwrapped_optim.state_dict()
>>> rekeyed_osd = FSDP.rekey_optim_state_dict(osd, OptimStateKeyType.PARAM_NAME, nonwrapped_model)
>>> wrapped_model, wrapped_optim = ...
>>> sharded_osd = FSDP.shard_full_optim_state_dict(rekeyed_osd, wrapped_model)
>>> wrapped_optim.load_state_dict(sharded_osd)
Возвращает:

Словарь состояния оптимизатора с повторно назначенными ключами параметров, заданными в optim_state_key_type.

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

Dict[str, Any]

static scatter_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None, group=None) [источник]

Рассылает полный словарь состояния оптимизатора с ранга 0 на все остальные ранги.

Возвращает шардированный словарь состояния оптимизатора на каждом ранге. Возвращаемое значение совпадает с результатом shard_full_optim_state_dict(); на ранге 0 первым аргументом должно быть возвращаемое значение full_optim_state_dict().

Пример:

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> model, optim = ...
>>> full_osd = FSDP.full_optim_state_dict(model, optim)  # only non-empty on rank 0
>>> # Define new model with possibly different world size
>>> new_model, new_optim, new_group = ...
>>> sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, new_model, group=new_group)
>>> new_optim.load_state_dict(sharded_osd)

Примечание

Для получения шардированного словаря состояния оптимизатора, который будет загружен, можно использовать как shard_full_optim_state_dict(), так и scatter_full_optim_state_dict(). Если полный словарь состояния оптимизатора находится в памяти CPU, первый метод требует, чтобы полный словарь находился в памяти CPU на каждом ранге; каждый ранг самостоятельно шардирует словарь без обмена данными. Второй метод требует, чтобы полный словарь в памяти CPU находился только на ранге 0; ранг 0 перемещает каждый фрагмент в память GPU (для NCCL) и надлежащим образом передаёт его рангам. Таким образом, первый метод требует больше совокупной памяти CPU, а второй — больше затрат на обмен данными.

Параметры:
  • full_optim_state_dict (Optional[Dict[str, Any]]) – Словарь состояния оптимизатора, соответствующий неразвёрнутым параметрам и содержащий полное нешардированное состояние оптимизатора на ранге 0; на ненулевых рангах этот аргумент игнорируется.
  • model (torch.nn.Module) – Корневой модуль (который может быть экземпляром FullyShardedDataParallel или не быть им), параметры которого соответствуют состоянию оптимизатора в full_optim_state_dict.
  • optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – Данные, переданные оптимизатору, представляющие собой либо list групп параметров, либо итерируемый объект параметров; если None, то этот метод предполагает, что входные данные были model.parameters(). Этот аргумент устарел и больше не требуется. (По умолчанию: None)
  • optim (Optional[torch.optim.Optimizer]) – Оптимизатор, который загрузит словарь состояния, возвращённый этим методом. Предпочтительно использовать этот аргумент вместо optim_input. (По умолчанию: None)
  • group (dist.ProcessGroup) – Группа процессов модели или None, если используется группа процессов по умолчанию. (По умолчанию: None)
Возвращает:

Полный словарь состояния оптимизатора, в котором теперь используются развёрнутые параметры вместо неразвёрнутых и оставлена только часть состояния оптимизатора, относящаяся к этому рангу.

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

Dict[str, Any]

static set_state_dict_type(module, state_dict_type, state_dict_config=None, optim_state_dict_config=None) [источник]

Задаёт state_dict_type для всех дочерних модулей FSDP целевого модуля.

Также принимает (необязательные) настройки словарей состояния модели и оптимизатора. Целевой модуль не обязательно должен быть модулем FSDP. Если целевой модуль является модулем FSDP, его state_dict_type также будет изменён.

Примечание

Этот API следует вызывать только для модуля верхнего уровня (корневого модуля).

Примечание

Этот API позволяет прозрачно использовать стандартный API state_dict для создания контрольных точек модели в случаях, когда корневой модуль FSDP обёрнут другим nn.Module. Например, следующий код обеспечивает вызов state_dict для всех экземпляров, не относящихся к FSDP, и передачу вызовов реализации sharded_state_dict для FSDP:

Пример:

>>> model = DDP(FSDP(...))
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.SHARDED_STATE_DICT,
>>>     state_dict_config = ShardedStateDictConfig(offload_to_cpu=True),
>>>     optim_state_dict_config = OptimStateDictConfig(offload_to_cpu=True),
>>> )
>>> param_state_dict = model.state_dict()
>>> optim_state_dict = FSDP.optim_state_dict(model, optim)
Параметры:
  • module (torch.nn.Module) – Корневой модуль.
  • state_dict_type (StateDictType) – требуемый тип state_dict_type.
  • state_dict_config (Optional[StateDictConfig]) – настройки целевого state_dict_type.
  • optim_state_dict_config (Optional[OptimStateDictConfig]) – настройки словаря состояния оптимизатора.
Возвращает:

Объект StateDictSettings, содержащий предыдущий тип state_dict и настройки модуля.

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

StateDictSettings

static shard_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None) [исходный код]

Разбивает полный словарь состояния оптимизатора на шарды.

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

Пример:

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> model, optim = ...
>>> full_osd = FSDP.full_optim_state_dict(model, optim)
>>> torch.save(full_osd, PATH)
>>> # Define new model with possibly different world size
>>> new_model, new_optim = ...
>>> full_osd = torch.load(PATH)
>>> sharded_osd = FSDP.shard_full_optim_state_dict(full_osd, new_model)
>>> new_optim.load_state_dict(sharded_osd)

Примечание

Для получения шардированного словаря состояния оптимизатора для загрузки можно использовать как shard_full_optim_state_dict(), так и scatter_full_optim_state_dict(). Если полный словарь состояния оптимизатора находится в памяти CPU, первый метод требует, чтобы полный словарь находился в памяти CPU на каждом ранге; каждый ранг самостоятельно разбивает словарь на шарды без обмена данными. Второй метод требует, чтобы полный словарь в памяти CPU был только у ранга 0; ранг 0 перемещает каждый шард в память GPU (для NCCL) и передает его соответствующим рангам. Таким образом, первый метод требует больше суммарной памяти CPU, а второй — больше затрат на обмен данными.

Параметры:
  • full_optim_state_dict (Dict[str, Any]) – Словарь состояния оптимизатора, соответствующий несглаженным параметрам и содержащий полное нешардированное состояние оптимизатора.
  • model (torch.nn.Module) – Корневой модуль (который может быть экземпляром FullyShardedDataParallel или не быть им), параметры которого соответствуют состоянию оптимизатора в full_optim_state_dict.
  • optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – Входные данные, передаваемые оптимизатору: список групп параметров list или итерируемый объект параметров; если None, метод предполагает, что входные данные имели вид model.parameters(). Этот аргумент устарел, передавать его больше не нужно. (По умолчанию: None)
  • optim (Optional[torch.optim.Optimizer]) – Оптимизатор, который загрузит словарь состояния, возвращаемый этим методом. Предпочтительно использовать этот аргумент вместо optim_input. (По умолчанию: None)
Возвращает:

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

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

Dict[str, Any]

static sharded_optim_state_dict(model, optim, group=None) [исходный код]

Возвращает словарь состояния оптимизатора в шардированной форме.

Этот API похож на full_optim_state_dict(), но разбивает все состояния с ненулевой размерностью на ShardedTensor для экономии памяти. Этот API следует использовать только в том случае, если модель state_dict получена с помощью менеджера контекста with state_dict_type(SHARDED_STATE_DICT):.

Подробные сведения об использовании приведены в разделе full_optim_state_dict().

Предупреждение

Возвращаемый словарь состояния содержит ShardedTensor и не может напрямую использоваться обычным optim.load_state_dict.

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

dict[str, Any]

static state_dict_type(module, state_dict_type, state_dict_config=None, optim_state_dict_config=None) [исходный код]

Задает state_dict_type для всех дочерних модулей FSDP целевого модуля.

Этот менеджер контекста выполняет те же функции, что и set_state_dict_type(). Подробные сведения см. в документации к set_state_dict_type().

Пример:

>>> model = DDP(FSDP(...))
>>> with FSDP.state_dict_type(
>>>     model,
>>>     StateDictType.SHARDED_STATE_DICT,
>>> ):
>>>     checkpoint = model.state_dict()
Параметры:
  • module (torch.nn.Module) – Корневой модуль.
  • state_dict_type (StateDictType) – требуемый state_dict_type.
  • state_dict_config (Optional[StateDictConfig]) – конфигурация state_dict модели для целевого state_dict_type.
  • optim_state_dict_config (Optional[OptimStateDictConfig]) – конфигурация state_dict оптимизатора для целевого state_dict_type.
Тип возвращаемого значения:

Generator

static summon_full_params(module, recurse=True, writeback=True, rank0_only=False, offload_to_cpu=False, with_grads=False) [исходный код]

Предоставляет доступ к полным параметрам экземпляров FSDP с помощью этого менеджера контекста.

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

Примечание

Можно использовать для внутренних модулей FSDP.

Примечание

Нельзя использовать во время прямого или обратного прохода. Также нельзя запускать прямой или обратный проход изнутри этого контекста.

Примечание

После выхода из менеджера контекста параметры снова становятся локальными шардами; поведение хранилища такое же, как при прямом проходе.

Примечание

Полные параметры можно изменять, но после выхода из менеджера контекста сохранится только часть, соответствующая локальному шарду параметров (если только не задано writeback=False — в этом случае изменения будут отброшены). Если FSDP не разбивает параметры на шарды, что в настоящее время происходит только при world_size == 1 или конфигурации NO_SHARD, изменения сохраняются независимо от writeback.

Примечание

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

Предупреждение

Обратите внимание: сочетание rank0_only=True и writeback=True в настоящее время не поддерживается и приведет к ошибке. Это связано с тем, что формы параметров модели внутри контекста будут различаться между рангами, а их изменение может привести к несогласованности между рангами после выхода из контекста.

Предупреждение

Обратите внимание: сочетание offload_to_cpu и rank0_only=False приведет к избыточному копированию полных параметров в память CPU для графических процессоров, расположенных на одном компьютере, что может вызвать нехватку памяти CPU. Рекомендуется использовать offload_to_cpu с rank0_only=True.

Параметры:
  • recurse (bool, Optional) – рекурсивно предоставляет все параметры вложенных экземпляров FSDP (по умолчанию: True).
  • writeback (bool, Optional) – если False, изменения параметров отбрасываются после выхода из менеджера контекста; отключение этого параметра может немного повысить эффективность (по умолчанию: True)
  • rank0_only (bool, Optional) – если True, полные параметры материализуются только на глобальном ранге 0. Это означает, что внутри контекста полные параметры будут только у ранга 0, а у остальных рангов будут шардированные параметры. Обратите внимание: сочетание rank0_only=True и writeback=True не поддерживается, поскольку формы параметров модели внутри контекста будут различаться между рангами, а их изменение может привести к несогласованности между рангами после выхода из контекста.
  • offload_to_cpu (bool, Optional) – если True, полные параметры выгружаются в память CPU. Обратите внимание: сейчас эта выгрузка выполняется только в том случае, если параметры шардированы (что не так только при world_size = 1 или конфигурации NO_SHARD). Во избежание избыточного копирования параметров модели в одну и ту же память CPU рекомендуется использовать offload_to_cpu с rank0_only=True.
  • with_grads (bool, Optional) – если True, градиенты также снимаются с шардирования вместе с параметрами. В настоящее время это поддерживается только при передаче use_orig_params=True конструктору FSDP и offload_to_cpu=False этому методу. (По умолчанию: False)
Тип возвращаемого значения:

Generator

class torch.distributed.fsdp.BackwardPrefetch(value) [исходный код]

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

  • BACKWARD_PRE: обеспечивает максимальное перекрытие, но и наибольшее увеличение использования памяти. Следующий набор параметров загружается до вычисления градиентов для текущего набора параметров. Это позволяет перекрыть следующую операцию all-gather и текущие вычисления градиентов; в пиковый момент в памяти находятся текущий набор параметров, следующий набор параметров и текущий набор градиентов.
  • BACKWARD_POST: обеспечивает меньшее перекрытие, но требует меньше памяти. Следующий набор параметров загружается после вычисления градиентов для текущего набора параметров. Это позволяет перекрыть текущую операцию reduce-scatter и следующие вычисления градиентов; текущий набор параметров освобождается до выделения памяти для следующего набора. В пиковый момент в памяти находятся только следующий набор параметров и текущий набор градиентов.
  • Аргумент backward_prefetch в FSDP принимает значение None, отключающее предварительную загрузку для обратного прохода. Перекрытия при этом нет, а использование памяти не увеличивается. В целом мы не рекомендуем этот параметр, поскольку он может значительно снизить пропускную способность.

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

class torch.distributed.fsdp.ShardingStrategy(value) [исходный код]

Задает стратегию шардирования для распределенного обучения с помощью FullyShardedDataParallel.

  • FULL_SHARD: параметры, градиенты и состояния оптимизатора шардируются. Для параметров эта стратегия снимает шардирование (с помощью all-gather) перед прямым проходом, снова разбивает их на шарды после прямого прохода, снимает шардирование перед вычислениями обратного прохода и снова разбивает на шарды после обратного прохода. Градиенты синхронизируются и разбиваются на шарды (с помощью reduce-scatter) после обратного прохода. Шардированные состояния оптимизатора обновляются локально на каждом ранге.
  • SHARD_GRAD_OP: градиенты и состояния оптимизатора шардируются во время вычислений, а параметры дополнительно шардируются вне вычислений. Для параметров эта стратегия снимает шардирование перед прямым проходом, не разбивает их на шарды после прямого прохода и делает это только после обратного прохода. Шардированные состояния оптимизатора обновляются локально на каждом ранге. Внутри no_sync() параметры не разбиваются на шарды после обратного прохода.
  • NO_SHARD: параметры, градиенты и состояния оптимизатора не шардируются, а реплицируются на всех рангах, подобно API DistributedDataParallel в PyTorch. Градиенты синхронизируются (с помощью all-reduce) после обратного прохода. Нешардированные состояния оптимизатора обновляются локально на каждом ранге.
  • HYBRID_SHARD: применяет FULL_SHARD внутри узла и реплицирует параметры между узлами. Это уменьшает объем передаваемых данных, поскольку ресурсоемкие операции all-gather и reduce-scatter выполняются только внутри узла, что может повысить производительность для моделей среднего размера.
  • _HYBRID_SHARD_ZERO2: применяет SHARD_GRAD_OP внутри узла и реплицирует параметры между узлами. Этот вариант похож на HYBRID_SHARD, но может обеспечить еще более высокую пропускную способность, поскольку нешардированные параметры не освобождаются после прямого прохода, что позволяет избежать операций all-gather перед обратным проходом.
class torch.distributed.fsdp.MixedPrecision(param_dtype=None, reduce_dtype=None, buffer_dtype=None, keep_low_precision_grads=False, cast_forward_inputs=False, cast_root_forward_inputs=True, _module_classes_to_ignore=(<class 'torch.nn.modules.batchnorm._BatchNorm'>, )) [исходный код]

Настраивает смешанную точность, реализованную средствами FSDP.

Переменные:
  • param_dtype (Optional[torch.dtype]) – Задает тип данных параметров модели во время прямого и обратного проходов, а значит, и тип данных для вычислений в этих проходах. Вне прямого и обратного проходов шардированные параметры хранятся в полной точности (например, для шага оптимизатора), а при сохранении контрольной точки модели параметры всегда сохраняются в полной точности. (По умолчанию: None)
  • reduce_dtype (Optional[torch.dtype]) – Задает тип данных для редукции градиентов (то есть reduce-scatter или all-reduce). Если это значение равно None, а param_dtype не равно None, используется значение param_dtype, и редукция градиентов по-прежнему выполняется с низкой точностью. Это значение может отличаться от param_dtype, например, чтобы выполнять редукцию градиентов с полной точностью. (По умолчанию: None)
  • buffer_dtype (Optional[torch.dtype]) – Задает тип данных буферов. FSDP не разбивает буферы на шарды, а преобразует их в buffer_dtype при первом прямом проходе и далее хранит в этом типе данных. При сохранении контрольной точки модели буферы сохраняются в полной точности, за исключением LOCAL_STATE_DICT. (По умолчанию: None)
  • keep_low_precision_grads (bool) – Если False, после обратного прохода FSDP преобразует градиенты в полную точность перед шагом оптимизатора. Если True, FSDP оставляет градиенты в типе данных, использованном для редукции градиентов; это может сэкономить память при использовании пользовательского оптимизатора, поддерживающего работу с низкой точностью. (По умолчанию: False)
  • cast_forward_inputs (bool) – Если True, этот модуль FSDP преобразует позиционные и именованные аргументы прямого прохода в param_dtype. Это позволяет обеспечить совпадение типов данных параметров и входных данных для вычислений прямого прохода, как того требуют многие операции. Может потребоваться установить значение True, если смешанная точность применяется только к некоторым модулям FSDP; в этом случае подмодулю FSDP со смешанной точностью нужно повторно преобразовать типы входных данных. (По умолчанию: False)
  • cast_root_forward_inputs (bool) – Если True, корневой модуль FSDP преобразует позиционные и именованные аргументы прямого прохода в param_dtype, переопределяя значение cast_forward_inputs. На некорневые модули FSDP это не влияет. (По умолчанию: True)
  • _module_classes_to_ignore (collections.abc.Sequence[type[torch.nn.modules.module.Module]]) – (Sequence[Type[nn.Module]]): Задает классы модулей, которые следует исключить из смешанной точности при использовании auto_wrap_policy: к модулям этих классов FSDP будет применяться отдельно, без смешанной точности (это означает, что итоговая конфигурация FSDP будет отличаться от заданной политикой). Если auto_wrap_policy не задан, этот параметр ни на что не влияет. API экспериментальный и может измениться. (По умолчанию: (_BatchNorm,))

Примечание

Этот API является экспериментальным и может измениться.

Примечание

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

Примечание

В summon_full_params параметры принудительно переводятся в полную точность, а буферы — нет.

Примечание

Нормализация по слоям и пакетам накапливает значения в float32, даже если входные данные имеют низкую точность, например float16 или bfloat16. Отключение смешанной точности FSDP только для таких модулей нормализации означает, что аффинные параметры сохраняются в float32. Однако это приводит к выполнению отдельных операций all-gather и reduce-scatter для таких модулей нормализации, что может быть неэффективно. Поэтому, если позволяет рабочая нагрузка, предпочтительно применять смешанную точность и к этим модулям.

Примечание

По умолчанию, если пользователь передает модель с модулями _BatchNorm и задает auto_wrap_policy, к модулям пакетной нормализации FSDP применяется отдельно, без смешанной точности. См. аргумент _module_classes_to_ignore.

Примечание

По умолчанию MixedPrecision имеет значения cast_root_forward_inputs=True и cast_forward_inputs=False. Для корневого экземпляра FSDP его cast_root_forward_inputs имеет приоритет над cast_forward_inputs. Для некорневых экземпляров FSDP их значения cast_root_forward_inputs игнорируются. Настройки по умолчанию подходят для типичного случая, когда у всех экземпляров FSDP одинаковая конфигурация MixedPrecision и входные данные нужно преобразовать в param_dtype в начале прямого прохода модели.

Примечание

Для вложенных экземпляров FSDP с разными конфигурациями MixedPrecision рекомендуется задавать индивидуальные значения cast_forward_inputs, чтобы настраивать преобразование входных данных перед прямым проходом каждого экземпляра. В этом случае, поскольку преобразование выполняется перед прямым проходом каждого экземпляра FSDP, родительский экземпляр FSDP должен запускать свои подмодули, не относящиеся к FSDP, до подмодулей FSDP, чтобы избежать изменения типа данных активаций из-за другой конфигурации MixedPrecision.

Пример:

>>> model = nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3))
>>> model[1] = FSDP(
>>>     model[1],
>>>     mixed_precision=MixedPrecision(param_dtype=torch.float16, cast_forward_inputs=True),
>>> )
>>> model = FSDP(
>>>     model,
>>>     mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, cast_forward_inputs=True),
>>> )

В примере выше показан рабочий вариант. Если же заменить model[1] на model[0], то есть если подмодуль с другой конфигурацией MixedPrecision первым выполнит прямой проход, то model[1] ошибочно получит активации типа float16 вместо активаций типа bfloat16.

class torch.distributed.fsdp.CPUOffload(offload_params=False) [исходный код]

Настраивает выгрузку данных в память CPU.

Переменные:

offload_params (bool) – Задает, следует ли выгружать параметры в память CPU, когда они не участвуют в вычислениях. Если True, градиенты также выгружаются в память CPU, а шаг оптимизатора выполняется на CPU.

class torch.distributed.fsdp.StateDictConfig(offload_to_cpu=False) [исходный код]

StateDictConfig — базовый класс для всех классов конфигурации state_dict. Для настройки параметров соответствующего типа state_dict, поддерживаемого FSDP, пользователям следует создавать экземпляр дочернего класса (например, FullStateDictConfig).

Переменные:

offload_to_cpu (bool) – Если True, FSDP выгружает значения словаря состояния в память CPU; если False, FSDP оставляет их в памяти GPU. (По умолчанию: False)

class torch.distributed.fsdp.FullStateDictConfig(offload_to_cpu=False, rank0_only=False) [исходный код]

FullStateDictConfig — класс конфигурации, предназначенный для использования с StateDictType.FULL_STATE_DICT. При сохранении полных словарей состояния рекомендуется включить и offload_to_cpu=True, и rank0_only=True, чтобы сэкономить память GPU и память CPU соответственно. Этот класс конфигурации предназначен для использования через менеджер контекста state_dict_type() следующим образом:

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> fsdp = FSDP(model, auto_wrap_policy=...)
>>> cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
>>> with FSDP.state_dict_type(fsdp, StateDictType.FULL_STATE_DICT, cfg):
>>>     state = fsdp.state_dict()
>>> # `state` will be empty on non rank 0 and contain CPU tensors on rank 0.
>>> # To reload checkpoint for inference, finetuning, transfer learning, etc:
>>> model = model_fn()  # Initialize model in preparation for wrapping with FSDP
>>> if dist.get_rank() == 0:
>>> # Load checkpoint only on rank 0 to avoid memory redundancy
>>>     state_dict = torch.load("my_checkpoint.pt")
>>>     model.load_state_dict(state_dict)
>>> # All ranks initialize FSDP module as usual. `sync_module_states` argument
>>> # communicates loaded checkpoint states from rank 0 to rest of the world.
>>> fsdp = FSDP(
...     model,
...     device_id=torch.cuda.current_device(),
...     auto_wrap_policy=...,
...     sync_module_states=True,
... )
>>> # After this point, all ranks have FSDP model with loaded checkpoint.
Переменные:

rank0_only (bool) – Если True, полный словарь состояния сохраняет только ранг 0, а остальные ранги сохраняют пустой словарь. Если False, полный словарь состояния сохраняют все ранги. (По умолчанию: False)

class torch.distributed.fsdp.ShardedStateDictConfig(offload_to_cpu=False, _use_dtensor=False) [исходный код]

ShardedStateDictConfig — класс конфигурации, предназначенный для использования с StateDictType.SHARDED_STATE_DICT.

Переменные:

_use_dtensor (bool) – Если True, FSDP сохраняет значения словаря состояния как DTensor; если False, FSDP сохраняет их как ShardedTensor. (По умолчанию: False)

Предупреждение

_use_dtensor — закрытое поле ShardedStateDictConfig, используемое FSDP для определения типа значений словаря состояния. Пользователям не следует вручную изменять _use_dtensor.

class torch.distributed.fsdp.LocalStateDictConfig(offload_to_cpu: bool = False) [исходный код]
class torch.distributed.fsdp.OptimStateDictConfig(offload_to_cpu=True) [исходный код]

OptimStateDictConfig — базовый класс для всех классов конфигурации optim_state_dict. Чтобы настроить параметры для соответствующего типа optim_state_dict, поддерживаемого FSDP, пользователям следует создать экземпляр дочернего класса (например, FullOptimStateDictConfig).

Переменные:

offload_to_cpu (bool) – Если True, FSDP переносит тензорные значения словаря состояния на CPU; если False, FSDP оставляет их на исходном устройстве (GPU, если не включена выгрузка параметров на CPU). (По умолчанию: True)

class torch.distributed.fsdp.FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=False) [исходный код]
Переменные:

rank0_only (bool) – Если True, полный словарь состояния сохраняет только rank 0, а ранги, отличные от нуля, сохраняют пустой словарь. Если False, полный словарь состояния сохраняют все ранги. (По умолчанию: False)

class torch.distributed.fsdp.ShardedOptimStateDictConfig(offload_to_cpu=True, _use_dtensor=False) [исходный код]

ShardedOptimStateDictConfig — класс конфигурации, предназначенный для использования с StateDictType.SHARDED_STATE_DICT.

Переменные:

_use_dtensor (bool) – Если True, FSDP сохраняет значения словаря состояния как DTensor; если False, FSDP сохраняет их как ShardedTensor. (По умолчанию: False)

Предупреждение

_use_dtensor — закрытое поле ShardedOptimStateDictConfig, используемое FSDP для определения типа значений словаря состояния. Пользователям не следует изменять _use_dtensor вручную.

class torch.distributed.fsdp.LocalOptimStateDictConfig(offload_to_cpu: bool = False) [исходный код]
class torch.distributed.fsdp.StateDictSettings(state_dict_type: torch.distributed.fsdp.api.StateDictType, state_dict_config: torch.distributed.fsdp.api.StateDictConfig, optim_state_dict_config: torch.distributed.fsdp.api.OptimStateDictConfig) [исходный код]
torch.distributed.fsdp.wrap.always_wrap_policy(*args, **kwargs) [исходный код]

Простая рекурсивная политика оборачивания, которая всегда возвращает True. Это означает, что каждый подмодуль оборачивается классом-оберткой, указанным в _recursive_wrap().

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

bool

torch.distributed.fsdp.wrap.enable_wrap(*, wrapper_cls, **wrapper_kwargs) [исходный код]

Менеджер контекста для оборачивания модулей с помощью обертки.

Полезен, если вы хотите применять одни и те же аргументы конфигурации ко всем оборачиваемым дочерним модулям. Особенно важный сценарий использования — оборачивание больших слоев, чтобы во время инициализации они разбивались на шарды (на месте) и не приводили к нехватке системной памяти. Большие слои могут указывать на необходимость разбиения на шарды с помощью аннотации wrap; этот менеджер контекста позволяет задать точную конфигурацию для таких вложенных экземпляров.

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

with enable_wrap(wrapper_cls, **params):
    # Wraps layer in FSDP by default if within context
    self.l1 = wrap(torch.nn.Linear(5, 5))
Параметры:
  • wrapper_cls (Any) – Класс, которым аннотация wrap будет wrap модули, например FullyShardedDataParallel.
  • **wrapper_kwargs (Any) – Параметры конфигурации, которые будут переданы всем экземплярам wrap внутри контекста
Тип возвращаемого значения:

Generator[None, None, None]

torch.distributed.fsdp.wrap.lambda_auto_wrap_policy(module, recurse, nonwrapped_numel, lambda_fn) [исходный код]

Удобная политика автоматического оборачивания подмодулей на основе произвольной пользовательской функции. Если lambda_fn(submodule) == True`, подмодуль будет обернут как единый блок wrapper_cls.

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

Первые три параметра требуются для _recursive_wrap().

Параметры:
  • module (nn.Module) – Рассматриваемый текущий модуль.
  • recurse (bool) – Если False, эта функция должна решить, следует ли оборачивать module как экземпляр FSDP. Если True, функция продолжает рекурсивный обход дерева модулей в рамках DFS.
  • nonwrapped_numel (int) – Число элементов параметров, которые еще не обернуты.
  • lambda_fn (Callable[[nn.Module], bool]) – Если функция возвращает True, этот модуль будет обернут.
Тип возвращаемого значения:

bool

torch.distributed.fsdp.wrap.size_based_auto_wrap_policy(module, recurse, nonwrapped_numel, min_num_params=100000000, force_leaf_modules=None, exclude_wrap_modules=None) [исходный код]

Политика автоматического оборачивания на основе размера.

Параметры:
  • module (nn.Module) – Рассматриваемый текущий модуль.
  • recurse (bool) – Если False, эта функция должна решить, следует ли оборачивать module как экземпляр FSDP. Если True, функция продолжает рекурсивный обход дерева модулей в рамках DFS.
  • nonwrapped_numel (int) – Число элементов параметров, которые еще не обернуты.
  • min_num_params (int) – Настраиваемый параметр политики, задающий порог размера, при превышении которого модуль готов к оборачиванию. Единица измерения — число элементов.
  • force_leaf_modules (Optional[set[type[nn.Module]]]) – Множество типов модулей, которые следует оставить листьями, то есть их дочерние модули никогда не будут оборачиваться.
  • exclude_wrap_modules (Optional[set[type[nn.Module]]]) – Множество типов модулей, которые следует исключить из оборачивания.
Возвращает:

Следует ли оборачивать module.

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

bool

torch.distributed.fsdp.wrap.transformer_auto_wrap_policy(module, recurse, nonwrapped_numel, transformer_layer_cls) [исходный код]

См. _module_wrap_policy(), где transformer_layer_cls совпадает с module_classes. Обратите внимание: общие параметры должны быть обернуты в одном экземпляре FSDP, поэтому эта политика автоматического оборачивания помогает поместить общие эмбеддинги в один экземпляр FSDP в моделях-трансформерах.

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

bool

torch.distributed.fsdp.wrap.wrap(module, **wrap_overrides) [исходный код]

Помечает модуль как подлежащий оборачиванию. Помеченные модули будут обернуты, только если находятся внутри менеджера контекста enable_wrap(). Это позволяет инициализировать модуль как с оберткой, так и без нее, не меняя код.

Класс, которым эта функция оборачивает переданный nn.Module, задается аргументом wrapper_cls, переданным в enable_wrap. В enable_wrap и wrap можно передать именованные аргументы, задающие способ создания экземпляра wrapper_cls. Если в enable_wrap и wrap есть одинаковые именованные аргументы, приоритет имеет аргумент, переданный в wrap.

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

with enable_wrap(wrapper_cls=FSDP, **fsdp_config):
    # Wraps layer in FSDP by default if within context
    self.l1 = wrap(torch.nn.Linear(5, 5))
Параметры:
  • module (nn.Module) – модуль для оборачивания (если он находится в контексте enable_wrap())
  • **wrap_overrides (Any) – переопределения конфигурации, имеющие приоритет над значениями, заданными в контексте enable_wrap()
Тип возвращаемого значения:
Module

© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/fsdp.html

Spec-Zone.ru

Настройки Оффлайн Что нового Помощь О нас
Spec-Zone .ru
спецификации, руководства, описания, API