Hugging Face's logo

Accelerate documentation

Работа с большими моделями

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.15.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Работа с большими моделями

Диспетчеризация и выгрузка

init_empty_weights

accelerate.init_empty_weights

< >

( include_buffers: bool | None = None )

Параметры

  • include_buffers (bool, необязательный) — Следует ли также размещать все буферы на мета-устройстве при инициализации.

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

Пример:

import torch.nn as nn
from accelerate import init_empty_weights

# Initialize a model with 100 billions parameters in no time and without using any RAM.
with init_empty_weights():
    tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

Any model created under this context manager has no weights. As such you can’t do something like model.to(some_device) with it. To load weights inside your empty model, see load_checkpoint_and_dispatch(). Make sure to overwrite the default device_map param for load_checkpoint_and_dispatch(), otherwise dispatch is not called.

cpu_offload

accelerate.cpu_offload

< >

( model: Moduleexecution_device: torch.device | None = Noneoffload_buffers: bool = Falsestate_dict: dict[str, torch.Tensor] | None = Nonepreload_module_classes: list[str] | None = None )

Параметры

  • model (torch.nn.Module) — Модель, которую нужно выгрузить.
  • execution_device (torch.device, необязательный) — Устройство, на котором будет выполняться прямой проход модели (должно быть GPU). По умолчанию используется устройство первого параметра модели.
  • offload_buffers (bool, необязательный, по умолчанию False) — Определяет, следует ли выгружать буферы вместе с параметрами модели.
  • state_dict (Dict[str, torch.Tensor], необязательный) — Словарь состояния модели, который будет храниться на CPU.
  • preload_module_classes (List[str], необязательный) — Список классов, экземпляры которых должны загружать все свои веса (в том числе в подмодулях) в начале forward. Это следует использовать только для классов, у которых есть подмодули, зарегистрированные, но не вызываемые напрямую во время forward, например, если зарегистрирован dense линейный слой, но во время forward dense.weight и dense.bias используются в некоторых операциях вместо прямого вызова dense.

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

cpu_offload_with_hook

accelerate.cpu_offload_with_hook

< >

( model: Moduleexecution_device: int | str | torch.device | None = Noneprev_module_hook: accelerate.hooks.UserCpuOffloadHook | None = None )

Параметры

  • model (torch.nn.Module) — Модель, которую нужно выгрузить.
  • execution_device(str, int or torch.device, optional) — Устройство, на котором должна выполняться модель. По умолчанию используется устройство MPS, если оно доступно, затем устройство 0, если есть ускоритель, и в противном случае — CPU.
  • prev_module_hook (UserCpuOffloadHook, необязательный) — Хук, возвращаемый этой функцией для предыдущей модели в выполняемом вами конвейере. Если передан, его метод offload будет вызван непосредственно перед forward модели, к которой привязан данный хук.

Переносит модель на CPU и возвращает её на устройство выполнения при запуске. Отличие от cpu_offload() состоит в том, что после forward модель остаётся на устройстве выполнения и переносится обратно только при вызове метода offload возвращённого hook. Полезно для конвейеров, которые запускают модель в цикле.

Пример:

