FullyShardedDataParallel

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) [source]

Обёртка для фрагментации параметров модуля по всем рабочим процессам параллельной обработки данных. Она вдохновлена работой Xu и соавт., а также этапом 3 ZeRO из 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 фрагментирует и преобразует параметры модуля, что может не сохранить исходные переменные параметров. Таким образом, ранее инициализированный оптимизатор может содержать устаревшие ссылки на параметры.

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

Если целевой CUDA-устройство имеет ID dev_id, то либо (1) module должно быть размещено на этом устройстве, (2) устройство должно быть установлено с помощью torch.cuda.set_device(dev_id), или (3) dev_id должно быть передано в аргумент конструктора device_id. Вычислительное устройство этого экземпляра FSDP будет целевым устройством. Для (1) и (3) инициализация FSDP всегда происходит на GPU. Для (2) инициализация FSDP происходит на текущем устройстве module, которое может быть процессором.

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

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

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

Изменение имен исходных переменных параметров после построения приведёт к неопределённому поведению.

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

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

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

Начиная с PyTorch 1.12, FSDP предлагает ограниченную поддержку общих параметров (например, установка веса одного слоя Linear другому). В частности, модули, которые используют общие параметры, должны быть обернуты в единицу FSDP. Если для вашего случая использования требуется улучшенная поддержка общих параметров, пожалуйста, обратитесь на https://github.com/pytorch/pytorch/issues/77724

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

У FSDP есть некоторые ограничения на замораживание параметров (т.е. установка param.requires_grad=False). Для use_orig_params=False, каждый экземпляр FSDP должен управлять параметрами, которые все заморожены или все не заморожены. Для use_orig_params=True, FSDP поддерживает смешение замороженных и незамороженных, но мы рекомендуем этого не делать, так как тогда использование памяти градиента будет выше, чем ожидалось (а именно, эквивалентно тому, как если бы эти параметры не были заморожены). Это означает, что в идеале замороженные параметры должны быть изолированы в свои собственные nn.Module и обернуты отдельно с помощью FSDP.

Примечание

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

Примечание

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

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

Пользователь не должен изменять параметры между прямым и обратным проходом без использования контекста summon_full_params(), так как внесённые изменения могут не сохраниться. Более того, для use_orig_params=False, доступ к исходным параметрам между прямым и обратным проходами может вызвать ошибку доступа к памяти.

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

Для use_orig_params=True, ShardingStrategy.SHARD_GRAD_OP предоставляет нефрагментированные параметры, а не фрагментированные параметры, после прямого прохода, так как он не освобождает нефрагментированные, в отличие от ShardingStrategy.FULL_SHARD. Одним из нюансов является то, что, поскольку градиенты всегда фрагментированы или None, ShardingStrategy.SHARD_GRAD_OP не предоставит фрагментированные градиенты с нефрагментированными параметрами после прямого прохода. Если вы хотите проверить градиенты, попробуйте summon_full_params() с with_grads=True.

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

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

Примечание

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

Примечание

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

Параметры

• module (nn.Module) – Это модуль, который будет обернут с помощью FSDP.
• process_group (Необязательно[Union[ProcessGroup, Tuple[ProcessGroup, ProcessGroup]]]) – Это группа процессов, по которой модель фрагментируется, и, следовательно, та, которая используется для коллективных коммуникаций FSDP all-gather и reduce-scatter. Если None, то FSDP использует группу процессов по умолчанию. Для гибридных стратегий фрагментации, таких как ShardingStrategy.HYBRID_SHARD, пользователи могут передать кортеж из групп процессов, представляющий группы, по которым фрагментировать и реплицировать соответственно. Если None, то FSDP создает группы процессов для пользователя, чтобы фрагментировать внутри узла и реплицировать между узлами. (По умолчанию: None)
• sharding_strategy (Необязательно[ShardingStrategy]) – Это настраивает стратегию фрагментации, которая может обеспечить компромисс между экономией памяти и издержками на коммуникации. См. ShardingStrategy для подробностей. (По умолчанию: FULL_SHARD)
• cpu_offload (Необязательно[CPUOffload]) – Это настраивает выгрузку на CPU. Если это установлено на None, то выгрузка на CPU не происходит. См. CPUOffload для подробностей. (По умолчанию: None)

