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 设备上。

一个上下文管理器,在该管理器下,模型的所有参数都会在 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。

为模型启用完整的 CPU 卸载。因此,模型的所有参数都将被卸载,并且仅保留一份模型状态字典的副本。在前向传播过程中,参数将从该状态字典中提取出来,并在需要时放置到作为参数传入的执行设备上,随后再次被卸载。

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 或 torch.device,可选) — 模型应在其上运行的设备。如果可用,默认使用 MPS 设备;否则,如果存在加速器设备,则使用设备 0;最后回退到 CPU。
  • prev_module_hook (UserCpuOffloadHook,可选) — 此函数为流水线中前一个模型返回的钩子。如果传入该参数,其 offload 方法将在挂载此钩子的模型执行 forward 之前被调用。

将模型卸载到 CPU,并在执行时将其放回执行设备。与 cpu_offload() 的区别在于,前向传播后模型会保留在执行设备上,只有当调用返回的 hook 的 offload 方法时才会再次卸载。适用于在循环中运行模型的流水线。

示例:

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。

为模型启用完整的磁盘卸载。因此,模型的所有参数都将作为内存映射数组卸载到指定文件夹中。在前向传播过程中,参数将从该文件夹中访问,并在需要时以原样放置到传入的执行设备上,随后再次被卸载。

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 上的模型部分的 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], 可选) — 一个类列表,其实例应在 forward 开始时加载所有权重(包括子模块中的权重)。此参数仅应用于那些在 forward 期间注册了子模块但未直接调用它们的类,例如,如果注册了一个 dense 线性层,但在 forward 中,某些操作使用的是 dense.weight 和 dense.bias,而不是直接调用 dense。
  • force_hooks (bool,可选,默认为 False) — 是否强制将设备钩子附加到模型上,即使所有层都调度到单个设备上。

根据给定的设备映射来调度模型。模型的层可能分布在多个 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) — 要加载的文件夹检查点。它可以是:
    • 指向包含完整模型状态字典的文件的路径
    • 指向包含分片检查点索引的 .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 或硬盘上的层,是否同时卸载缓冲区以及参数。
  • dtype (str 或 torch.dtype,可选) — 如果提供,加载时权重将被转换为该类型。
  • offload_state_dict (bool, 可选) — 如果为 True,则会将 CPU 状态字典临时卸载到硬盘上,以避免当 CPU 状态字典的权重加上最大分片无法放入 CPU 内存时发生内存不足。如果所选的设备映射包含 "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: 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) — 要加载的文件夹检查点。它可以是:
    • 指向包含完整模型状态字典的文件的路径
    • 指向包含分片检查点索引的 .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, 可选) — 是否启用将 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 中的张量逐个广播到其他 rank。其他 rank 将接收这些张量,并根据模型中的本地分片进行分片(如适用)。

在模型中加载(可能分片的)检查点,并在加载时将权重发送到指定设备。

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]], 可选) — 如果提供,则指定某些特定权重需要考虑的特殊数据类型(将覆盖用于所有权重的默认 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 上占用空间最大的层重新加载回 GPU。
  • 如果需要将数据卸载到 CPU,我们不会超出 CPU 上可用的 RAM。
  • 如果需要将层卸载到磁盘,CPU 上始终会保留足够的空间,以便将已卸载到磁盘的层中体积最大的那一层重新加载回 CPU。

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

钩子

ModelHook

类 accelerate.hooks.ModelHook

< >

( )

一个包含回调的钩子,这些回调将在模型的前向方法执行之前和之后运行。与 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。

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

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

类 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) — 是否在 init_hook 事件期间将子模块放置在 execution_device 上。

一个通用的 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 附加到给定模型中所有具有直接参数和/或缓冲区(buffers)的子模块。

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) — 该模块的叶子模块将在计算时转换为高精度 dtype,在存储时转换为低精度 dtype。
  • storage_dtype (torch.dtype) — 用于在 forward pass 前后将模块转换以进行存储的 dtype。
  • compute_dtype (torch.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 hook 的模块。

返回值

bool

如果模块具有 offload 钩子且 offloading 已启用,则返回 True,否则返回 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