model_1, hook_1 = cpu_offload_with_hook(model_1, device)
model_2, hook_2 = cpu_offload_with_hook(model_2, device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()

disk_offload

accelerate.disk_offload

< >

( model: Moduleoffload_dir: str | os.PathLikeexecution_device: torch.device | None = Noneoffload_buffers: bool = Falsepreload_module_classes: list[str] | None = None )

Параметры

  • model (torch.nn.Module) — Модель для выгрузки.
  • offload_dir (str или os.PathLike) — Папка, в которую выгружаются веса модели (или где они уже выгружены).
  • execution_device (torch.device, необязательный) — Устройство, на котором будет выполняться прямой проход модели (должно быть GPU). По умолчанию используется устройство первого параметра модели.
  • offload_buffers (bool, необязательный, по умолчанию False) — Определяет, следует ли выгружать буферы вместе с параметрами модели.
  • preload_module_classes (List[str], необязательный) — Список классов, экземпляры которых должны загружать все свои веса (в том числе в подмодулях) в начале forward. Это следует использовать только для классов, у которых есть подмодули, зарегистрированные, но не вызываемые напрямую во время forward, например, если зарегистрирован dense линейный слой, но во время forward dense.weight и dense.bias используются в некоторых операциях вместо прямого вызова dense.

Включает полный offload модели на диск. В результате все параметры модели будут offload-ированы в виде memory-mapped массива в указанную папку. Во время forward pass параметры будут извлекаться из этой папки и по мере необходимости перемещаться на устройство выполнения, переданное в качестве аргумента, после чего снова offload-ироваться.

dispatch_model

accelerate.dispatch_model

< >

( model: Moduledevice_map: dictmain_device: torch.device | None = Nonestate_dict: dict[str, torch.Tensor] | None = Noneoffload_dir: str | os.PathLike | None = Noneoffload_index: dict[str, str] | None = Noneoffload_buffers: bool = Falseskip_keys: str | list[str] | None = Nonepreload_module_classes: list[str] | None = Noneforce_hooks: bool = False )

Параметры

  • model (torch.nn.Module) — Модель для диспетчеризации.
  • device_map (Dict[str, Union[str, int, torch.device]]) — Словарь, сопоставляющий имена модулей в state_dict моделей с устройствами, на которые они должны быть размещены. Обратите внимание, что значение "disk" принимается, даже если оно не является корректным для torch.device.
  • main_device (str, int или torch.device, необязательный) — Основное устройство выполнения. По умолчанию используется первое устройство в device_map, отличное от "cpu" или "disk".
  • state_dict (Dict[str, torch.Tensor], необязательный) — Словарь состояния той части модели, которая будет храниться на CPU.
  • offload_dir (str или os.PathLike) — Папка, в которую выгружаются веса модели (или где они уже выгружены).
  • offload_index (Dict, необязательный) — Словарь, сопоставляющий имена весов с их информацией (dtype/ shape или имя файла safetensors). По умолчанию используется индекс, сохранённый в save_folder.
  • offload_buffers (bool, необязательный, по умолчанию False) — Определяет, следует ли выгружать буферы вместе с параметрами модели.
  • skip_keys (str или List[str], необязательный) — Список ключей, которые следует игнорировать при перемещении входных или выходных данных между устройствами.
  • preload_module_classes (List[str], необязательный) — Список классов, экземпляры которых должны загружать все свои веса (в том числе в подмодулях) в начале forward. Это следует использовать только для классов, у которых есть подмодули, зарегистрированные, но не вызываемые напрямую во время forward, например, если зарегистрирован dense линейный слой, но во время forward dense.weight и dense.bias используются в некоторых операциях вместо прямого вызова dense.
  • force_hooks (bool, необязательный, по умолчанию False) — Определяет, следует ли принудительно прикреплять device hooks к модели, даже если все слои распределены на одном устройстве.

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

load_checkpoint_and_dispatch

accelerate.load_checkpoint_and_dispatch

< >

( model: Modulecheckpoint: str | os.PathLikedevice_map: str | dict[str, int | str | torch.device] | None = Nonemax_memory: dict[int | str, int | str] | None = Noneno_split_module_classes: list[str] | None = Noneoffload_folder: str | os.PathLike | None = Noneoffload_buffers: bool = Falsedtype: str | torch.dtype | None = Noneoffload_state_dict: bool | None = Noneskip_keys: str | list[str] | None = Nonepreload_module_classes: list[str] | None = Noneforce_hooks: bool = Falsestrict: bool = Falsefull_state_dict: bool = Truebroadcast_from_rank0: bool = False )

Параметры

  • model (torch.nn.Module) — модель, в которую мы хотим загрузить контрольную точку.
  • checkpoint (str или os.PathLike) — Папка контрольной точки для загрузки. Она может быть:
    • путь к файлу, содержащему полный state dict модели
    • путь к файлу .json, содержащему индекс для шардированного контрольного пункта
    • путь к папке, содержащей уникальный файл .index.json и фрагменты контрольной точки.
  • device_map (Dict[str, Union[int, str, torch.device]], необязательный) — Карта, указывающая, на какое устройство следует поместить каждый подмодуль. Не обязательно указывать имя каждого параметра или буфера: как только имя модуля указано, все его подмодули будут отправлены на то же устройство.

    Чтобы Accelerate автоматически вычислял наиболее оптимизированный device_map, установите device_map="auto". Подробнее о каждом варианте см. здесь. По умолчанию значение равно None, что означает, что dispatch_model() не будет вызываться.

  • max_memory (Dict, необязательный) — Словарь, сопоставляющий идентификатор устройства с максимальным объемом памяти. Если не задан, по умолчанию используется максимальная доступная память для каждого GPU и доступная оперативная память CPU.
  • no_split_module_classes (List[str], необязательный) — Список имён классов слоёв, которые никогда не следует разделять между устройствами (например, любой слой, имеющий остаточное соединение).
  • offload_folder (str или os.PathLike, необязательный) — Папка, в которую будут выгружаться веса, если device_map содержит значение "disk".
  • offload_buffers (bool, необязательный, по умолчанию False) — Определяет, следует ли выгружать на CPU или диск буферы в дополнение к параметрам в тех слоях, которые выгружаются на CPU или диск.
  • dtype (str или torch.dtype, необязательный) — Если указано, веса будут преобразованы в этот тип при загрузке.
  • offload_state_dict (bool, необязательный) — Если True, CPU state dict будет временно выгружаться на диск, чтобы избежать нехватки оперативной памяти CPU, если размер CPU state dict и крупнейшего шарда не помещается в неё. По умолчанию устанавливается в True, если в выбранной device map есть значения "disk".
  • skip_keys (str или List[str], необязательный) — Список ключей, которые следует игнорировать при перемещении входных или выходных данных между устройствами.
  • preload_module_classes (List[str], необязательный) — Список классов, экземпляры которых должны загружать все свои веса (в том числе в подмодулях) в начале forward. Это следует использовать только для классов, у которых есть подмодули, зарегистрированные, но не вызываемые напрямую во время forward, например, если зарегистрирован dense линейный слой, но во время forward dense.weight и dense.bias используются в некоторых операциях вместо прямого вызова dense.
  • force_hooks (bool, необязательный, по умолчанию False) — Определяет, следует ли принудительно прикреплять device hooks к модели, даже если все слои распределены на одном устройстве.
  • strict (bool, необязательный, по умолчанию False) — Определяет, следует ли строго требовать, чтобы ключи в state_dict контрольной точки совпадали с ключами state_dict модели.
  • full_state_dict (bool, необязательный, по умолчанию True) — если установлено значение True, все тензоры в загруженном state_dict будут собраны. В загруженном state_dict не будет ShardedTensor и DTensor.
  • broadcast_from_rank0 (False, необязательный, по умолчанию False) — если параметр установлен в True, необходимо инициализировать распределённую группу ProcessGroup. rank0 должен получить полный state_dict и последовательно рассылать тензоры из state_dict по одному на остальные ранги. Остальные ранги получают тензоры и, при необходимости, разбивают их на фрагменты в соответствии с локальными фрагментами в модели.

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

Пример:

>>> from accelerate import init_empty_weights, load_checkpoint_and_dispatch
>>> from huggingface_hub import hf_hub_download
>>> from transformers import AutoConfig, AutoModelForCausalLM

>>> # Download the Weights
>>> checkpoint = "EleutherAI/gpt-j-6B"
>>> weights_location = hf_hub_download(checkpoint, "pytorch_model.bin")

>>> # Create a model and initialize it with empty weights
>>> config = AutoConfig.from_pretrained(checkpoint)
>>> with init_empty_weights():
...     model = AutoModelForCausalLM.from_config(config)

>>> # Load the checkpoint and dispatch it to the right devices
>>> model = load_checkpoint_and_dispatch(
...     model, weights_location, device_map="auto", no_split_module_classes=["GPTJBlock"]
... )

load_checkpoint_in_model

accelerate.load_checkpoint_in_model

< >

( model: Modulecheckpoint: str | os.PathLikedevice_map: dict[str, int | str | torch.device] | None = Noneoffload_folder: str | os.PathLike | None = Nonedtype: str | torch.dtype | None = Noneoffload_state_dict: bool = Falseoffload_buffers: bool = Falsekeep_in_fp32_modules: list[str] | None = Noneoffload_8bit_bnb: bool = Falsestrict: bool = Falsefull_state_dict: bool = Truebroadcast_from_rank0: bool = False )

Параметры

  • model (torch.nn.Module) — Модель, в которую мы хотим загрузить контрольную точку.
  • checkpoint (str или os.PathLike) — Папка контрольной точки для загрузки. Она может быть:
    • путь к файлу, содержащему полный state dict модели
    • путь к файлу .json, содержащему индекс для шардированного контрольного пункта
    • путь к папке, содержащей уникальный файл .index.json и фрагменты контрольной точки.
    • путь к папке, содержащей уникальный файл pytorch_model.bin или model.safetensors.
  • device_map (Dict[str, Union[int, str, torch.device]], необязательный) — Карта, указывающая, на какое устройство следует поместить каждый подмодуль. Не требуется указывать её для каждого параметра/буфера по имени: как только имя модуля указано в карте, все его подмодули будут отправлены на то же устройство.
  • offload_folder (str или os.PathLike, необязательный) — Папка, в которую будут выгружаться веса, если device_map содержит значение "disk".
  • dtype (str или torch.dtype, необязательный) — Если указано, веса будут преобразованы в этот тип при загрузке.
  • offload_state_dict (bool, необязательный, по умолчанию False) — Если True, словарь состояния CPU будет временно выгружен на диск, чтобы избежать нехватки оперативной памяти CPU, если размер словаря состояния CPU и крупнейшего фрагмента не помещается в неё.
  • offload_buffers (bool, необязательный, по умолчанию False) — Включать ли буферы в веса, выгружаемые на диск.
  • keep_in_fp32_modules(List[str], необязательный) — Список модулей, которые мы сохраняем в типе данных torch.float32.
  • offload_8bit_bnb (bool, необязательный) — Включать ли offload 8-битных модулей на cpu/disk.
  • strict (bool, необязательный, по умолчанию False) — Определяет, следует ли строго требовать, чтобы ключи в state_dict контрольной точки совпадали с ключами state_dict модели.
  • full_state_dict (bool, необязательный, по умолчанию True) — если установлено значение True, все тензоры в загруженном state_dict будут собраны. В загруженном state_dict не будет ShardedTensor и DTensor.
  • broadcast_from_rank0 (False, необязательный, по умолчанию False) — если параметр установлен в True, необходимо инициализировать распределённую группу ProcessGroup. rank0 должен получить полный state_dict и последовательно рассылать тензоры из state_dict по одному на остальные ранги. Остальные ранги получают тензоры и, при необходимости, разбивают их на фрагменты в соответствии с локальными фрагментами в модели.

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

Once loaded across devices, you still need to call dispatch_model() on your model to make it able to run. To group the checkpoint loading and dispatch in one single call, use load_checkpoint_and_dispatch().

infer_auto_device_map

accelerate.infer_auto_device_map

< >

( model: Modulemax_memory: dict[int | str, int | str] | None = Noneno_split_module_classes: list[str] | None = Nonedtype: str | torch.dtype | None = Nonespecial_dtypes: dict[str, str | torch.dtype] | None = Noneverbose: bool = Falseclean_result: bool = Trueoffload_buffers: bool = Falsefallback_allocation: bool = False )

Параметры

  • model (torch.nn.Module) — Модель для анализа.
  • max_memory (Dict, необязательный) — Словарь, сопоставляющий идентификатор устройства с максимальным объёмом памяти. Если не задано, по умолчанию используется максимально доступная память. Пример: max_memory={0: "1GB"}.
  • no_split_module_classes (List[str], необязательный) — Список имён классов слоёв, которые никогда не следует разделять между устройствами (например, любой слой, имеющий остаточное соединение).
  • dtype (str или torch.dtype, необязательный) — Если указано, веса будут преобразованы в этот тип при загрузке.
  • special_dtypes (Dict[str, Union[str, torch.device]], необязательный) — Если указано, специальные dtypes, которые следует учитывать для некоторых конкретных весов (перекрывает dtype, используемый по умолчанию для всех весов).
  • verbose (bool, необязательный, по умолчанию False) — Определяет, следует ли выводить отладочные сообщения при построении device_map функцией.
  • clean_result (bool, необязательный, по умолчанию True) — Очистить итоговый device_map, объединяя все подмодули, размещаемые на одном и том же устройстве.
  • offload_buffers (bool, необязательный, по умолчанию False) — Определяет, следует ли выгружать на CPU или диск буферы в дополнение к параметрам в тех слоях, которые выгружаются на CPU или диск.
  • fallback_allocation (bool, необязательный, по умолчанию False) — Если обычное выделение памяти не удаётся, попытаться выделить модуль, который укладывается в лимит размера, с помощью обхода в ширину (BFS).

Вычислите карту устройств для заданной модели, отдавая приоритет GPU, затем выгружая на CPU и, наконец, на диск, так чтобы:

  • мы не превышаем доступную память ни на одном из GPU.
  • если требуется offload на CPU, на GPU 0 всегда остаётся место для возврата на GPU слоя, выгруженного на CPU, с наибольшим размером.
  • если требуется offload на CPU, мы не превышаем доступный объём ОЗУ на CPU.
  • если требуется offload на диск, на CPU всегда остаётся место для возврата слоя, offload которого на диск выполнен, и который имеет наибольший размер.

Все вычисления выполняются на основе анализа размеров и dtypes параметров модели. В результате модель может находиться на meta-устройстве (как если бы она была инициализирована внутри контекстного менеджера init_empty_weights).

Хуки

ModelHook

класс accelerate.hooks.ModelHook

< >

( )

Хук, содержащий обратные вызовы, которые выполняются непосредственно перед и после метода forward модели. Отличие от существующих хуков PyTorch состоит в том, что они получают в качестве аргумента kwargs.

Атрибут класса:

  • no_grad (bool, необязательный, по умолчанию False) — следует ли выполнять фактический прямой проход в контекстном менеджере torch.no_grad().

detach_hook

< >

( module )

Параметры

  • module (torch.nn.Module) — Модуль, отсоединённый от этого хука.

Выполняется при отключении хука от модуля.

init_hook

< >

( module )

Параметры

  • module (torch.nn.Module) — модуль, привязанный к этому хуку.

Выполняется при подключении хука к модулю.

post_forward

< >

( moduleoutput ) → Any

Параметры

  • module (torch.nn.Module) — модуль, прямой проход которого был выполнен непосредственно перед этим событием.
  • output (Any) — Выход модуля.

Возвращает

Any

Обработанный output.

Выполняется сразу после метода forward модели.

pre_forward

< >

( module*args**kwargs ) → Tuple[Tuple[Any], Dict[Str, Any]]

Параметры

  • module (torch.nn.Module) — модуль, прямой проход которого будет выполнен сразу после этого события.
  • args (Tuple[Any]) — Позиционные аргументы, переданные в модуль.
  • kwargs (Dict[Str, Any]) — Именованные аргументы, передаваемые модулю.

Возвращает

Tuple[Tuple[Any], Dict[Str, Any]]

Кортеж с обработанными args и kwargs.

Выполняется непосредственно перед методом forward модели.

AlignDevicesHook

класс accelerate.hooks.AlignDevicesHook

< >

( execution_device: int | str | torch.device | None = Noneoffload: bool = Falseio_same_device: bool = Falseweights_map: collections.abc.Mapping | None = Noneoffload_buffers: bool = Falseplace_submodules: bool = Falseskip_keys: str | list[str] | None = Nonetied_params_map: dict[int, dict[torch.device, torch.Tensor]] | None = None )

Параметры

  • execution_device (torch.device, необязательный) — Устройство, на котором должны быть размещены входные данные и веса модели перед прямым проходом.
  • offload (bool, необязательный, по умолчанию False) — Следует ли выгружать веса после прямого прохода.
  • io_same_device (bool, необязательный, по умолчанию False) — Определяет, следует ли размещать вывод на том же устройстве, что и вход.
  • weights_map (Mapping[str, torch.Tensor], необязательный) — Если веса модели выгружены, (возможно, ленивая) карта из имён параметров в значения тензоров.
  • offload_buffers (bool, необязательный, по умолчанию False) — Указывает, следует ли включать буферы соответствующего модуля при выгрузке.
  • place_submodules (bool, необязательный, по умолчанию False) — Определяет, следует ли размещать подмодули на execution_device во время события init_hook.

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

SequentialHook

класс accelerate.hooks.SequentialHook

< >

( *hooks )

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

LayerwiseCastingHook

класс accelerate.hooks.LayerwiseCastingHook

< >

( storage_dtype: dtypecompute_dtype: dtypenon_blocking: bool )

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

Добавление хуков

add_hook_to_module

accelerate.hooks.add_hook_to_module

< >

( module: Modulehook: ModelHookappend: bool = False ) → torch.nn.Module

Параметры

  • module (torch.nn.Module) — Модуль, к которому нужно подключить хук.
  • hook (ModelHook) — Хук, который нужно подключить.
  • append (bool, необязательный, по умолчанию False) — Определяет, следует ли объединять хук с уже существующим (если модуль уже содержит хук) или нет.

Возвращает

torch.nn.Module

Тот же модуль, с подключённым хуком (модуль изменяется на месте, поэтому результат можно отбросить).

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

Если модуль уже содержит хук, он будет заменён новым хуком, переданным по умолчанию. Чтобы объединить два хука в цепочку, передайте append=True, тогда текущий и новый хуки будут объединены в экземпляр класса SequentialHook.

attach_execution_device_hook

accelerate.hooks.attach_execution_device_hook

< >

( module: Moduleexecution_device: int | str | torch.deviceskip_keys: str | list[str] | None = Nonepreload_module_classes: list[str] | None = Nonetied_params_map: dict[int, dict[torch.device, torch.Tensor]] | None = None )

Параметры

  • module (torch.nn.Module) — Модуль, к которому мы хотим привязать хуки.
  • execution_device (int, str или torch.device) — Устройство, на котором должны быть размещены входные данные и веса модели перед прямым проходом.
  • skip_keys (str или List[str], необязательный) — Список ключей, которые следует игнорировать при перемещении входных или выходных данных между устройствами.
  • preload_module_classes (List[str], необязательный) — Список классов, экземпляры которых должны загружать все свои веса (в том числе в подмодулях) в начале forward. Это следует использовать только для классов, у которых есть подмодули, зарегистрированные, но не вызываемые напрямую во время forward, например, если зарегистрирован dense линейный слой, но во время forward dense.weight и dense.bias используются в некоторых операциях вместо прямого вызова dense.
  • tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], необязательный, по умолчанию None) — Сопоставление указателей данных со словарями устройств и уже распределённых связанных весов. Для данного устройства выполнения этот параметр полезен для повторного использования первого доступного указателя общего веса для всех остальных, вместо дублирования памяти.

