Hugging Face 的徽标

Accelerate 文档

处理大型模型

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

处理大型模型

调度和卸载

init_empty_weights

accelerate.init_empty_weights

< >

( include_buffers: bool = None )

参数

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

一个上下文管理器，在此管理器下，模型的所有参数都在 meta 设备上初始化，从而创建一个空模型。当仅初始化模型就会耗尽可用 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, 可选) — 模型的前向传递将在其上执行的设备（应为 GPU）。 将默认为模型的第一个参数设备。
  • offload_buffers (bool, 可选, 默认为 False) — 是否将缓冲区与模型参数一起卸载。
  • state_dict (Dict[str, torch.Tensor], 可选) — 模型的 state dict，将保留在 CPU 上。
  • preload_module_classes (List[str], 可选) — 类列表，其实例应在其 forward 的开头加载所有权重（即使在子模块中）。 这应仅用于具有已注册但未在 forward 期间直接调用的子模块的类，例如，如果注册了一个 dense 线性层，但在 forward 时，使用 dense.weight 和 dense.bias 进行某些操作，而不是直接调用 dense。

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

cpu_offload_with_hook

accelerate.cpu_offload_with_hook

< >

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

参数

  • model (torch.nn.Module) — 要卸载的模型。
  • execution_device(str, int 或 torch.device, 可选) — 模型应在其上执行的设备。 如果 MPS 设备可用，则默认为 MPS 设备，如果存在 GPU，则默认为 GPU 0，最后默认为 CPU。
  • prev_module_hook (UserCpuOffloadHook, 可选) — 此函数为正在运行的管道中的先前模型发送回的钩子。 如果传递，则将在附加此钩子的模型的前向传递之前调用其 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, 可选) — 模型前向传递将在其上执行的设备（应为 GPU）。将默认为模型的第一个参数设备。
  • offload_buffers (bool, 可选, 默认为 False) — 是否卸载带有模型参数的缓冲区。
  • preload_module_classes (List[str], 可选) — 一个类列表，这些类的实例应在正向传递开始时加载其所有权重（即使在子模块中）。这应该仅用于具有已注册但未在正向传递期间直接调用的子模块的类，例如，如果注册了 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, 可选) — 主要执行设备。 将默认为 device_map 中第一个不同于 "cpu" 或 "disk" 的设备。
  • state_dict (Dict[str, torch.Tensor], 可选) — 将保留在 CPU 上的模型部分的 state dict。
  • 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 甚至磁盘上。

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 )

参数

  • 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 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], 可选) — 类实例应在其前向传播开始时加载所有权重（甚至在子模块中）的类列表。 这应仅用于具有已注册但未在前向传播期间直接调用的子模块的类，例如，如果注册了 dense 线性层，但在前向传播时，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"]
... )

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 )

参数

  • 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 dtype 中保留的模块列表。
  • offload_8bit_bnb (bool, 可选) — 是否启用在 cpu/磁盘上卸载 8 位模块。
  • strict (bool, 可选, 默认为 False) — 是否严格强制检查点 state_dict 中的键与模型 state_dict 的键匹配。

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

一旦跨设备加载，您仍然需要在您的模型上调用 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]], optional) — 如果提供，则为某些特定权重考虑的特殊数据类型（将覆盖用作所有权重默认数据类型的 dtype）。
  • verbose (bool, optional, defaults to False) — 是否提供调试语句，以帮助函数构建 device_map。（默认为 False）
  • clean_result (bool, optional, defaults to True) — 通过将所有放置在同一设备上的子模块分组在一起，来清理生成的 device_map。（默认为 True）
  • offload_buffers (bool, optional, defaults to False) — 在 CPU 或硬盘上卸载的层中，是否也卸载缓冲区以及参数。（默认为 False）
  • fallback_allocation (bool, optional, defaults to False) — 当常规分配失败时，尝试使用 BFS 分配一个符合大小限制的模块。（默认为 False）

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

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

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

Hooks

ModelHook

class accelerate.hooks.ModelHook

< >

( )

一个 hook，包含在模型的前向传播方法之前和之后执行的回调函数。与 PyTorch 现有 hook 的区别在于，它们会传递 kwargs。

类属性

  • no_grad (bool, optional, defaults to False) — 是否在 torch.no_grad() 上下文管理器下执行实际的前向传播。（默认为 False）

detach_hook

