处理大型模型

分派和卸载模型

accelerate.init_empty_weights

( include_buffers: bool = None )

参数

include_buffers (bool, optional) — 是否在初始化时也将所有缓冲区放在元设备上。

模型在该上下文管理器下初始化，此管理器包含元设备上的所有参数，因此创建了一个空模型。当仅初始化模型会耗尽可用 RAM 时很有用。

示例

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

在此上下文管理器下创建的任何模型都没有权重。因此，你不能对它执行如下操作：model.to(some_device)。要在空模型内部加载权重，请参见 load_checkpoint_and_dispatch()。请务必覆盖 load_checkpoint_and_dispatch() 的默认 device_map 参数，否则不会调用 dispatch。

accelerate.cpu_offload

< 源 >

( model: 模块 execution_device: 可选 = 无 offload_buffers: bool = False state_dict: 可选 = 无 preload_module_classes: 可选 = 无 )

参数

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]，可选) — 列出所有实例需要在正向开始时加载其所有权重（甚至在子模块中）的类。仅适用于在其子模块已注册但在正向过程中未直接调用的类，例如，注册了dense线性层，但在正向过程中，在某些操作中使用dense.weight和dense.bias代替直接调用dense。

对模型激活完全的 CPU 卸载。因此，模型的所有参数都将被卸载，且仅保留模型状态字典的一个副本。在正向传递期间，将从该状态字典中提取参数，并将它们放在按需要传递的执行设备上，然后再次卸载。

accelerate.cpu_offload_with_hook

< 来源 >

( 模型: 模块执行设备: 联合 = 无上一个模块钩: 可选 = 无 )

参数

模型(torch.nn.Module) — 要卸载的模型。
execution_device(str，int或torch.device，可选) — 执行模型的设备。如果可用，将默认为 MPS 设备，如果存在 GPU，则为 GPU 0，最后为 CPU。
`prev_module_hook` (UserCpuOffloadHook, _可选_ )—此函数返回的挂钩，用于您运行的管道中的上一个模型。如果通过，它的卸载方法将在此挂钩附加的模型前向传递之前调用。

在 CPU 上卸载模型并执行时将其放回执行设备。与 `cpu_offload()` 的区别在于，模型在前向传播后仍保留在执行设备上，并且仅在调用的返回 `hook` 的 `offload` 方法时才重新卸载。对于在循环中运行模型的管道很有用。

示例

model_1, hook_1 = cpu_offload_with_hook(model_1, cuda_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, cuda_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()

`accelerate.disk_offload`

< 源代码 >

( 模型: 模块 offload_dir: 联合执行设备: 可选 = 无 offload_buffers: bool = 假预加载模块类: 可选 = 无 )

参数

模型 (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。

激活模型的全磁盘卸载。结果，该模型的所有参数将作为一个给定文件夹中的内存映射数组被卸载。在 forward pass 期间，参数将被从该文件夹访问并按照需要放置在执行的设备上，然后再次卸载。

accelerate.dispatch_model

< 源代码 >

( model: 模型 device_map: 字典 main_device: 可选 = 无 state_dict: 可选 = 无 offload_dir: 联合 = 无 offload_index: 可选 = 无 offload_buffers: 布尔值 = false skip_keys: 联合 = 无 preload_module_classes: 可选 = 无 force_hooks: 布尔值 = false )

参数

模型 (torch.nn.Module) — 要分派的模型。
device_map (Dict[str, Union[str, int, torch.device]]) — 用模型中模块 state_dict 的名称映射它们应该去的设备的词典。请注意，即便是 torch.device 的不当值，也接受 "disk"。
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]，可选) — 应该在前向开始时加载所有其权重（即使是在子模块中）的类的列表。这应仅用于具有子模块的类，这些类已注册但未在正向过程中直接调用，例如，如果注册了dense线性图层，但在正向过程中，dense.weight和dense.bias在某些操作中用作替代直接调用dense。
force_hooks (bool, 可选，默认为 False) — 无论所有层是否分发给某个单一设备，都要强制将设备连接器附加至模型。