Рекурсивно подключает AlignDevicesHook ко всем подмодулям указанной модели, чтобы убедиться, что у них настроено правильное устройство выполнения

attach_align_device_hook

accelerate.hooks.attach_align_device_hook

< >

( module: Moduleexecution_device: torch.device | None = Noneoffload: bool = Falseweights_map: collections.abc.Mapping | None = Noneoffload_buffers: bool = Falsemodule_name: str = ''skip_keys: str | list[str] | None = Nonepreload_module_classes: list[str] | None = Nonetied_params_map: dict[int, dict[torch.device, torch.Tensor]] | None = None )

Параметры

  • module (torch.nn.Module) — Модуль, к которому мы хотим привязать хуки.
  • execution_device (torch.device, необязательный) — Устройство, на котором должны быть размещены входные данные и веса модели перед прямым проходом.
  • offload (bool, необязательный, по умолчанию False) — Следует ли выгружать веса после прямого прохода.
  • weights_map (Mapping[str, torch.Tensor], необязательный) — Если веса модели выгружены, (возможно, ленивая) карта из имён параметров в значения тензоров.
  • offload_buffers (bool, необязательный, по умолчанию False) — Указывает, следует ли включать буферы соответствующего модуля при выгрузке.
  • module_name (str, необязательный, по умолчанию "") — Имя модуля.
  • skip_keys (str или List[str], необязательный) — Список ключей, которые следует игнорировать при перемещении входных или выходных данных между устройствами.
  • preload_module_classes (List[str], необязательный) — Список классов, экземпляры которых должны загружать все свои веса (в том числе в подмодулях) в начале forward. Это следует использовать только для классов, у которых есть подмодули, зарегистрированные, но не вызываемые напрямую во время forward, например, если зарегистрирован dense линейный слой, но во время forward dense.weight и dense.bias используются в некоторых операциях вместо прямого вызова dense.
  • tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], необязательный, по умолчанию None) — Сопоставление указателей данных со словарями устройств и уже распределённых связанных весов. Для данного устройства выполнения этот параметр полезен для повторного использования первого доступного указателя общего веса для всех остальных, вместо дублирования памяти.

