Accelerate documentation
处理大型模型
处理大型模型
分发与卸载
init_empty_weights
accelerate.init_empty_weights
< source >( include_buffers: bool | None = None )
一个上下文管理器,在该管理器下,模型的所有参数都会在 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
< source >( 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
< source >( model: Moduleexecution_device: int | str | torch.device | None = Noneprev_module_hook: accelerate.hooks.UserCpuOffloadHook | None = None )
将模型卸载到 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
< source >( 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
< source >( 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
< source >( 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
< source >( 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
< source >( 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
一个包含回调的钩子,这些回调将在模型的前向方法执行之前和之后运行。与 PyTorch 中现有的钩子不同,它们会接收并传递 kwargs。
类属性:
- no_grad (
bool,可选,默认为False) — 是否在torch.no_grad()上下文管理器下执行实际的前向传播。
在钩子从模块上分离时执行。
在钩子附加到模块时执行。
post_forward
< source >( moduleoutput ) → Any
在模型的前向方法执行后立即执行。
pre_forward
< source >( module*args**kwargs ) → Tuple[Tuple[Any], Dict[Str, Any]]
在模型的前向方法执行之前立即执行。
AlignDevicesHook
类 accelerate.hooks.AlignDevicesHook
< source >( 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
一个可以包含多个钩子并在每个事件时遍历它们的钩子。
LayerwiseCastingHook
类 accelerate.hooks.LayerwiseCastingHook
< source >( storage_dtype: dtypecompute_dtype: dtypenon_blocking: bool )
一个钩子,用于将模块的权重转换为高精度数据类型以进行计算,并转换为低精度数据类型以进行存储。此过程可能会导致输出质量下降,但可以显著减少内存占用。
添加钩子
add_hook_to_module
accelerate.hooks.add_hook_to_module
< source >( module: Modulehook: ModelHookappend: bool = False ) → torch.nn.Module
为指定模块添加一个钩子。这将重写该模块的 forward 方法以包含该钩子。若要移除此行为并恢复原始的 forward 方法,请使用 remove_hook_from_module。
如果模块中已存在一个钩子,默认情况下会用传入的新钩子将其替换。要将两个钩子串联起来,请传入
append=True,这样会将当前钩子和新钩子链接为SequentialHook类的一个实例。
attach_execution_device_hook
accelerate.hooks.attach_execution_device_hook
< source >( 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
< source >( 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
< source >( 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
< source >( 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(...)用户还可以传入希望避免被降级的模块。
移除钩子
remove_hook_from_module
accelerate.hooks.remove_hook_from_module
< source >( module: Modulerecurse = False ) → torch.nn.Module
移除通过 add_hook_to_module 附加到模块的任何钩子。
remove_hook_from_submodules
accelerate.hooks.remove_hook_from_submodules
< source >( module: Module )
递归移除给定模型的所有子模块上附加的所有钩子。
实用工具
has_offloaded_params
accelerate.utils.has_offloaded_params
< source >( module: Module ) → bool
通过检查给定模块是否附加了启用卸载的 AlignDevicesHook,来检查该模块是否具有已卸载的参数
align_module_device
accelerate.utils.align_module_device
< source >( module: Moduleexecution_device: torch.device | None = None )
上下文管理器,用于将模块的参数移动到指定的执行设备。