根据特定设备的映射调度模型。此模型的层可能散落在 GPU 跨的不同区域，甚至被卸载到 CPU 或磁盘中。

accelerate.load_checkpoint_and_dispatch

< 源代码 >

( 模型: 模块检查点: 联合设备映射: 联合 = 无最大内存: 可选 = 无不拆分模块类: 可选 = 无卸载文件夹: 联合 = 无卸载缓冲区: 布尔值 = 假 dtype: 联合 = 无卸载状态字典: 可选 = 无跳过键: 联合 = 无预加载模块类: 可选 = 无强制挂钩: 布尔值 = 假严格: 布尔值 = 假 )

参数

模型 (torch.nn.Module) — 要在其中加载检查点的模型。
checkpoint (str 或 os.PathLike) — 要加载的文件夹检查点。可以是：
- 包含整个模型状态字典的文件路径
- 包含指向分片检查点的索引的 .json 文件路径
- 包含唯一的 .index.json 文件和检查点分片的文件夹路径。
device_map (Dict[str, Union[int, str, torch.device]]，可选) — 指定每个子模块应转至何处的映射。无需细化到每个参数/缓冲区名称，一旦给定模块名称在其内部，其每个子模块都将发送到同一设备。

要让加速计算最优化的 device_map，设置 device_map="auto"。有关每个选项的更多信息，请参阅此处。默认为 None，这意味着 dispatch_model() 将不会被调用。
max_memory (Dict, 可选项) — 从设备标识符到最大内存的字典，若未设置，每个 GPU 将默认为可用最大内存，CPU RAM 也为可用内存。
no_split_module_classes (List[str], 可选项) — 一些不应该在跨设备中分割的层类名的列表（例如，拥有残差连接的层）。
offload_folder (str或os.PathLike，可选) — 如果device_map包含任何值"disk"，我们将卸载权重的文件夹。
offload_buffers (bool，可选，默认为False) — 在卸载至 CPU 或硬盘的层中，是否卸载缓冲区以及参数。
dtype (str 或 torch.dtype，可选) — 如果提供此参数，权重在加载时会转换成该类型。
offload_state_dict (bool，可选) — 如果为 True，将暂时卸下硬盘上的 CPU 状态字典，以避免 CPU RAM 耗尽（如果 CPU 状态字典的权重 + 最大碎片不合适）。如果选定的设备映射包含 "disk" 值，则默认值为 True。
skip_keys (字符串或列表[字符串]，可选) — 在设备之间移动输入或输出时需要忽略的键列表。
preload_module_classes (列表[字符串]，可选) — 应在 forward 开始时加载所有权重（即使在子模块中）的类的列表。这仅应用于具有子模块的类，这些子模块已注册但未在 forward 期间直接调用，例如，如果注册了dense线性层，但在 forward 中，dense.weight和dense.bias在某些操作中使用，而不是直接调用dense。
force_hooks (bool, 可选，默认为 false) — 不管所有层分派到哪个设备，是否强制将设备挂钩附加到模型。
strict (bool, 可选，默认为 false) — 是否严格执行检查点的 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"]
... )

accelerate.load_checkpoint_in_model

< 源 >

( model: 模块 checkpoint: 联合 device_map: 可选 = None offload_folder: 联合 = None dtype: 联合 = None offload_state_dict: bool = False offload_buffers: bool = False keep_in_fp32_modules: 列表 = None offload_8bit_bnb: bool = False strict: bool = False )

参数