Рекурсивно подключает AlignDevicesHook ко всем подмодулям заданной модели, которые имеют прямые параметры и/или буферы.

attach_align_device_hook_on_blocks

accelerate.hooks.attach_align_device_hook_on_blocks

< >

( module: Moduleexecution_device: torch.device | dict[str, torch.device] | None = Noneoffload: bool | dict[str, bool] = Falseweights_map: collections.abc.Mapping | None = Noneoffload_buffers: bool = Falsemodule_name: str = ''skip_keys: str | list[str] | None = Nonepreload_module_classes: list[str] | None = Nonetied_params_map: dict[int, dict[torch.device, torch.Tensor]] | None = None )

Параметры

  • module (torch.nn.Module) — Модуль, к которому мы хотим привязать хуки.
  • execution_device (torch.device или Dict[str, torch.device], необязательный) — Устройство, на котором должны размещаться входные данные и веса модели перед прямым проходом. Может быть одним устройством для всего модуля или словарем, сопоставляющим имя модуля с устройством.
  • offload (bool, необязательный, по умолчанию False) — Определяет, следует ли выгружать веса после прямого прохода. Может быть одним булевым значением для всего модуля или словарем, сопоставляющим имена модулей с булевыми значениями.
  • weights_map (Mapping[str, torch.Tensor], необязательный) — Если веса модели выгружены, (возможно, ленивая) карта из имён параметров в значения тензоров.
  • offload_buffers (bool, необязательный, по умолчанию False) — Указывает, следует ли включать буферы соответствующего модуля при выгрузке.
  • module_name (str, необязательный, по умолчанию "") — Имя модуля.
  • skip_keys (str или List[str], необязательный) — Список ключей, которые следует игнорировать при перемещении входных или выходных данных между устройствами.
  • preload_module_classes (List[str], необязательный) — Список классов, экземпляры которых должны загружать все свои веса (в том числе в подмодулях) в начале forward. Это следует использовать только для классов, у которых есть подмодули, зарегистрированные, но не вызываемые напрямую во время forward, например, если зарегистрирован dense линейный слой, но во время forward dense.weight и dense.bias используются в некоторых операциях вместо прямого вызова dense.
  • tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], необязательный, по умолчанию None) — Сопоставление указателей данных со словарями устройств и уже распределённых связанных весов. Для данного устройства выполнения этот параметр полезен для повторного использования первого доступного указателя общего веса для всех остальных, вместо дублирования памяти.

