Accelerate

( 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 参数，否则不会调用调度。

cpu_offload

accelerate.cpu_offload

( model: Module execution_device: typing.Optional[torch.device] = None offload_buffers: bool = False state_dict: typing.Optional[dict[str, torch.Tensor]] = None preload_module_classes: typing.Optional[list[str]] = None )

参数

model (torch.nn.Module) — 要卸载的模型。
execution_device (torch.device, optional) — 模型前向传播将在其上执行的设备（应为 GPU）。将默认为模型的第一个参数所在的设备。
offload_buffers (bool, optional, 默认为 False) — 是否将缓冲区与模型参数一起卸载。
state_dict (Dict[str, torch.Tensor], optional) — 将保存在 CPU 上的模型的状态字典。
preload_module_classes (List[str], optional) — 一个类名列表，这些类的实例应在前向传播开始时加载其所有权重（即使在子模块中）。这应该只用于那些具有已注册但在前向传播期间未直接调用的子模块的类，例如，如果注册了一个 `dense` 线性层，但在前向传播中，`dense.weight` 和 `dense.bias` 在某些操作中使用，而不是直接调用 `dense`。

为一个模型激活完全 CPU 卸载。因此，模型的所有参数将被卸载，并且只保留模型状态字典的一个副本。在前向传播期间，参数将从该状态字典中提取，并根据需要放到指定的执行设备上，然后再重新卸载。

cpu_offload_with_hook

accelerate.cpu_offload_with_hook

( model: Module execution_device: typing.Union[int, str, torch.device, NoneType] = None prev_module_hook: typing.Optional[accelerate.hooks.UserCpuOffloadHook] = None )

参数

model (torch.nn.Module) — 要卸载的模型。
execution_device(str, int 或 torch.device, optional) — 模型应在其上执行的设备。如果 MPS 设备可用，则默认为 MPS 设备；否则，如果有 GPU，则为 GPU 0；最后为 CPU。
prev_module_hook (UserCpuOffloadHook, optional) — 此函数为正在运行的流水线中前一个模型返回的钩子。如果传递了此参数，其 offload 方法将在附加此钩子的模型的前向传播之前被调用。

将模型卸载到 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()

disk_offload

accelerate.disk_offload

( model: Module offload_dir: typing.Union[str, os.PathLike] execution_device: typing.Optional[torch.device] = None offload_buffers: bool = False preload_module_classes: typing.Optional[list[str]] = None )

参数

model (torch.nn.Module) — 要卸载的模型。
offload_dir (str 或 os.PathLike) — 用于卸载模型权重的文件夹（或者模型权重已经卸载到的位置）。
execution_device (torch.device, optional) — 模型前向传播将在其上执行的设备（应为 GPU）。将默认为模型的第一个参数所在的设备。
offload_buffers (bool, optional, 默认为 False) — 是否将缓冲区与模型参数一起卸载。
preload_module_classes (List[str], optional) — 一个类名列表，这些类的实例应在前向传播开始时加载其所有权重（即使在子模块中）。这应该只用于那些具有已注册但在前向传播期间未直接调用的子模块的类，例如，如果注册了一个 `dense` 线性层，但在前向传播中，`dense.weight` 和 `dense.bias` 在某些操作中使用，而不是直接调用 `dense`。

为模型激活完全磁盘卸载。因此，模型的所有参数都将作为内存映射数组卸载到给定文件夹中。在前向传播期间，将从该文件夹访问参数，并根据需要将其放置到指定的执行设备上，然后再次卸载。

dispatch_model

accelerate.dispatch_model

( model: Module device_map: dict main_device: typing.Optional[torch.device] = None state_dict: typing.Optional[dict[str, torch.Tensor]] = None offload_dir: typing.Union[str, os.PathLike, NoneType] = None offload_index: typing.Optional[dict[str, str]] = None offload_buffers: bool = False skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None force_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, optional) — 主要执行设备。将默认为 `device_map` 中第一个不同于 `"cpu"` 或 `"disk"` 的设备。
state_dict (Dict[str, torch.Tensor], optional) — 将保存在 CPU 上的模型部分的状态字典。
offload_dir (str 或 os.PathLike) — 用于卸载模型权重的文件夹（或者模型权重已经卸载到的位置）。
offload_index (Dict, optional) — 一个从权重名称到其信息（`dtype`/`shape` 或 safetensors 文件名）的字典。将默认为 `save_folder` 中保存的索引。
offload_buffers (bool, optional, 默认为 False) — 是否将缓冲区与模型参数一起卸载。
skip_keys (str 或 List[str], optional) — 在设备之间移动输入或输出时要忽略的键的列表。
preload_module_classes (List[str], optional) — 一个类名列表，这些类的实例应在前向传播开始时加载其所有权重（即使在子模块中）。这应该只用于那些具有已注册但在前向传播期间未直接调用的子模块的类，例如，如果注册了一个 `dense` 线性层，但在前向传播中，`dense.weight` 和 `dense.bias` 在某些操作中使用，而不是直接调用 `dense`。
force_hooks (bool, optional, 默认为 False) — 是否强制将设备钩子附加到模型，即使所有层都被调度到单个设备。

根据给定的设备映射调度模型。模型的层可能分布在不同的 GPU 上，卸载到 CPU 甚至磁盘上。

load_checkpoint_and_dispatch

accelerate.load_checkpoint_and_dispatch

( model: Module checkpoint: typing.Union[str, os.PathLike] device_map: typing.Union[str, dict[str, typing.Union[int, str, torch.device]], NoneType] = None max_memory: typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None no_split_module_classes: typing.Optional[list[str]] = None offload_folder: typing.Union[str, os.PathLike, NoneType] = None offload_buffers: bool = False dtype: typing.Union[str, torch.dtype, NoneType] = None offload_state_dict: typing.Optional[bool] = None skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None force_hooks: bool = False strict: bool = False full_state_dict: bool = True broadcast_from_rank0: bool = False )

参数

model (torch.nn.Module) — 我们希望加载检查点的模型。
checkpoint (str 或 os.PathLike) — 要加载的文件夹检查点。它可以是：
- 包含完整模型状态字典的文件的路径
- 包含分片检查点索引的 .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 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 状态字典的权重加上最大分片的权重不适合时耗尽 CPU RAM。如果选择的设备映射包含 "disk" 值，则默认为 True。
skip_keys (str 或 List[str], 可选) — 在设备之间移动输入或输出时要忽略的键的列表。
preload_module_classes (List[str], 可选) — 一个类的列表，其实例应在 forward 开始时加载所有权重（即使在子模块中）。这只应用于那些已注册但在 forward 期间未直接调用的子模块的类，例如，如果注册了一个 dense 线性层，但在 forward 时，在某些操作中使用了 dense.weight 和 dense.bias，而不是直接调用 dense。
force_hooks (bool, 可选, 默认为 False) — 是否强制将设备钩子附加到模型上，即使所有层都分派到单个设备。
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 中的张量逐一广播到其他 rank。其他 rank 将接收张量并根据模型中的本地分片进行分片（如果适用）。

将一个（可能是分片的）检查点加载到模型中，加载时可能会将权重发送到给定的设备，并添加各种钩子，使该模型能够正常运行（即使跨设备拆分）。

示例

>>> 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: Module checkpoint: typing.Union[str, os.PathLike] device_map: typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None offload_folder: typing.Union[str, os.PathLike, NoneType] = None dtype: typing.Union[str, torch.dtype, NoneType] = None offload_state_dict: bool = False offload_buffers: bool = False keep_in_fp32_modules: list = None offload_8bit_bnb: bool = False strict: bool = False full_state_dict: bool = True broadcast_from_rank0: 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 状态字典的权重加上最大分片的权重不适合时耗尽 CPU RAM。
offload_buffers (bool, 可选, 默认为 False) — 是否将缓冲区包含在卸载到磁盘的权重中。
keep_in_fp32_modules(List[str], 可选) — 保持 torch.float32 数据类型的模块列表。
offload_8bit_bnb (bool, 可选) — 是否在 cpu/磁盘上启用 8-bit 模块的卸载。
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 中的张量逐一广播到其他 rank。其他 rank 将接收张量并根据模型中的本地分片进行分片（如果适用）。

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

跨设备加载后，您仍需要在模型上调用 dispatch_model() 才能使其运行。要将检查点加载和分派组合在一次调用中，请使用 load_checkpoint_and_dispatch()。

infer_auto_device_map

accelerate.infer_auto_device_map

( model: Module max_memory: typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None no_split_module_classes: typing.Optional[list[str]] = None dtype: typing.Union[str, torch.dtype, NoneType] = None special_dtypes: typing.Optional[dict[str, typing.Union[str, torch.dtype]]] = None verbose: bool = False clean_result: bool = True offload_buffers: bool = False fallback_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]], 可选) — 如果提供，将为某些特定权重考虑特殊的数据类型（将覆盖用作所有权重默认值的 dtype）。
verbose (bool, 可选, 默认为 False) — 是否在函数构建 device_map 时提供调试语句。
clean_result (bool, 可选, 默认为 True) — 通过将所有放置在同一设备上的子模块分组来清理生成的 device_map。
offload_buffers (bool, 可选, 默认为 False) — 在卸载到 CPU 或硬盘的层中，是否也将缓冲区与参数一起卸载。
fallback_allocation (bool, 可选, 默认为 False) — 当常规分配失败时，尝试使用 BFS 分配一个适合大小限制的模块。

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