model (torch.nn.Module) — 要加载检查点的模型。
checkpoint (str 或 os.PathLike) — 要加载的文件夹检查点。它可以是：
- 包含完整模型状态字典的文件的路径
- 包含分片检查点索引的 .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 RAM 耗尽，如果 CPU 状态字典的权重加上最大的分片不合适。
offload_buffers (布尔值，可选，默认为 False) — 是否在权重中包含卸载到磁盘的缓冲。
keep_in_fp32_modules(List[str], 可选) — 将我们保留在 torch.float32 数据类型的模块列表。
strict (bool, 可选，默认为 False) — 是否严格强制检查检查点 state_dict 中的键与模型 state_dict 中的键匹配。

加载一个（可能是分片）检查点在一个模型中，潜在地在加载时将权重发送到一个给定的设备。

一旦加载到各个设备，您仍然需要对您的模型调用 dispatch_model() 以使其能够运行。要在一个调用中对检查点加载和分派进行分组，请使用 load_checkpoint_and_dispatch()。

accelerate.infer_auto_device_map

< 源代码 >

（模型：模块最大内存：可选 = 无不拆分模块类：可选 = 无数据类型：联合 = 无特殊数据类型：可选 = 无详细：bool = 假清除结果：bool = 真卸载缓冲区：bool = 假）

参数

模型 (torch.nn.Module) — 要分析的模型。
max_memory (字典，可选) — 设备标识符到最大内存的字典。如果未设置，将默认为可用的最大内存。示例：max_memory={0: "1GB"}.
no_split_module_classes ([str]列表，可选) — 层类名称的列表，它们绝不会跨设备拆分（例如，任何具有残差连接的层）。
dtype (str 或 torch.dtype，可选) —如果提供了此项，加载时会将权重转换为该类型。
special_dtypes (Dict[str, Union[str, torch.device]]，可选) —如果提供了此项，将考虑某些特定权重的特殊数据类型（将覆盖用作所有权重的默认数据类型）。
详细 (bool，可选，默认为 False) — 在函数构建设备映射时是否提供调试语句。
clean_result (bool，可选，默认为 True) — 通过将所有放在同一台设备上的子模块分组，来清理生成的设备映射。
offload_buffers (bool, 可选，默认为 False) — 对于卸载到 CPU 或硬盘上的层，是否卸载缓冲区以及参数。

为给定模型计算设备映射，优先考虑 GPU，然后卸载到 CPU 最后卸载到磁盘，以便

我们不会超过任何 GPU 的可用内存。
如果需要卸载到 CPU，GPU 0 上始终有剩余的空间可以将卸载到 CPU 上且具有最大大小的层放回 GPU 0。
如果需要卸载到 CPU，我们不会超过 CPU 上可用的 RAM。
如果需要卸载到磁盘，CPU 上始终有剩余的空间可以将卸载到磁盘上且具有最大大小的层放回 CPU。

所有计算都是通过分析模型参数的大小和数据类型完成的。结果，该模型可以在元设备上（如同在 init_empty_weights 上下文管理器中初始化时一样）。

模型挂钩

钩子类

类 accelerate.hooks.ModelHook

< 源代码 >

( )

在模型的正向方法执行前和执行后会被调用的钩子。与 PyTorch 现有的钩子之间的区别是它们通过关键字参数传递。

类属性

no_grad (bool, 可选, 默认为 False) — 是否在 torch.no_grad() 上下文管理器下执行实际的前向传递。

detach_hook

< 源代码 >

( 模块 )

参数

模块 (torch.nn.Module) — 从此钩子中分离的模块。

当钩子从一个模块中分离时执行。

init_hook

< 源代码 >

( 模块 )

参数

模块 (torch.nn.Module) — 附加到此挂钩的模块。

当挂钩附加到模块时要执行。

post_forward

( 模块输出 ) → Any

参数

模块 (torch.nn.Module) — 就在此事件之前执行了其前向传递的模块。
输出 (Any) — 模块的输出。

返回值

Any

经过处理的 output。

将在模型的前向方法后立即执行。

pre_forward

< 源 >

( 模块 *args **kwargs ) → 元组[元组[任何], 字典[字符串, 任何]]

参数