Подключает AlignDevicesHook ко всем блокам указанной модели по мере необходимости.

attach_layerwise_casting_hooks

accelerate.big_modeling.attach_layerwise_casting_hooks

< >

( module: Modulestorage_dtype: dtypecompute_dtype: dtypeskip_modules_pattern: str | tuple[str, ...] | None = Noneskip_modules_classes: tuple[type[torch.nn.Module], ...] | None = Nonenon_blocking: bool = False )

Параметры

  • module (torch.nn.Module) — Модуль, дочерние модули которого будут приведены к типу данных высокой точности для вычислений и к типу данных низкой точности для хранения.
  • storage_dtype (torch.dtype) — dtype, к которому модуль приводится до/после прямого прохода для хранения.
  • compute_dtype (torch.dtype) — Тип данных (dtype), к которому приводится модуль во время прямого прохода для вычислений.
  • skip_modules_pattern (tuple[str, ...], по умолчанию None) — Список шаблонов для сопоставления с именами модулей, которые нужно пропустить в процессе послойного приведения типов. Если установлено в None при том, что skip_modules_classes также равно None, послойное приведение типов применяется непосредственно к модулю, а не к его внутренним подмодулям.
  • skip_modules_classes (tuple[type[torch.nn.Module], ...], по умолчанию None) — Список классов модулей, которые нужно пропустить при послойном приведении типов.
  • non_blocking (bool, по умолчанию False) — Если True, операции приведения типов весов выполняются без блокировки.