< >

( module )

参数

  • module (torch.nn.Module) — 从此 hook 分离的模块。

当 hook 从模块分离时执行。

init_hook

< >

( module )

参数

  • module (torch.nn.Module) — 附加到此 hook 的模块。

当 hook 附加到模块时执行。

post_forward

< >

( module output ) → Any

参数

  • module (torch.nn.Module) — 其前向传播在此事件之前刚刚执行的模块。
  • output (Any) — 模块的输出。

返回值

Any

处理后的 output。

在模型的前向传播方法之后立即执行。

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 的元组。

在模型的前向传播方法之前立即执行。

AlignDevicesHook

class accelerate.hooks.AlignDevicesHook

< >

( execution_device: typing.Union[str, torch.device, int, 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, optional) — 输入和模型权重在前向传播之前应放置在其上的设备。（可选）
  • offload (bool, optional, defaults to False) — 是否在前向传播后卸载权重。（默认为 False）
  • io_same_device (bool, optional, defaults to False) — 输出是否应放置在与输入相同的设备上。（默认为 False）
  • weights_map (Mapping[str, torch.Tensor], optional) — 当模型权重被卸载时，从参数名称到张量值的（可能是惰性的）映射。（可选）
  • offload_buffers (bool, optional, defaults to False) — 卸载时是否包含关联模块的缓冲区。（默认为 False）
  • place_submodules (bool, optional, defaults to False) — 是否在 init_hook 事件期间将子模块放置在 execution_device 上。（默认为 False）

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

SequentialHook

class accelerate.hooks.SequentialHook

< >

( *hooks )

一个可以包含多个 hook 并在每个事件中迭代它们的 hook。

添加钩子

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) — The module where we want to attach the hooks.
  • execution_device (torch.device or Dict[str, torch.device], optional) — The device on which inputs and model weights should be placed before the forward pass. It can be one device for the whole module, or a dictionary mapping module name to device.
  • offload (bool, optional, defaults to False) — Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole module, or a dictionary mapping module name to boolean.
  • weights_map (Mapping[str, torch.Tensor], optional) — When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
  • offload_buffers (bool, optional, defaults to False) — Whether or not to include the associated module’s buffers when offloading.
  • module_name (str, optional, defaults to "") — The name of the module.
  • skip_keys (str or List[str], optional) — A list of keys to ignore when moving inputs or outputs between devices.
  • preload_module_classes (List[str], optional) — A list of classes whose instances should load all their weights (even in the submodules) at the beginning of the forward. This should only be used for classes that have submodules which are registered but not called directly during the forward, for instance if a dense linear layer is registered, but at forward, dense.weight and dense.bias are used in some operations instead of calling dense directly.
  • tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], optional, defaults to None) — A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution device, this parameter is useful to reuse the first available pointer of a shared weight for all others, instead of duplicating memory.

Attaches AlignDevicesHook to all blocks of a given model as needed.

移除 Hooks

remove_hook_from_module

accelerate.hooks.remove_hook_from_module

< >

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

参数

  • module (torch.nn.Module) — The module to attach a hook to.
  • recurse (bool, optional) — Whether to remove the hooks recursively

返回值

torch.nn.Module

The same module, with the hook detached (the module is modified in place, so the result can be discarded).

Removes any hook attached to a module via add_hook_to_module.

remove_hook_from_submodules

accelerate.hooks.remove_hook_from_submodules

< >

( module: Module )

参数

  • module (torch.nn.Module) — The module on which to remove all hooks.

Recursively removes all hooks attached on the submodules of a given model.

实用工具

has_offloaded_params

accelerate.utils.has_offloaded_params

< >

( module: Module ) → bool

参数

  • module (torch.nn.Module) — The module to check for an offload hook.

返回值

bool

True if the module has an offload hook and offloading is enabled, False otherwise.

Checks if a module has offloaded parameters by checking if the given module has a AlignDevicesHook attached with offloading enabled

align_module_device

accelerate.utils.align_module_device

< >

( module: Module execution_device: typing.Optional[torch.device] = None )

参数

  • module (torch.nn.Module) — Module with parameters to align.
  • execution_device (torch.device, optional) — If provided, overrides the module’s execution device within the context. Otherwise, use hook execution device or pass

Context manager that moves a module’s parameters to the specified execution device.

< > Update on GitHub

←Logging Pipeline parallelism→