• auto_wrap_policy (Необязательно[Union[Callable[[nn.Module, bool, int], bool], ModuleWrapPolicy]]) –

  Это определяет политику применения 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 (Необязательно[BackwardPrefetch]) – Это настраивает явную предварительную выборку обратного прохода all-gathers. Если None, то FSDP не выполняет предварительную выборку обратного прохода, и нет перекрытия коммуникации и вычислений в обратном проходе. См. BackwardPrefetch для подробностей. (По умолчанию: BACKWARD_PRE)
• mixed_precision (Необязательно[MixedPrecision]) – Это настраивает нативное смешанное точность для FSDP. Если это установлено на None, то смешанная точность не используется. В противном случае можно задать типы данных для параметров, буферов и уменьшения градиентов. См. MixedPrecision для подробностей. (По умолчанию: None)
• ignored_modules (Необязательно[Iterable[torch.nn.Module]]) – Модули, чьи собственные параметры и параметры и буферы дочерних модулей игнорируются этим экземпляром. Ни один из модулей напрямую в ignored_modules не должен быть экземпляром FullyShardedDataParallel, и любые дочерние модули, которые уже являются построенными экземплярами FullyShardedDataParallel экземпляров, не будут игнорироваться, если они вложены в этот экземпляр. Этот аргумент может быть использован для избежания фрагментации определённых параметров на уровне модуля при использовании auto_wrap_policy или если фрагментация параметров не управляется FSDP. (По умолчанию: None)

• param_init_fn (Необязательно[Callable[[nn.Module], None]]) –

  Функция, которая определяет, как модули, которые в данный момент находятся на мета-устройстве, должны быть инициализированы на фактическом устройстве. С версии 1.12, FSDP обнаруживает модули с параметрами или буферами на мета-устройстве через is_meta и либо применяет param_init_fn при указании, либо вызывает nn.Module.reset_parameters() в противном случае. В обоих случаях реализация должна только инициализировать параметры/буферы модуля, а не подмодулей. Это делается для избежания повторной инициализации. Кроме того, FSDP также поддерживает отложенную инициализацию через API torchdistX (https://github.com/pytorch/torchdistX), где отложенные модули инициализируются вызовом param_init_fn при указании или по умолчанию в torchdistX materialize_module(). Если param_init_fn указано, то оно применяется ко всем модулям на мета-устройстве, что означает, что оно, вероятно, должно работать с типом модуля. 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 (Необязательно[Union[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, соответственно. В соответствии с текущим алгоритмом, разделенная форма всегда 1D, теряя исходную структуру тензора. Исходный параметр может содержать все, некоторые или ни одного своих данных для данного ранга. В случае отсутствия данных его данные будут похожи на пустой тензор размера 0. Пользователи не должны создавать программы, полагающиеся на то, какие данные присутствуют для данного исходного параметра в его разделенной форме. Необходимо установить True для использования torch.compile(). Установка этого значения в False предоставляет пользователю доступ к внутренним FlatParameter FSDP через 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.

apply(fn) [source]

Применяет fn рекурсивно к каждому подмодулю (как возвращает .children()) а также к self. Типичное использование включает инициализацию параметров модели (см. также torch.nn.init).

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

Parameters

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

Returns

self

Return type

Module

clip_grad_norm_(max_norm, norm_type=2.0) [source]

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

Parameters

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

Returns

Общая норма параметров (представленных как один вектор).

Return type

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) [source]

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

Parameters

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

Returns

См. shard_full_optim_state_dict().

Return type

Dict[str, Any]

forward(*args, **kwargs) [source]

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

Return type

Any

static fsdp_modules(module, root_only=False) [source]