模块 (torch.nn.Module) — 在此事件后将执行其前向传播的模块。
args (元组 [Any]) — 传递给模块的指定参数。
kwargs (字典 [Str, Any]) — 传递给模块的关键字参数。

返回值

元组 [元组 [Any], 字典 [Str, Any]]

一个元组，包含处理过的 args 和 kwargs。

在模型的前向方法正执行之前执行。

类 accelerate.hooks.AlignDevicesHook

< 源代码 >

( execution_device: Union = None offload: bool = False io_same_device: bool = False weights_map: Optional = None offload_buffers: bool = False place_submodules: bool = False skip_keys: Union = None tied_params_map: Optional = None )

参数

execution_device (torch.device, 可选) — 输入和模型权重在正向传递前放置到的设备。
offload (bool, 可选，默认为 False) — 正向传递后权重是否卸载。
io_same_device (bool, 可选，默认为 False) — 输出是否应用于与输入相同设备的位置。
weights_map (Mapping[str, torch.Tensor], 可选) — 当模型权重被卸载时，一个（可能延迟的）从参数名到张量值的映射。
offload_buffers (布尔型，可选，默认为 False) - 在卸载时是否包括相关模块的缓冲区。
place_submodules (布尔型，可选，默认为 False) — 在 init_hook 事件期间是否将子模块放置在 execution_device 上。

通用 ModelHook，用于确保在相关模块的前向传递中输入和模型权重位于同一设备上，在完成前向传递后卸载权重。

类 accelerate.hooks.SequentialHook

< 来源 >

( *hooks )

一个包含几个钩子的钩子，并在每个事件中遍历它们。

添加钩子

accelerate.hooks.add_hook_to_module

< 源代码 >

( module: 模块 hook: 模型钩子 append: bool = False ) → torch.nn.Module

参数

模块 (torch.nn.Module) — 要附加钩子的模块。
hook (ModelHook) — 要附加的钩子。
append (bool, 可选，默认为 False) — 钩子是否应该与现有钩子（如果模块已包含钩子）链接或不链接。

返回值

torch.nn.Module

附加了钩子的同一模块（此模块在原地修改，因此可以丢弃结果）。

向给定模块添加钩子。这将重写模块的 forward 方法以包括钩子，为去除此行为并恢复原始 forward 方法，请使用 remove_hook_from_module。

如果模块已经包含钩子，这将用默认传递的新钩子替换它。要将两个钩子连接在一起，请传递 append=True，以便将当前钩子和新钩子链接到 SequentialHook 类的一个实例中。

accelerate.hooks.attach_execution_device_hook

< 源码 >

( 模块：模块执行设备：联合跳过密钥：联合 = 无预加载模块类：可选 = 无绑定的参数映射：可选 = 无 )

参数

模块 (torch.nn.Module) — 希望将钩子附加到的模块。
execution_device (int、str 或 torch.device) — 在前向传递之前，应该将输入和模型权重放在该设备上。
skip_keys (str 或 List[str]，可选) — 在设备之间移动输入或输出时，要忽略的键列表。
preload_module_classes (List[str], 可选) — 一个列表，该列表的实例应在其前向开始时加载它们的所有权重（即使在子模块中）。这只应适用于具有子模块的类，这些类已注册但未在正向期间直接调用，例如，如果注册了dense线性层，但在正向时，在某些操作中将使用dense.weight和dense.bias，而不是直接调用 dense。
tied_params_map (可选[Dict[int, Dict[torch.device, torch.Tensor]]], 可选, 默认值为 None) — 数据指针映射到已经分发的权重设备的字典。对于给定的执行设备，此参数可用于对除所有其他设备外的共享权重的第一个可用指针进行重用，而不是复制内存。

递归地将 AlignDevicesHook 附加到给定模型的所有子模块，以确保它们具有正确的执行设备

accelerate.hooks.attach_align_device_hook

< 源代码 >