我们不会超过任何 GPU 的可用内存。
如果需要卸载到 CPU，GPU 0 上始终留有空间来放回卸载到 CPU 的最大尺寸的层。
如果需要卸载到 CPU，我们不会超过 CPU 上可用的 RAM。
如果需要卸载到磁盘，CPU 上始终留有空间来放回卸载到磁盘的最大尺寸的层。

所有的计算都是通过分析模型参数的大小和数据类型来完成的。因此，模型可以位于元设备上（就像在 init_empty_weights 上下文管理器中初始化时一样）。

钩子 (Hooks)

ModelHook

class 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

( module output ) → 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

class accelerate.hooks.AlignDevicesHook

( execution_device: typing.Union[int, str, torch.device, NoneType] = None offload: bool = False io_same_device: bool = False weights_map: typing.Optional[collections.abc.Mapping] = None offload_buffers: bool = False place_submodules: bool = False skip_keys: typing.Union[str, list[str], NoneType] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = 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) — 是否在 init_hook 事件期间将子模块放置在 execution_device 上。

一个通用的 ModelHook，确保输入和模型权重在关联模块的前向传播中位于同一设备上，并可能在前向传播后卸载权重。

SequentialHook

class accelerate.hooks.SequentialHook

( *hooks )

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

LayerwiseCastingHook

