FullyShardedDataParallel

class torch.distributed.fsdp.FullyShardedDataParallel(module, process_group=None, sharding_strategy=None, cpu_offload=None, auto_wrap_policy=None, backward_prefetch=None, mixed_precision=None, ignored_modules=None, param_init_fn=None, device_id=None, sync_module_states=False, forward_prefetch=False, limit_all_gathers=False) [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 фрагментирует параметры на месте, и это нарушит работу ранее инициализированных оптимизаторов.

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

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

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

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

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

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

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

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

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

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

Примечание

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

Параметры:

• module (nn.Module) – модуль, который будет обернут с помощью FSDP.
• process_group (Необязательно[ProcessGroup]) – группа процессов для фрагментации.
• sharding_strategy (Необязательно[ShardingStrategy]) – конфигурация алгоритма фрагментации, различные алгоритмы фрагментации имеют компромисс между экономией памяти и задержкой обмена. FULL_SHARD будет выбран, если sharding_strategy не указан.
• cpu_offload (Необязательно[CPUOffload]) – конфигурация выгрузки на CPU. В настоящее время поддерживается только выгрузка параметров и градиентов на CPU. Ее можно включить, передав cpu_offload=CPUOffload(offload_params=True). Обратите внимание, что в настоящее время это неявным образом включает выгрузку градиентов на CPU, чтобы параметры и градиенты находились на одном устройстве для работы с оптимизатором. Этот API может быть изменен. По умолчанию None , в этом случае выгрузка не будет выполнена.

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

  Функция, определяющая политику рекурсивного обертывания слоев с помощью FSDP. Обратите внимание, что эта политика в настоящее время будет применяться только к дочерним модулям переданного модуля. Остальные модули всегда обертываются в возвращенном корневом экземпляре FSDP. size_based_auto_wrap_policy , написанная на torch.distributed.fsdp.wrap, является примером auto_wrap_policy функции, эта политика оборачивает слои с количеством параметров больше 100 млн. transformer_auto_wrap_policy , написанная на torch.distributed.fsdp.wrap, является примером auto_wrap_policy функции для архитектур моделей типа трансформеров. Пользователи могут предоставить настраиваемую auto_wrap_policy функцию, которая должна принимать следующие аргументы: module: nn.Module, recurse: bool, unwrapped_params: int, и возвращать bool , указывающую, должен ли переданный module` быть обернут (если recurse=False) или следует ли рекурсивно переходить вниз по подграфу дочерних элементов module (если recurse=True). Дополнительные настраиваемые аргументы также могут быть добавлены к настраиваемой auto_wrap_policy функции. Хорошей практикой является вывод отформатированной модели и проверка того, соответствует ли отформатированная модель ожиданиям приложения, а затем соответствующая корректировка.

  Пример:

  >>> def custom_auto_wrap_policy(
  >>>     module: nn.Module,
  >>>     recurse: bool,
  >>>     unwrapped_params: int,
  >>>     # These are customizable for this policy function.
  >>>     min_num_params: int = int(1e8),
  >>> ) -> bool:
  >>>     return unwrapped_params >= min_num_params
  >>> # Configure a custom min_num_params
  >>> my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=1e5)
  

• backward_prefetch (Необязательно[BackwardPrefetch]) – Это экспериментальная функция, которая может быть изменена в ближайшем будущем. Она позволяет пользователям включить два различных алгоритма backward_prefetch, чтобы помочь перекрытию обратного обмена и вычислений. Плюсы и минусы каждого алгоритма описаны в классе BackwardPrefetch.
• mixed_precision (Необязательно[MixedPrecision]) – экземпляр MixedPrecision , описывающий конфигурацию обучения с переменной точностью, которая будет использоваться. MixedPrecision поддерживает настройку типа данных для обмена параметрами, буферами и градиентами. Обратите внимание, что только данные с плавающей точкой преобразуются в сниженную точность. Это позволяет пользователям потенциально экономить память и ускорять обучение, жертвуя точностью во время обучения модели. Если None, то применение переменной точности не производится. Обратите внимание, что если mixed_precision включено для модели FSDP, которая содержит BatchNorm с auto_wrap_policy, FSDP позаботится о выключении переменной точности для BatchNorm блоков, обернув их отдельно в свой собственный блок FSDP с mixed_precision=None . Это делается потому, что некоторые BatchNorm ядра в настоящее время не реализуют поддержку сниженных типов. Если модуль обертывается индивидуально, пользователи должны позаботиться о настройке mixed_precision=None для BatchNorm блоков. (По умолчанию: None)
• ignored_modules (Необязательно[Iterable[torch.nn.Module]]) – модули, чьи собственные параметры и параметры дочерних модулей и буферы игнорируются этим экземпляром. Ни один из модулей непосредственно в ignored_modules не должен быть экземплярами FullyShardedDataParallel, и любые дочерние модули, которые уже являются построенными экземплярами FullyShardedDataParallel экземплярами, не будут игнорироваться, если они вложены в этот экземпляр. Этот аргумент может быть использован для того, чтобы избежать фрагментации определенных параметров на уровне модуля при использовании auto_wrap_policy или если фрагментация параметров не управляется FSDP. (По умолчанию: None)

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

  Функция, определяющая, как модули, которые в настоящее время находятся на мета-устройстве, должны быть инициализированы на фактическом устройстве. Обратите внимание, что начиная с версии v1.12, мы обнаруживаем модули на мета-устройстве с помощью is_meta проверки и применяем стандартную инициализацию, которая вызывает reset_parameters метод на переданном nn.Module , если param_init_fn не указан, в противном случае мы выполняем param_init_fn для инициализации переданного nn.Module. В частности, это означает, что если is_meta=True для любых параметров модуля для модулей, которые будут обернуты с помощью FSDP, и param_init_fn не указан, мы предполагаем, что ваш модуль правильно реализует reset_paramters() и вызовет ошибки, если нет. Обратите внимание, что дополнительно мы поддерживаем модули, инициализированные с помощью API torchdistX (https://github.com/pytorch/torchdistX). В этом случае отложенные модули будут инициализированы с помощью функции инициализации по умолчанию, которая вызывает materialize_module torchdistX, или переданную param_init_fn, если она не None. Такая же Callable применяется для инициализации всех мета-модулей. Обратите внимание, что эта функция инициализации применяется до выполнения любой логики фрагментации FSDP.

  Пример:

  >>> module = MyModule(device="meta")
  >>> def my_init_fn(module):
  >>>     # responsible for initializing a module, such as with reset_parameters
  >>>     ...
  >>> 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 или объект torch.device, описывающий CUDA устройство, на которое должен быть перемещен модуль FSDP, определяя, где происходит инициализация, такая как фрагментация. Если этот аргумент не указан, и module находится на CPU, мы выводим предупреждение, в котором упоминается, что этот аргумент может быть указан для более быстрого запуска. Если указано, результирующие экземпляры FSDP будут находиться на этом устройстве, включая перемещение параметров игнорируемых модулей при необходимости. Обратите внимание, что если device_id указан, но module уже находится на другом CUDA устройстве, будет выброшено исключение. (По умолчанию: None)
• sync_module_states (bool) – Если True, каждый индивидуально обернутый блок FSDP будет транслировать параметры модулей с ранга 0, чтобы гарантировать их одинаковость на всех рангах после инициализации. Это помогает гарантировать, что параметры модели одинаковы на всех рангах перед началом обучения, но добавляет накладные расходы на обмен в __init__, так как по меньшей мере одна трансляция вызывается для каждого индивидуально обернутого блока FSDP. Это также может помочь загрузить контрольные точки, взятые state_dict и загруженные load_state_dict с эффективным использованием памяти. См. документацию для FullStateDictConfig для примера этого. (По умолчанию: False)
• forward_prefetch (bool) – Если True, тогда FSDP *явно* преобразует следующий предстоящий all-gather во время выполнения прямого прохода. Это может улучшить перекрытие обмена и вычислений для задач с интенсивной нагрузкой на CPU. Это следует использовать только для моделей со статической схемой, так как порядок прямого прохода фиксируется на основе выполнения первой итерации. (По умолчанию: False)
• limit_all_gathers (bool) – Если False, то FSDP позволяет потоку CPU планировать all-gather без дополнительной синхронизации. Если True, тогда FSDP явно синхронизирует поток CPU, чтобы предотвратить слишком много текущих all-gather. Это bool влияет только на стратегии фрагментации, которые планируют all-gather. Включение этого может помочь снизить количество попыток malloc CUDA.

apply(fn) [source]

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

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

Параметры:

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

Возвращаемое значение:

self

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

Module

clip_grad_norm_(max_norm, norm_type=2.0) [source]

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

Параметры:

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

Возвращаемое значение:

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

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

None

Примечание

Это аналогично torch.nn.utils.clip_grad_norm_, но обрабатывает разбиение и несколько устройств на ранг под капотом. Стандартная утилита torch здесь неприменима, потому что каждый ранг имеет только частичный вид всех градиентов в модели, поэтому её вызов для FSDP-моделей приведет к тому, что к каждому подмножеству параметров модели будет применено разное масштабирование.

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

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

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

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) – См. :meth:shard_full_optim_state_dict.

Возвращаемое значение:

См. shard_full_optim_state_dict().

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

Dict[str, Any]

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

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

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

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]

load_state_dict(state_dict, *args, **kwargs) [source]

Точка входа для всех трех API FSDP load_state_dict. По умолчанию, вызов load_state_dict для модуля FSDP приведет к попытке FSDP загрузить «полный» словарь состояния, т.е. словарь состояния, содержащий полные, нефрагментированные, несглаженные исходные параметры модуля. Это требует от FSDP загрузки полного контекста параметров на каждом ранге, что может привести к ошибке OOM на GPU. В результате доступен API state_dict_type() для настройки между реализациями load_state_dict. Таким образом, пользователь может использовать контекстный менеджер with self.state_dict_type(self, StateDictType.LOCAL_STATE_DICT) для загрузки локального контрольного пункта словаря состояния, который восстановит только локальные фрагменты модуля. В настоящее время поддерживаются только реализации StateDictType.LOCAL_STATE_DICT и StateDictType.FULL_STATE_DICT (по умолчанию). См. state_dict() для документации по созданию контрольного пункта FSDP.

Пример:

>>> import torch
>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> from torch.distributed.fsdp import StateDictType
>>> torch.cuda.set_device(device_id)
>>> my_module = nn.Linear(...)
>>> sharded_module = FSDP(my_module)
>>> checkpoint = torch.load(PATH)
>>> full_state_dict = checkpoint['full_state_dict']
>>> with FSDP.state_dict_type(sharded_module, StateDictType.FULL_STATE_DICT):
>>>     sharded_module.load_state_dict(full_state_dict)
>>> full_dict.keys()
>>> odict_keys(['weight', 'bias'])
>>> # using local state dict
>>> local_state_dict = checkpoint['local_state_dict']
>>> with FSDP.state_dict_type(sharded_module, StateDictType.LOCAL_STATE_DICT):
>>>     sharded_module.load_state_dict(local_state_dict)
>>> local_dict.keys()
>>> odict_keys(['flat_param', 'inner.flat_param'])

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

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

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

NamedTuple

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 будет накапливать полные градиенты модели (вместо фрагментов градиента) до момента синхронизации.

Примечание

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

Возвращаемый тип:

Generator

property params_with_grad: List[Parameter]

Рекурсивно возвращает список всех параметров модуля, у которых есть градиент.

register_comm_hook(state, hook) [source]

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

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

Хук связи FSDP должен быть зарегистрирован до запуска начального прохода вперёд и только один раз.

Параметры:

• state (объект) –

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

Параметры:

• full_optim_state_dict (Optional[Dict[str, Any]]) – Словарь состояния оптимизатора, соответствующий нераскладываемым параметрам и содержащий полный нефрагментированный словарь состояния оптимизатора, если на ранге 0; аргумент игнорируется на рангах отличных от 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 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, optim_input=None, 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]

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

Это точка входа для всех трех FSDP state_dict API: полной, локальной и фрагментированной. Для полного словаря состояния (StateDictType.FULL_STATE_DICT) FSDP пытается дефрагментировать модель на всех рангах, что может привести к ошибке OOM, если полная модель не поместится на одной видеокарте. В этом случае пользователи могут передать FullStateDictConfig для сохранения контрольной точки только на ранге 0 и/или для выгрузки ее в оперативную память процессора послойно, что позволяет создавать гораздо большие контрольные точки. Если полная модель не помещается в оперативной памяти процессора, пользователи вместо этого могут взять локальный словарь состояния (StateDictType.LOCAL_STATE_DICT), который сохраняет только локальный фрагмент модели. Фрагментированный словарь состояния (StateDictType.SHARDED_STATE_DICT) сохраняет параметры модели в виде ShardedTensor. Тип state_dict может быть настроен с помощью контекстного менеджера state_dict_type().

Пример:

>>> import torch
>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> from torch.distributed.fsdp import StateDictType
>>> torch.cuda.set_device(device_id)
>>> my_module = nn.Linear(...)
>>> sharded_module = FSDP(my_module)
>>> full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
>>> with FSDP.state_dict_type(sharded_module, StateDictType.FULL_STATE_DICT, full_state_dict_config):
>>>     full_dict = sharded_module.state_dict()
>>> full_dict.keys()
>>> odict_keys(['weight', 'bias'])
>>> # using local state dict
>>> with FSDP.state_dict_type(sharded_module, StateDictType.LOCAL_STATE_DICT):
>>>     local_dict = sharded_module.state_dict()
>>> local_dict.keys()
>>> odict_keys(['flat_param', 'inner.flat_param'])

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

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

static state_dict_type(module, state_dict_type, state_dict_config=None) [source]

Контекстный менеджер для установки state_dict_type всех дочерних модулей FSDP целевого модуля. Целевой модуль не обязательно должен быть модулем FSDP. Если целевой модуль является модулем FSDP, его state_dict_type также будет изменен.

Примечание

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

Примечание

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

Пример:

>>> model = DDP(FSDP(...))
>>> with FSDP.state_dict_type(model, StateDictType.LOCAL_STATE_DICT):
>>>     checkpoint = model.state_dict()

Параметры:

• module (torch.nn.Module) – Корневой модуль.
• state_dict_type (StateDictType) – требуемый state_dict_type для установки.

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

Generator

static summon_full_params(module, recurse=True, writeback=True, rank0_only=False, offload_to_cpu=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=True приведут к тому, что полные параметры будут избыточно скопированы в оперативную память ЦП для графических процессоров, расположенных на одной машине, что может привести к риску исчерпания оперативной памяти ЦП. Рекомендуется использовать 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, полные параметры загружаются в ЦП. Обратите внимание, что эта загрузка в настоящее время происходит только в том случае, если параметр фрагментирован (что не так только для world_size = 1 или NO_SHARD конфигурации). Рекомендуется использовать offload_to_cpu с rank0_only=True для избежания избыточных копий параметров модели, загружаемых в ту же оперативную память ЦП.

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

Generator