Возвращает все вложенные экземпляры 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) [source]

Объединяет полное состояние оптимизатора на ранге 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 представляющие либо список групп параметров, либо итерируемый объект параметров; если 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) [source]

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

Возвращает

Словарь StateDictSettings, содержащий тип state_dict и конфигурации state_dict / optim_state_dict, которые в настоящее время установлены.

Исключения

• AssertionError` если StateDictSettings для разных –
• подмодулей FSDP отличаются. –

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

StateDictSettings

property module: Module

Возвращает обернутый модуль (как DistributedDataParallel).

named_buffers(*args, **kwargs) [source]

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

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

Iterator[Tuple[str, Tensor]]

named_parameters(*args, **kwargs) [source]

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

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

Iterator[Tuple[str, Parameter]]

no_sync() [source]

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

Примечание

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

Примечание

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

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

Генератор

static optim_state_dict(model, optim, optim_state_dict=None, group=None) [source]

Преобразует словарь состояния оптимизатора optim для model, которые фрагментированы FSDP, в один из трёх типов: 1) полный словарь состояния оптимизатора, 2) фрагментированный словарь состояния оптимизатора, 3) локальный словарь состояния оптимизатора.

Для полного словаря состояния оптимизатора все состояния распаковываются и не фрагментируются. Только ранг 0 и только ЦП могут быть указаны через state_dict_type() для избежания OOM.

Для фрагментированного словаря состояния оптимизатора все состояния распаковываются, но фрагментируются. Только ЦП может быть указан через state_dict_type() для дальнейшей экономии памяти.

Для локального словаря состояния преобразования не будут выполняться. Но состояние будет преобразовано из 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(
>>>     optim_state_dict, model, optim
>>> )
>>> optim.load_state_dict(optim_state_dict)

Параметры

• model (torch.nn.Module) – Корневой модуль (который может быть или не быть экземпляром FullyShardedDataParallel), параметры которого были переданы в оптимизатор optim.
• optim (torch.optim.Optimizer) – Оптимизатор для параметров model.
• optim_state_dict (Dict[str, Any]) – целевой словарь состояния оптимизатора для преобразования. Если значение равно 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) [source]

Принимая optim_state_dict, преобразованный через optim_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(
>>>     optim_state_dict, model, optim
>>> )
>>> 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) [source]

Регистрирует коммуникационный хук, который предоставляет гибкий хук пользователям, где они могут указать, как 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) [source]

Переключает словарь состояния оптимизатора 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) [source]

Рассылает полный словарь состояния оптимизатора с ранга 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() могут использоваться для получения словаря состояния раздробленного оптимизатора для загрузки. Предполагая, что полный словарь состояния оптимизатора находится в оперативной памяти ЦП, первый требует, чтобы каждый ранг имел полный словарь в оперативной памяти ЦП, где каждый ранг индивидуально дробит словарь без какой-либо связи, в то время как второй требует, чтобы только ранг 0 имел полный словарь в оперативной памяти ЦП, где ранг 0 перемещает каждый фрагмент в оперативную память видеокарты (для NCCL) и передает его соответствующим рангам. Таким образом, первый имеет более высокую общую стоимость оперативной памяти ЦП, в то время как второй имеет более высокую стоимость связи.

Параметры

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

Устанавливает тип 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 (Необязательно[StateDictConfig]) – конфигурация для целевого state_dict_type.

Возвращает

StateDictSettings, содержащие предыдущий тип словаря состояния и конфигурацию для модуля.

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

StateDictSettings

static shard_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None) [source]

Разделяет полный словарь состояния оптимизатора 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() могут быть использованы для получения разделенного словаря состояния оптимизатора для загрузки. Предполагая, что полный словарь состояния оптимизатора находится в оперативной памяти процессора, первый метод требует, чтобы каждый ранг имел полный словарь в оперативной памяти процессора, где каждый ранг индивидуально разделяет словарь без каких-либо коммуникаций, а второй метод требует, чтобы только ранг 0 имел полный словарь в оперативной памяти процессора, где ранг 0 перемещает каждый фрагмент в оперативную память графического процессора (для NCCL) и передает его соответствующим рангам. Следовательно, первый метод имеет более высокую общую стоимость оперативной памяти процессора, а второй метод имеет более высокую стоимость коммуникаций.

Параметры

• 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]]]) – Входной параметр для оптимизатора, представляющий собой либо список групп параметров, либо итерируемый список параметров; если None, то этот метод предполагает, что вход был model.parameters(). Этот аргумент устарел, и нет необходимости передавать его. (По умолчанию: None)
• optim (Optional[torch.optim.Optimizer]) – Оптимизатор, который загрузит словарь состояния, возвращаемый этим методом. Это предпочтительный аргумент по сравнению с optim_input. (По умолчанию: None)

Возвращает

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

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

Dict[str, Any]

static sharded_optim_state_dict(model, optim, group=None) [source]

API похожа на full_optim_state_dict(), но эта API разбивает все состояния с размерностью, отличной от нуля, на 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) [source]

Контекстный менеджер для установки 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) [source]

Менеджер контекста для экспонирования полных параметров экземпляров 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 приведут к тому, что полные параметры будут избыточно копироваться в оперативную память ЦП для графических процессоров, которые находятся на одной машине, что может привести к риску переполнения памяти ЦП. Рекомендуется использовать offload_to_cpu с rank0_only=True.

Параметры

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

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

Генератор

class torch.distributed.fsdp.BackwardPrefetch(value) [source]

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

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

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

class torch.distributed.fsdp.ShardingStrategy(value) [source]

Это определяет стратегию фрагментации, которая будет использоваться для распределенного обучения с помощью 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'>, )) [source]

Это настраивает обучение с плавающей точностью, встроенное в FSDP.

Переменные

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

Примечание

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

Примечание

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

Примечание

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

Примечание

Нормализация слоев (layer norm) и батч-нормализация (batch norm) накапливают значения в 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) [source]

Это настраивает перенос вычислений на ЦП.

Переменные

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

class torch.distributed.fsdp.StateDictConfig(offload_to_cpu=False, use_dtensor=False) [source]

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

Переменные

• offload_to_cpu (bool) – Если True, то FSDP переносит значения словаря состояния на ЦП, а если False, то FSDP сохраняет их на видеокарте. (По умолчанию: False)
• use_dtensor (bool) – Если True, то FSDP сохраняет значения словаря состояния как DTensor если значение фрагментировано, а если False, то FSDP сохраняет их как ShardedTensor. (По умолчанию: False)

class torch.distributed.fsdp.FullStateDictConfig(offload_to_cpu=False, use_dtensor=False, rank0_only=False) [source]

FullStateDictConfig — это класс конфигурации, предназначенный для использования с StateDictType.FULL_STATE_DICT. Мы рекомендуем включить как offload_to_cpu=True, так и rank0_only=True, когда сохраняются полные словари состояния, чтобы сэкономить память на видеокарте и памяти на ЦП соответственно. Этот класс конфигурации предназначен для использования через менеджер контекста 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 on CPU 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: bool = False, use_dtensor: bool = False) [source]

class torch.distributed.fsdp.LocalStateDictConfig(offload_to_cpu: bool = False, use_dtensor: bool = False) [source]

class torch.distributed.fsdp.OptimStateDictConfig(offload_to_cpu=True, use_dtensor=False) [source]

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

Переменные

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

class torch.distributed.fsdp.FullOptimStateDictConfig(offload_to_cpu=True, use_dtensor=False, rank0_only=False) [source]

Переменные

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

class torch.distributed.fsdp.ShardedOptimStateDictConfig(offload_to_cpu: bool = True, use_dtensor: bool = False) [source]

class torch.distributed.fsdp.LocalOptimStateDictConfig(offload_to_cpu: bool = False, use_dtensor: bool = False) [source]

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) [source]