class accelerate.hooks.LayerwiseCastingHook

( storage_dtype: dtype compute_dtype: dtype non_blocking: bool )

一个将模块权重转换为高精度 dtype 进行计算，并转换为低精度 dtype 进行存储的钩子。此过程可能会导致输出质量下降，但可以显著减少内存占用。

添加钩子

add_hook_to_module

accelerate.hooks.add_hook_to_module

( module: Module hook: ModelHook append: 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: Module execution_device: typing.Union[int, str, torch.device] skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None )

参数

module (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 (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], 可选, 默认为 None) — 一个从数据指针到设备字典的映射，该字典包含已调度的绑定权重。对于给定的执行设备，此参数可用于为所有共享权重重用第一个可用的指针，而不是复制内存。

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

attach_align_device_hook

accelerate.hooks.attach_align_device_hook

( module: Module execution_device: typing.Optional[torch.device] = None offload: bool = False weights_map: typing.Optional[collections.abc.Mapping] = None offload_buffers: bool = False module_name: str = '' skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = 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], 可选) — 一个类列表，其类的实例应在前向传播开始时加载所有权重（即使在子模块中）。这应该只用于那些具有已注册但在前向传播期间未直接调用的子模块的类，例如，如果注册了一个 dense 线性层，但在前向传播时，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: Module execution_device: typing.Union[torch.device, dict[str, torch.device], NoneType] = None offload: typing.Union[bool, dict[str, bool]] = False weights_map: Mapping = None offload_buffers: bool = False module_name: str = '' skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = 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