Accelerate 文档

处理大型模型

Accelerate

加入 Hugging Face 社区

并获取增强的文档体验

在模型、数据集和 Spaces 上协作

通过加速推理获得更快的示例

在文档主题之间切换

开始使用

处理大型模型

分发和卸载

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

cpu_offload

accelerate.cpu_offload

< 源代码 >

( model: Module execution_device: Optional = None offload_buffers: bool = False state_dict: Optional = None preload_module_classes: Optional = 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], 可选) — 一个类列表，这些类的实例应在前向传播开始时加载其所有权重（即使在子模块中也是如此）。这应仅用于具有已注册但未在前向传播期间直接调用的子模块的类，例如，如果注册了一个 dense 线性层，但在前向传播时，dense.weight 和 dense.bias 用于某些操作，而不是直接调用 dense。

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

cpu_offload_with_hook

accelerate.cpu_offload_with_hook

< 源代码 >

( model: Module execution_device: Union = None prev_module_hook: Optional = 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: Union execution_device: Optional = None offload_buffers: bool = False preload_module_classes: Optional = 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

< source >

( model: Module device_map: Dict main_device: Optional = None state_dict: Optional = None offload_dir: Union = None offload_index: Optional = None offload_buffers: bool = False skip_keys: Union = None preload_module_classes: Optional = None force_hooks: bool = False )

参数

model (torch.nn.Module) — 要调度的模型。
device_map (Dict[str, Union[str, int, torch.device]]) — 一个字典，将模型 state_dict 中的模块名称映射到它们应前往的设备。请注意，即使 "disk" 不是 torch.device 的正确值，也接受 "disk"。
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

< source >

( model: Module checkpoint: Union device_map: Union = None max_memory: Optional = None no_split_module_classes: Optional = None offload_folder: Union = None offload_buffers: bool = False dtype: Union = None offload_state_dict: Optional = None skip_keys: Union = None preload_module_classes: Optional = None force_hooks: bool = False strict: 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], 可选) — 类列表，其实例应在正向传播开始时加载其所有权重（即使在子模块中）。这仅应用于具有已注册但未在正向传播期间直接调用的子模块的类，例如，如果注册了 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

< source >

( model: Module checkpoint: Union device_map: Optional = None offload_folder: Union = None dtype: Union = 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

< source >

( model: Module max_memory: Optional = None no_split_module_classes: Optional = None dtype: Union = None special_dtypes: Optional = None verbose: bool = False clean_result: bool = True offload_buffers: 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]], 可选) — 如果提供，则为某些特定权重考虑的特殊数据类型（将覆盖用于所有权重的默认数据类型）。
verbose (bool, 可选, 默认为 False) — 是否在函数构建 device_map 时提供调试语句。
clean_result (bool, 可选, 默认为 True) — 通过将所有位于同一设备上的子模块分组在一起，来清理生成的 device_map。
offload_buffers (bool, 可选, 默认为 False) — 在 CPU 或硬盘上卸载的层中，是否也卸载缓冲区以及参数。

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

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

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

钩子

ModelHook

class accelerate.hooks.ModelHook

< source >

( )

一个包含回调函数的钩子，用于在模型的前向传播方法之前和之后执行。与 PyTorch 现有钩子的不同之处在于，它们会沿着 kwargs 传递。

类属性

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

detach_hook

< source >

( module )

参数

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

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

init_hook

< source >

( module )

参数

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

当钩子附加到模块时执行。

post_forward

< source >

( module output ) → Any

参数

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

返回值

Any

处理后的 output。

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

pre_forward

< source >

( 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

< source >

( 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 (bool, 可选, 默认为 False) — 卸载时是否包含关联模块的缓冲区。
place_submodules (bool, 可选, 默认为 False) — 是否在 init_hook 事件期间将子模块放置在 execution_device 上。

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

class accelerate.hooks.SequentialHook
< source >
( *hooks )

A hook that can contain several hooks and iterates through them at each event.

Adding Hooks

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, optional, defaults to False) — 指示是否应将此钩子与现有钩子（如果模块已包含钩子）链接在一起。(可选，默认为 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: Union skip_keys: Union = None preload_module_classes: Optional = None tied_params_map: Optional = None )

参数

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

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

attach_align_device_hook

accelerate.hooks.attach_align_device_hook

< 源代码 >

( module: Module execution_device: Optional = None offload: bool = False weights_map: Optional = None offload_buffers: bool = False module_name: str = '' skip_keys: Union = None preload_module_classes: Optional = None tied_params_map: Optional = None )

参数

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

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

attach_align_device_hook_on_blocks

accelerate.hooks.attach_align_device_hook_on_blocks

< 源代码 >

( module: Module execution_device: Union = None offload: Union = False weights_map: Mapping = None offload_buffers: bool = False module_name: str = '' skip_keys: Union = None preload_module_classes: Optional = None tied_params_map: Optional = 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], 可选) — 一个类列表，这些类的实例应在正向传播开始时加载其所有权重（即使在子模块中）。这应该仅用于具有已注册但未在正向传播期间直接调用的子模块的类，例如，如果注册了一个 dense 线性层，但在正向传播时，使用了 dense.weight 和 dense.bias 在某些操作中，而不是直接调用 dense。
tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], 可选, 默认为 None) — 数据指针到设备字典的映射，其中设备字典映射到已分发的绑定权重。对于给定的执行设备，此参数对于为所有其他共享权重重用第一个可用指针而不是复制内存很有用。

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

移除钩子

remove_hook_from_module

accelerate.hooks.remove_hook_from_module

< source >

( module: Module recurse = 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

< source >

( module: Module )

参数

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

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

< > 更新在 GitHub 上

←日志记录流水线并行→

Accelerate

处理大型模型

分发和卸载

init_empty_weights

accelerate.init_empty_weights

cpu_offload

accelerate.cpu_offload

cpu_offload_with_hook

accelerate.cpu_offload_with_hook

disk_offload

accelerate.disk_offload

dispatch_model

accelerate.dispatch_model

load_checkpoint_and_dispatch

accelerate.load_checkpoint_and_dispatch

load_checkpoint_in_model

accelerate.load_checkpoint_in_model

infer_auto_device_map

accelerate.infer_auto_device_map

钩子

ModelHook

class accelerate.hooks.ModelHook

detach_hook

init_hook

post_forward

pre_forward

AlignDevicesHook

class accelerate.hooks.AlignDevicesHook

class accelerate.hooks.SequentialHook < source > ( *hooks ) A hook that can contain several hooks and iterates through them at each event.

class accelerate.hooks.SequentialHook

Adding Hooks

add_hook_to_module

accelerate.hooks.add_hook_to_module

attach_execution_device_hook

accelerate.hooks.attach_execution_device_hook

attach_align_device_hook

accelerate.hooks.attach_align_device_hook

attach_align_device_hook_on_blocks

accelerate.hooks.attach_align_device_hook_on_blocks

移除钩子

remove_hook_from_module

accelerate.hooks.remove_hook_from_module

remove_hook_from_submodules

accelerate.hooks.remove_hook_from_submodules

class accelerate.hooks.SequentialHook
< source >
( *hooks )

A hook that can contain several hooks and iterates through them at each event.