Применяет послойное приведение типов к заданному модулю. Ожидаемый модуль — это PyTorch nn.Module. Это полезно для снижения требований к памяти, когда нет необходимости полностью квантовать модель. Параметры модели можно хранить, например, в формате torch.float8_e4m3fn и приводить к более высокой точности, такой как torch.bfloat16, во время прямого прохода, а затем снова приводить к torch.float8_e4m3fn для экономии памяти.

Пример:

>>> from accelerate.hooks import attach_layerwise_casting_hooks
>>> from transformers import AutoModelForCausalLM
>>> import torch

>>> # Model
>>> checkpoint = "EleutherAI/gpt-j-6B"
>>> model = AutoModelForCausalLM.from_pretrained(checkpoint)

>>> # Attach hooks and perform inference
>>> attach_layerwise_casting_hooks(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)
>>> with torch.no_grad():
...     model(...)

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

>>> attach_layerwise_casting_hooks(
...     model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16, skip_modules_pattern=["norm"]
... )

Удаление хуков

remove_hook_from_module

accelerate.hooks.remove_hook_from_module

< >

( module: Modulerecurse = False ) → torch.nn.Module

Параметры

  • module (torch.nn.Module) — Модуль, к которому нужно подключить хук.
  • recurse (bool, необязательный) — Указывает, следует ли удалять хуки рекурсивно

Возвращает

torch.nn.Module

Тот же модуль, с отключённым хуком (модуль изменяется на месте, поэтому результат можно отбросить).

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

remove_hook_from_submodules

accelerate.hooks.remove_hook_from_submodules

< >

( module: Module )

Параметры

  • module (torch.nn.Module) — модуль, с которого нужно удалить все хуки.

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

Утилиты

has_offloaded_params

accelerate.utils.has_offloaded_params

< >

( module: Module ) → bool

Параметры

  • module (torch.nn.Module) — Модуль, для которого проверяется наличие хука offload.

Возвращает

bool

True, если у модуля есть хук offload и offloading включён, False в противном случае.

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

align_module_device

accelerate.utils.align_module_device

< >

( module: Moduleexecution_device: torch.device | None = None )

Параметры

  • module (torch.nn.Module) — Модуль с параметрами для выравнивания.
  • execution_device (torch.device, необязательный) — Если указан, переопределяет устройство выполнения модуля в пределах контекста. В противном случае используется устройство выполнения хука или значение по умолчанию

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

Update on GitHub

←Логирование Обработчики аргументов Kwargs→