Параметры

• checkpoint_dir (str) — Каталог, содержащий контрольные точки FSDP (может быть как для модели, так и для оптимизатора).
• output_path (str) — Путь для сохранения объединённой контрольной точки.
• safe_serialization (bool, необязательный, по умолчанию True) — Сохранять ли объединённые веса в формате safetensors (рекомендуется).
• remove_checkpoint_dir (bool, необязательный, по умолчанию False) — Удалять ли каталог контрольной точки после слияния.

Параметры

• fsdp_version (int, по умолчанию 1) — Версия FSDP, которую следует использовать. По умолчанию — 1. Если установлено значение 2, загрузчик ожидает, что конфигурация будет преобразована в формат FSDP2.
• sharding_strategy (Union[str, torch.distributed.fsdp.ShardingStrategy], по умолчанию 'FULL_SHARD') — Стратегия шардирования. Должна быть либо str, либо экземпляром torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy. Устарела, рекомендуется использовать reshard_after_forward.
• reshard_after_forward (Union[str, torch.distributed.fsdp.ShardingStrategy, bool], по умолчанию 'FULL_SHARD' для fsdp_version=1 и True для fsdp_version=2) — Стратегия шардирования. Должна быть bool, если fsdp_version установлено в 2, иначе str или экземпляр torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy.
• backward_prefetch (Union[str, torch.distributed.fsdp.BackwardPrefetch], по умолчанию 'NO_PREFETCH') — Стратегия предварительной загрузки в обратном направлении. Должно быть строкой str или экземпляром torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch.
• mixed_precision_policy (Optional[Union[dict, str, torch.distributed.fsdp.MixedPrecision, torch.distributed.fsdp.MixedPrecisionPolicy]], по умолчанию None) — Конфигурация для включения обучения с смешанной точностью с использованием FullyShardedDataParallel. Если передается dict, он должен содержать следующие ключи: param_dtype, reduce_dtype и buffer_dtype. Может быть экземпляром torch.distributed.fsdp.MixedPrecisionPolicy, если fsdp_version установлено в 2. Если передается str, он должен быть одним из следующих значений: fp8, fp16, bf16, fp32, и используется для установки param_dtype, reduce_dtype и buffer_dtype.
• auto_wrap_policy (Optional(Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]), defaults to NO_WRAP) -- A callable or string specifying a policy to recursively wrap layers with FSDP. If a string, it must be one of transformer_based_wrap, size_based_wrap, or no_wrap. See torch.distributed.fsdp.wrap.size_based_wrap_policy` for a direction on what it should look like.
• cpu_offload (Union[bool, torch.distributed.fsdp.CPUOffload, torch.distributed.fsdp.CPUOffloadPolicy], по умолчанию False) — Следует ли выгружать параметры на CPU. Значением может быть bool или экземпляр torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload или torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy, если fsdp_version установлено в 2.
• ignored_modules (Optional[Union[Iterable[torch.nn.Module], str]], по умолчанию None) — Список модулей, которые следует игнорировать при обёртке в FSDP. Если передана строка, модули будут сопоставляться по имени с использованием полного совпадения регулярного выражения. Если fsdp_version установлен в 2, модули преобразуются в параметры и используются.
• state_dict_type (Union[str, torch.distributed.fsdp.StateDictType], по умолчанию 'FULL_STATE_DICT') — Тип state dict для использования. Если передана строка, она должна быть одной из full_state_dict, local_state_dict или sharded_state_dict.
• state_dict_config (Optional[Union[torch.distributed.fsdp.FullStateDictConfig, torch.distributed.fsdp.ShardedStateDictConfig], по умолчанию None) — Конфигурация state dict, которую следует использовать. Если не передана, определяется на основе state_dict_type.
• optim_state_dict_config (Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig], по умолчанию None) — Конфигурация state dict для оптимизатора. Если не передана, определяется на основе state_dict_type.
• limit_all_gathers (bool, по умолчанию True) — Определяет, следует ли FSDP явно синхронизировать CPU-поток, чтобы предотвратить слишком большое количество незавершённых all-gather операций. Этот параметр bool влияет только на стратегии с шардированием, которые планируют all-gather операции. Включение этого параметра может помочь снизить количество повторных попыток выделения памяти CUDA.
• use_orig_params (bool, по умолчанию False) — Использовать ли исходные параметры для оптимизатора.
• param_init_fn (Optional[Callable[[torch.nn.Module], None], по умолчанию None) — Callable[torch.nn.Module] -> None, определяющая, как модули, находящиеся на meta-устройстве, должны быть инициализированы на реальном устройстве. Применимо только тогда, когда sync_module_states равно True. По умолчанию это lambda, вызывающая to_empty для модуля.
• sync_module_states (bool, по умолчанию False) — Определяет, следует ли каждому индивидуально обернутому модулю FSDP транслировать параметры модуля с ранга 0, чтобы гарантировать их идентичность на всех рангах после инициализации. По умолчанию False, если только cpu_ram_efficient_loading не равно True, в этом случае параметр принудительно включается.
• forward_prefetch (bool, по умолчанию False) — Определяет, следует ли FSDP явно выполнять предварительную загрузку (prefetch) следующего all-gather во время выполнения прямого прохода (forward pass). Используйте только со статическими графами.
• activation_checkpointing (bool, по умолчанию False) — Техника для снижения потребления памяти путём очистки активаций определённых слоёв и их повторного вычисления во время обратного прохода. По сути, это дополнительное время вычислений в обмен на уменьшенное потребление памяти.
• cpu_ram_efficient_loading (bool, по умолчанию None) — Если True, предобученную контрольную точку модели загружает только первый процесс, а все остальные процессы имеют пустые веса. Применимо только для Transformers. При использовании этого параметра sync_module_states должно быть True.
• transformer_cls_names_to_wrap (Optional[List[str]], по умолчанию None) — Список имён классов слоёв трансформера, которые нужно обернуть. Применимо только тогда, когда auto_wrap_policy равно transformer_based_wrap.
• min_num_params (Optional[int], по умолчанию None) — Минимальное количество параметров, которое должен иметь модуль для его обёртывания. Применимо только тогда, когда auto_wrap_policy равно size_based_wrap.

Параметры

• accelerator (Accelerator) — экземпляр акселератора
• model (torch.nn.Module) — Модель, в которую будет загружен state dict. Ожидается, что она находится на meta device, иначе может произойти скачок потребления VRAM
• full_sd (dict) — Полный state dict для загрузки, может находиться только на rank 0
• cpu_offload (bool, по умолчанию False) — Если True, перемещает разбитые параметры на CPU после распределения. Обязательно, если включено выделение памяти FSDP на CPU.

Параметры

• optimizer (torch.optim.Optimizer) — экземпляр оптимизатора, содержащий исходные параметры модели
• mapping (dict) — Соответствие между исходным параметром (указанным через data_ptr) и шардированным параметром

Параметры

• accelerator (Accelerator) — экземпляр акселератора
• model (torch.nn.Module) — Модель для подготовки