( module: 模块 execution_device: 可选 = 无 offload: 布尔值 = 错误 weights_map: 可选 = 无 offload_buffers: 布尔值 = 错误 module_name: 字符串 = '' skip_keys: 联合 = 无 preload_module_classes: 可选 = 无 tied_params_map: 可选 = 无 )

参数

模块 (torch.nn.Module) — 要附加钩子的模块。
execution_device (torch.device, 可选) — 正向传递之前输入和模型权重应置于其上的设备。
weights_map (Mapping[str, torch.Tensor], 可选) —当模型权重卸载后，从参数名到张量值的（潜在的延迟）映射。
offload_buffers (bool, 可选，默认值为 False) —卸载时是否包括关联模块的缓冲区。
skip_keys (str 或 List[str]，可选) — 在设备间移动输入或输出时的忽略键列表。
preload_module_classes (List[str]，可选) — 其实例应在其正向开始时加载其所有权重（即使在子模块中）的类的列表。这仅应用于子模块已注册但正向期间未直接调用的类的列表，例如如果注册了一个dense 线性层，但正向时，dense.weight 和 dense.bias 会在某些操作中使用，而不是直接调用 dense。
tied_params_map (可选[Dict[int, Dict[torch.device, torch.Tensor]]], 可选，默认为 None) — 设备字典的数据指针图已分派了权重。对于给定的执行设备，此参数可用于为所有其他设备重用共享权重的第一个可用指针，而不是重复内存。

递归地附加 AlignDevicesHook 到给定模型的所有具有直接参数和/或缓冲区的子模块。

accelerate.hooks.attach_align_device_hook_on_blocks

< 源代码 >

( 模块: 模块执行设备: 联合 = 无卸载: 联合 = 假权重映射: 映射 = 无卸载缓冲区: 布尔值 = 假模块名称: str = '' 跳过键: 联合 = 无预加载模块类: 可选 = 无关联参数映射: 可选 = 无 )

参数

模块 (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（可选的[字典[int，字典[torch.device，torch.Tensor]]], 可选的，默认为None）— 这是一种字典，它的数据指针到设备的字典已分配的共享权重。对于给定的执行设备，此参数可用于为所有其他设备重用共享权重的第一个可用指针，而不是重复内存。

根据需要将AlignDevicesHook附加到给定模型的所有块。

移除钩子

accelerate.hooks.remove_hook_from_module

< 源代码 >

( module: 模块 recurse = False ) → torch.nn.Module

参数

module (torch.nn.Module) — 要附加钩子的模块。
recurse (bool, optional) — 是否递归删除钩子

返回值

torch.nn.Module

使用 add_hook_to_module 附加到模块上的钩子（模块在原处修改，因此可以丢弃结果）。

移除通过 add_hook_to_module 附加到模块上的任何钩子。

accelerate.hooks.remove_hook_from_submodules

< 源代码 >

( 模块: 模块 )

参数

module (torch.nn.Module) — 要从中删除所有钩子的模块。

递归地删除附加在给定模型的子模块上的所有钩子。

< > 更新于 GitHub

加速

处理大型模型

分派和卸载模型

accelerate.init_empty_weights

accelerate.cpu_offload

accelerate.cpu_offload_with_hook

`accelerate.disk_offload`

accelerate.dispatch_model

accelerate.load_checkpoint_and_dispatch

accelerate.load_checkpoint_in_model

accelerate.infer_auto_device_map

模型挂钩

钩子类

类 accelerate.hooks.ModelHook

detach_hook

init_hook

post_forward

pre_forward

类 accelerate.hooks.AlignDevicesHook

类 accelerate.hooks.SequentialHook

添加钩子

accelerate.hooks.add_hook_to_module

accelerate.hooks.attach_execution_device_hook

accelerate.hooks.attach_align_device_hook

accelerate.hooks.attach_align_device_hook_on_blocks

移除钩子

accelerate.hooks.remove_hook_from_module

accelerate.hooks.remove_hook_from_submodules