Accelerator

Accelerator 是用于在任何类型的训练设置中启用分布式训练的主类。请阅读将 Accelerator 添加到您的代码中教程，以了解有关如何将 Accelerator 添加到脚本的更多信息。

Accelerator

class accelerate.Accelerator

参数

device_placement (bool，可选，默认为 True) — Accelerator 是否应将对象（数据加载器生成的张量、模型等）放置在设备上。
mixed_precision (str，可选) — 是否使用混合精度训练。可选择 'no'、'fp16'、'bf16' 或 'fp8'。将默认为环境变量 ACCELERATE_MIXED_PRECISION 中的值，该值将使用当前系统的 Accelerate 配置中的默认值或通过 accelerate.launch 命令传递的标志。'fp8' 需要安装 transformers-engine。
gradient_accumulation_steps (int，可选，默认为 1) — 在累积梯度之前应经过的步骤数。大于 1 的数字应与 Accelerator.accumulate 结合使用。如果未传递，将默认为环境变量 ACCELERATE_GRADIENT_ACCUMULATION_STEPS 中的值。也可以通过 GradientAccumulationPlugin 进行配置。
cpu (bool，可选) — 是否强制脚本在 CPU 上执行。如果设置为 True，将忽略可用的 GPU，并强制在单个进程上执行。
dataloader_config (DataLoaderConfiguration，可选) — 用于配置在分布式场景中如何处理数据加载器的配置。
deepspeed_plugin (DeepSpeedPlugin 或 str 的字典 — DeepSpeedPlugin，可选)：使用此参数调整与 DeepSpeed 相关的参数。此参数是可选的，可以直接使用 *accelerate config* 进行配置。如果使用多个插件，请使用每个插件配置的 key 属性从 accelerator.state.get_deepspeed_plugin(key) 访问它们。是 deepspeed_plugins 的别名。
fsdp_plugin (FullyShardedDataParallelPlugin，可选) — 使用此参数调整与 FSDP 相关的参数。此参数是可选的，可以直接使用 *accelerate config* 进行配置。
torch_tp_plugin (TorchTensorParallelPlugin，可选) — 已弃用：请改用带有 tp_size 的 parallelism_config。
megatron_lm_plugin (MegatronLMPlugin，可选) — 使用此参数调整与 MegatronLM 相关的参数。此参数是可选的，可以直接使用 *accelerate config* 进行配置。
rng_types (str 或 RNGType 列表) — 在您准备的数据加载器的每次迭代开始时要同步的随机数生成器列表。应为以下一项或多项：
- "torch"：基础 torch 随机数生成器
- "cuda"：CUDA 随机数生成器（仅限 GPU）
- "xla"：XLA 随机数生成器（仅限 TPU）
- "generator"：采样器的 torch.Generator（如果数据加载器中没有采样器，则为批采样器），或者如果底层数据集是该类型，则是可迭代数据集的 torch.Generator（如果存在）。
对于 PyTorch 版本 <=1.5.1，将默认为 ["torch"]，对于 PyTorch 版本 >= 1.6，将默认为 ["generator"]。
log_with (str、LoggerType 或 GeneralTracker 列表，可选) — 为实验跟踪设置的日志记录器列表。应为以下一项或多项：
- "all"
- "tensorboard"
- "wandb"
- "trackio"
- "aim"
- "comet_ml"
- "mlflow"
- "dvclive"
- "swanlab" 如果选择 "all"，将选择环境中所有可用的跟踪器并初始化它们。也可以接受 GeneralTracker 的实现以用于自定义跟踪器，并且可以与 "all" 结合使用。
project_config (ProjectConfiguration，可选) — 用于配置如何处理状态保存的配置。
project_dir (str，os.PathLike，可选) — 用于存储数据（如本地兼容日志记录器的日志和可能保存的检查点）的目录路径。
step_scheduler_with_optimizer (bool，可选，默认为 True) — 如果学习率调度器与优化器同时更新，则设置为 True；如果仅在特定情况下（例如每个 epoch 结束时）更新，则设置为 False。
kwargs_handlers (KwargsHandler 列表，可选) — KwargsHandler 列表，用于自定义如何创建与分布式训练、性能分析或混合精度相关的对象。有关更多信息，请参阅 kwargs。
dynamo_backend (str 或 DynamoBackend，可选，默认为 "no") — 设置为可能的 dynamo 后端之一，以使用 torch dynamo 优化您的训练。
dynamo_plugin (TorchDynamoPlugin，可选) — 用于配置如何处理 torch dynamo 的配置，如果需要的调整不仅仅是 backend 或 mode。
gradient_accumulation_plugin (GradientAccumulationPlugin，可选) — 用于配置如何处理梯度累积的配置，如果需要的调整不仅仅是 gradient_accumulation_steps。

为分布式训练或混合精度训练创建 accelerator 的实例。

可用属性

device (torch.device) — 要使用的设备。
distributed_type (DistributedType) — 分布式训练配置。
local_process_index (int) — 当前机器上的进程索引。
mixed_precision (str) — 配置的混合精度模式。
num_processes (int) — 用于训练的总进程数。
optimizer_step_was_skipped (bool) — 优化器更新是否被跳过（因为混合精度中的梯度溢出），在这种情况下学习率不应改变。
process_index (int) — 当前进程在所有进程中的总索引。
state (AcceleratorState) — 分布式设置状态。
sync_gradients (bool) — 梯度当前是否在所有进程中同步。
use_distributed (bool) — 当前配置是否用于分布式训练。

accumulate

< 源码 >

( *models )

参数

*models (torch.nn.Module 列表) — 使用 Accelerator.prepare 准备的 PyTorch 模块。传递给 accumulate() 的模型将在分布式训练的反向传播过程中跳过梯度同步。

一个上下文管理器，它将轻量级地包装并自动执行梯度累积。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator(gradient_accumulation_steps=1)
>>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)

>>> for input, output in dataloader:
...     with accelerator.accumulate(model):
...         outputs = model(input)
...         loss = loss_func(outputs)
...         loss.backward()
...         optimizer.step()
...         scheduler.step()
...         optimizer.zero_grad()

autocast

< 源码 >

( autocast_handler: AutocastKwargs = None )

如果启用了自动混合精度，将在此上下文管理器内的块中应用它。否则不会发生任何变化。

可以传入一个不同的 autocast_handler 来覆盖在 Accelerator 对象中设置的那个。这在 autocast 下的块中非常有用，当你想要恢复到 fp32 时。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator(mixed_precision="fp16")
>>> with accelerator.autocast():
...     train()

backward

< 源码 >

( loss **kwargs )

根据 GradientAccumulationPlugin 缩放梯度，并根据配置调用正确的 backward()。

应该用来替代 loss.backward()。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator(gradient_accumulation_steps=2)
>>> outputs = model(inputs)
>>> loss = loss_fn(outputs, labels)
>>> accelerator.backward(loss)

check_trigger

< 源码 >

( )

检查内部触发张量在任何进程中是否被设置为 1。如果是，将返回 True 并将触发张量重置为 0。

注意：不需要 wait_for_everyone()

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume later in the training script
>>> # `should_do_breakpoint` is a custom function to monitor when to break,
>>> # e.g. when the loss is NaN
>>> if should_do_breakpoint(loss):
...     accelerator.set_trigger()
>>> # Assume later in the training script
>>> if accelerator.check_trigger():
...     break

clear

< 源码 >

( *objects )

Accelerate.free_memory 的别名，释放对存储的内部对象的所有引用并调用垃圾回收器。您应该在两次使用不同模型/优化器的训练之间调用此方法。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer, scheduler = ...
>>> model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler)
>>> model, optimizer, scheduler = accelerator.clear(model, optimizer, scheduler)

clip_grad_norm_

< 源码 >

( parameters max_norm norm_type = 2 ) → torch.Tensor

torch.Tensor

参数梯度的总范数（视为单个向量）。

应该用来代替 torch.nn.utils.clip_grad_norm_。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator(gradient_accumulation_steps=2)
>>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)

>>> for input, target in dataloader:
...     optimizer.zero_grad()
...     output = model(input)
...     loss = loss_func(output, target)
...     accelerator.backward(loss)
...     if accelerator.sync_gradients:
...         accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)
...     optimizer.step()

clip_grad_value_

< 源码 >

( parameters clip_value )

应该用来代替 torch.nn.utils.clip_grad_value_。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator(gradient_accumulation_steps=2)
>>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)

>>> for input, target in dataloader:
...     optimizer.zero_grad()
...     output = model(input)
...     loss = loss_func(output, target)
...     accelerator.backward(loss)
...     if accelerator.sync_gradients:
...         accelerator.clip_grad_value_(model.parameters(), clip_value)
...     optimizer.step()

end_training

< 源码 >

( )

运行任何特殊的训练结束行为，例如仅在主进程上停止跟踪器或销毁进程组。如果使用实验跟踪，应始终在脚本结束时调用。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator(log_with="tensorboard")
>>> accelerator.init_trackers("my_project")
>>> # Do training
>>> accelerator.end_training()

free_memory

< 源码 >

( *objects )

将释放对存储的内部对象的所有引用并调用垃圾回收器。您应该在两次使用不同模型/优化器的训练之间调用此方法。同时会将 Accelerator.step 重置为 0。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer, scheduler = ...
>>> model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler)
>>> model, optimizer, scheduler = accelerator.free_memory(model, optimizer, scheduler)

gather

< 源码 >

( tensor ) → torch.Tensor，或嵌套的元组/列表/字典的 torch.Tensor

参数

tensor (torch.Tensor，或嵌套的元组/列表/字典的 torch.Tensor) — 要在所有进程中收集的张量。

torch.Tensor，或嵌套的元组/列表/字典的 torch.Tensor

收集到的张量。注意，结果的第一维是 *num_processes* 乘以输入张量的第一维。

在所有进程中收集 *tensor* 中的值，并在第一个维度上进行连接。在进行评估时，用于重新组合所有进程的预测非常有用。

注意：此收集操作在所有进程中进行。

示例

>>> # Assuming four processes
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> process_tensor = torch.tensor([accelerator.process_index])
>>> gathered_tensor = accelerator.gather(process_tensor)
>>> gathered_tensor
tensor([0, 1, 2, 3])

gather_for_metrics

< 源码 >

( input_data use_gather_object = False )

参数

input (torch.Tensor、object、嵌套的元组/列表/字典的 torch.Tensor，或嵌套的元组/列表/字典的 object) — 用于在所有进程中计算指标的张量或对象。
use_gather_object(bool) — 是否强制使用 gather_object 而不是 gather（如果传递的所有对象都不包含张量，则已经这样做了）。此标志对于收集不同大小的张量非常有用，因为我们不想沿着第一个维度进行填充和连接。与 GPU 张量一起使用时支持不佳且效率低下，因为它会因张量被 pickle 化而导致 GPU -> CPU 的传输。

收集 input_data，并在分布式系统上可能删除最后一个批次中的重复项。应用于收集用于指标计算的输入和目标。

示例

>>> # Assuming two processes, with a batch size of 5 on a dataset with 9 samples
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> dataloader = torch.utils.data.DataLoader(range(9), batch_size=5)
>>> dataloader = accelerator.prepare(dataloader)
>>> batch = next(iter(dataloader))
>>> gathered_items = accelerator.gather_for_metrics(batch)
>>> len(gathered_items)
9

get_state_dict

< 源码 >

( model unwrap = True ) → dict

参数

model (torch.nn.Module) — 通过 Accelerator.prepare() 传递的 PyTorch 模型。
unwrap (bool，可选，默认为 True) — 是否返回 model 的原始底层 state_dict，还是返回包装后的 state_dict。

字典

可能不包含完整精度的模型的 state_dict。

返回通过 Accelerator.prepare() 传递的模型的 state_dict，可能不包含完整精度。

示例

>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> net = torch.nn.Linear(2, 2)
>>> net = accelerator.prepare(net)
>>> state_dict = accelerator.get_state_dict(net)

get_tracker

< 源码 >

( name: str unwrap: bool = False ) → GeneralTracker

参数

name (str) — 跟踪器的名称，对应于 .name 属性。
unwrap (bool) — 是返回内部跟踪机制还是返回包装后的跟踪器（推荐）。

GeneralTracker

如果存在，则返回与 name 对应的跟踪器。

仅在主进程中，根据 name 从 self.trackers 返回一个 tracker。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator(log_with="tensorboard")
>>> accelerator.init_trackers("my_project")
>>> tensorboard_tracker = accelerator.get_tracker("tensorboard")

join_uneven_inputs

< 源码 >

( joinables even_batches = None )

参数

joinables (list[torch.distributed.algorithms.Joinable]) — 一个模型或优化器的列表，它们是 torch.distributed.algorithms.Joinable 的子类。最常见的是一个使用 Accelerator.prepare 准备用于 DistributedDataParallel 训练的 PyTorch Module。
even_batches (bool, 可选) — 如果设置，将覆盖在 Accelerator 中设置的 even_batches 的值。如果未提供，将使用默认的 Accelerator 值。

一个上下文管理器，它有助于在不均匀输入上进行分布式训练或评估，作为 torch.distributed.algorithms.join 的包装器。当总批次大小不能被数据集长度整除时，这很有用。

join_uneven_inputs 仅支持在多个 GPU 上进行分布式数据并行训练。对于任何其他配置，此方法将没有效果。

覆盖 even_batches 不会影响可迭代式数据加载器。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator(even_batches=True)
>>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

>>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
...     for input, output in dataloader:
...         outputs = model(input)
...         loss = loss_func(outputs)
...         loss.backward()
...         optimizer.step()
...         optimizer.zero_grad()

load_state

< 源码 >

( input_dir: str = None load_kwargs: dict | None = None **load_model_func_kwargs )

参数

input_dir (str 或 os.PathLike) — 所有相关权重和状态保存的文件夹名称。如果使用了 automatic_checkpoint_naming，则可以为 None，并将从最新的检查点加载。
load_kwargs (dict, 可选) — 传递给底层 load 函数的附加关键字参数，例如 state_dict 和 optimizer 的可选参数。
load_model_func_kwargs (dict, 可选) — 用于加载模型的附加关键字参数，可以传递给底层加载函数，例如 DeepSpeed 的 load_checkpoint 函数的可选参数或用于加载模型和优化器的 map_location。

加载模型、优化器、缩放器、RNG 生成器和已注册对象的当前状态。

应仅与 Accelerator.save_state() 结合使用。如果某个文件未注册用于检查点，则即使它存储在目录中也不会被加载。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer, lr_scheduler = ...
>>> model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
>>> accelerator.load_state("my_checkpoint")

local_main_process_first

< 源码 >

( )

让本地主进程先进入 with 块。

其他进程将在主进程退出后进入 with 块。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> with accelerator.local_main_process_first():
...     # This will be printed first by local process 0 then in a seemingly
...     # random order by the other processes.
...     print(f"This will be printed by process {accelerator.local_process_index}")

lomo_backward

< 源码 >

( loss: torch.Tensor learning_rate: float )

在 LOMO 优化器上运行反向传播。

main_process_first

< 源码 >

( )

让主进程先进入 with 块。

其他进程将在主进程退出后进入 with 块。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> with accelerator.main_process_first():
...     # This will be printed first by process 0 then in a seemingly
...     # random order by the other processes.
...     print(f"This will be printed by process {accelerator.process_index}")

maybe_context_parallel

< 源码 >

( buffers: list[torch.Tensor] | None = None buffer_seq_dims: list[int] | None = None no_restore_buffers: set[torch.Tensor] | None = None )

参数

buffers (list[torch.Tensor], 可选) — 将沿序列维度进行分片的缓冲区。常见示例包括输入、标签或位置嵌入缓冲区。此上下文管理器将就地修改这些缓冲区，退出上下文后，缓冲区将恢复到其原始状态。为避免不必要的恢复，您可以使用 no_restore_buffers 指定哪些缓冲区不需要恢复。
buffer_seq_dims (list[int], 可选) — buffers 的序列维度。
no_restore_buffers (set[torch.Tensor], 可选) — 此集合必须是 buffers 的子集。指定在上下文退出后，buffers 参数中的哪些缓冲区将不会被恢复。这些缓冲区将保持分片状态。

一个启用上下文并行训练的上下文管理器。

context_parallel 目前仅支持与 FSDP2 一起使用，并且需要 parallelism_config.cp_size >

如果这些条件中的任何一个不满足，此上下文管理器将没有效果，但为了减少代码更改，它不会引发异常。

此上下文管理器必须在每个训练步骤中重新创建，如下面的示例所示。

示例

>>> for batch in dataloader:
...     with accelerator.maybe_context_parallel(
...         buffers=[batch["input_ids"], batch["attention_mask"]],
...         buffer_seq_dims=[1, 1],
...         no_restore_buffers={batch["input_ids"]},
...     ):
...         outputs = model(batch)
...         ...

no_sync

< 源码 >

( model )

参数

model (torch.nn.Module) — 使用 Accelerator.prepare 准备的 PyTorch Module

一个上下文管理器，通过调用 torch.nn.parallel.DistributedDataParallel.no_sync 来禁用 DDP 进程间的梯度同步。

如果 model 不处于 DDP 模式，此上下文管理器不执行任何操作。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer)
>>> input_a = next(iter(dataloader))
>>> input_b = next(iter(dataloader))

>>> with accelerator.no_sync():
...     outputs = model(input_a)
...     loss = loss_func(outputs)
...     accelerator.backward(loss)
...     # No synchronization across processes, only accumulate gradients
>>> outputs = model(input_b)
>>> accelerator.backward(loss)
>>> # Synchronization across all processes
>>> optimizer.step()
>>> optimizer.zero_grad()

on_last_process

< 源码 >

( function: Callable[..., Any] )

参数

function (Callable) — 要装饰的函数。

一个装饰器，它将仅在最后一个进程上运行被装饰的函数。也可以使用 PartialState 类调用。

示例

# Assume we have 4 processes.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_last_process
def print_something():
    print(f"Printed on process {accelerator.process_index}")


print_something()
"Printed on process 3"

on_local_main_process

< 源码 >

( function: Callable[..., Any] = None )

参数

function (Callable) — 要装饰的函数。

一个装饰器，它将仅在本地主进程上运行被装饰的函数。也可以使用 PartialState 类调用。

示例

# Assume we have 2 servers with 4 processes each.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_local_main_process
def print_something():
    print("This will be printed by process 0 only on each server.")


print_something()
# On server 1:
"This will be printed by process 0 only"
# On server 2:
"This will be printed by process 0 only"

on_local_process

< 源码 >

( function: Callable[..., Any] = None local_process_index: int = None )

参数

function (Callable, 可选) — 要装饰的函数。
local_process_index (int, 可选) — 运行函数的本地进程的索引。

一个装饰器，它将仅在给定的本地进程索引上运行被装饰的函数。也可以使用 PartialState 类调用。

示例

# Assume we have 2 servers with 4 processes each.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_local_process(local_process_index=2)
def print_something():
    print(f"Printed on process {accelerator.local_process_index}")


print_something()
# On server 1:
"Printed on process 2"
# On server 2:
"Printed on process 2"

on_main_process

< 源码 >

( function: Callable[..., Any] = None )

参数

function (Callable) — 要装饰的函数。

一个装饰器，它将仅在主进程上运行被装饰的函数。也可以使用 PartialState 类调用。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()


>>> @accelerator.on_main_process
... def print_something():
...     print("This will be printed by process 0 only.")


>>> print_something()
"This will be printed by process 0 only"

on_process

< 源码 >

( function: Callable[..., Any] = None process_index: int = None )

参数

function (Callable, 可选) — 要装饰的函数。
process_index (int, 可选) — 运行函数的进程的索引。

一个装饰器，它将仅在给定的进程索引上运行被装饰的函数。也可以使用 PartialState 类调用。

示例

# Assume we have 4 processes.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_process(process_index=2)
def print_something():
    print(f"Printed on process {accelerator.process_index}")


print_something()
"Printed on process 2"

pad_across_processes

< 源码 >

( tensor dim = 0 pad_index = 0 pad_first = False ) → torch.Tensor 或嵌套的元组/列表/字典的 torch.Tensor

参数

tensor (嵌套的列表/元组/字典的 torch.Tensor) — 要收集的数据。
dim (int, 可选, 默认为 0) — 要填充的维度。
pad_index (int, 可选, 默认为 0) — 用来填充的值。
pad_first (bool, 可选, 默认为 False) — 是在开头还是结尾填充。

torch.Tensor，或嵌套的元组/列表/字典的 torch.Tensor

填充后的张量。

递归地将嵌套的列表/元组/字典中的张量从所有设备填充到相同的大小，以便它们可以安全地被收集。

示例

>>> # Assuming two processes, with the first processes having a tensor of size 1 and the second of size 2
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> process_tensor = torch.arange(accelerator.process_index + 1).to(accelerator.device)
>>> padded_tensor = accelerator.pad_across_processes(process_tensor)
>>> padded_tensor.shape
torch.Size([2])

prepare

< 源码 >

( *args device_placement = None )

参数

*args (对象列表) — 以下任何类型的对象：
- torch.utils.data.DataLoader: PyTorch Dataloader
- torch.nn.Module: PyTorch Module
- torch.optim.Optimizer: PyTorch Optimizer
- torch.optim.lr_scheduler.LRScheduler: PyTorch LR Scheduler
device_placement (list[bool], 可选) — 用于自定义是否为每个传递的对象执行自动设备放置。需要是一个与 args 长度相同的列表。与 DeepSpeed 或 FSDP 不兼容。

为分布式训练和混合精度准备 args 中传递的所有对象，然后按相同顺序返回它们。

如果模型仅用于推理且不使用任何混合精度，则无需准备。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume a model, optimizer, data_loader and scheduler are defined
>>> model, optimizer, data_loader, scheduler = accelerator.prepare(model, optimizer, data_loader, scheduler)

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume a model, optimizer, data_loader and scheduler are defined
>>> device_placement = [True, True, False, False]
>>> # Will place the first two items passed in automatically to the right device but not the last two.
>>> model, optimizer, data_loader, scheduler = accelerator.prepare(
...     model, optimizer, data_loader, scheduler, device_placement=device_placement
... )

prepare_data_loader

< 源码 >

( data_loader: torch.utils.data.DataLoader device_placement = None slice_fn_for_dispatch = None )

参数

data_loader (torch.utils.data.DataLoader) — 一个普通的 PyTorch DataLoader，用于准备
device_placement (bool, 可选) — 是否在准备好的 dataloader 中将批次放置在正确的设备上。默认为 self.device_placement。
slice_fn_for_dispatch (Callable, 可选) -- 如果传递，此函数将用于在 num_processes 之间切分张量。默认为 [slice_tensors()](/docs/accelerate/v1.10.0/en/package_reference/utilities#accelerate.utils.slice_tensors)。此参数仅在 dispatch_batches 设置为 True` 时使用，否则将被忽略。

为任何分布式设置中的训练准备一个 PyTorch DataLoader。建议改用 Accelerator.prepare()。

示例

>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> data_loader = torch.utils.data.DataLoader(...)
>>> data_loader = accelerator.prepare_data_loader(data_loader, device_placement=True)

prepare_model

< 源码 >

( model: torch.nn.Module device_placement: bool = None evaluation_mode: bool = False )

参数

model (torch.nn.Module) — 一个 PyTorch 模型，用于准备。如果模型仅用于推理且不使用任何混合精度，则无需准备
device_placement (bool, 可选) — 是否将模型放置在正确的设备上。默认为 self.device_placement。
evaluation_mode (bool, 可选, 默认为 False) — 是否仅为评估设置模型，只应用混合精度和 torch.compile（如果在 Accelerator 对象中配置）。

为任何分布式设置中的训练准备一个 PyTorch 模型。建议改用 Accelerator.prepare()。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume a model is defined
>>> model = accelerator.prepare_model(model)

prepare_optimizer

< 源码 >

( optimizer: torch.optim.Optimizer device_placement = None )

参数

optimizer (torch.optim.Optimizer) — 一个普通的 PyTorch 优化器，用于准备
device_placement (bool, optional) — 是否将优化器放置在正确的设备上。将默认为 self.device_placement。

为任何分布式设置中的训练准备 PyTorch 优化器。建议改用 Accelerator.prepare()。

示例

>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> optimizer = torch.optim.Adam(...)
>>> optimizer = accelerator.prepare_optimizer(optimizer, device_placement=True)

prepare_scheduler

< source >

( scheduler: LRScheduler )

参数

scheduler (torch.optim.lr_scheduler.LRScheduler) — 一个待准备的原生 PyTorch 调度器

为任何分布式设置中的训练准备 PyTorch 调度器。建议改用 Accelerator.prepare()。

示例

>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> optimizer = torch.optim.Adam(...)
>>> scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, ...)
>>> scheduler = accelerator.prepare_scheduler(scheduler)

print

< source >

( *args **kwargs )

print() 的直接替代品，每台服务器只打印一次。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> accelerator.print("Hello world!")

profile

< source >

( profile_handler: ProfileKwargs | None = None )

参数

profile_handler (ProfileKwargs, optional) — 用于此上下文管理器的性能分析处理器。如果未传入，将使用 Accelerator 对象中设置的处理器。

将对上下文管理器内的代码进行性能分析。如果设置了 profile_handler.output_trace_dir，性能分析结果将保存到 Chrome Trace 文件中。

可以传入一个不同的 profile_handler 来覆盖 Accelerator 对象中设置的处理器。

示例

# Profile with default settings
from accelerate import Accelerator
from accelerate.utils import ProfileKwargs

accelerator = Accelerator()
with accelerator.profile() as prof:
    train()
accelerator.print(prof.key_averages().table())


# Profile with the custom handler
def custom_handler(prof):
    print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=10))


kwargs = ProfileKwargs(schedule_option=dict(wait=1, warmup=1, active=1), on_trace_ready=custom_handler)
accelerator = Accelerator(kwarg_handler=[kwargs])
with accelerator.profile() as prof:
    for _ in range(10):
        train_iteration()
        prof.step()


# Profile and export to Chrome Trace
kwargs = ProfileKwargs(output_trace_dir="output_trace")
accelerator = Accelerator(kwarg_handler=[kwargs])
with accelerator.profile():
    train()

reduce

< source >

( tensor reduction = 'sum' scale = 1.0 ) → torch.Tensor，或 torch.Tensor 的嵌套元组/列表/字典

参数

tensor (torch.Tensor 或 torch.Tensor 的嵌套元组/列表/字典) — 需要在所有进程间进行规约的张量。
reduction (str, optional, 默认为 “sum”) — 规约类型，可以是 ‘sum’、‘mean’ 或 ‘none’ 之一。如果为 ‘none’，将不执行任何操作。
scale (float, optional, 默认为 1.0) — 规约后应用的默认缩放值，仅在 XLA 上有效。

torch.Tensor，或嵌套的元组/列表/字典的 torch.Tensor

规约后的张量。

根据 reduction 在所有进程间规约 tensor 中的值。

注意：所有进程都会获得规约后的值。

示例

>>> # Assuming two processes
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> process_tensor = torch.arange(accelerator.num_processes) + 1 + (2 * accelerator.process_index)
>>> process_tensor = process_tensor.to(accelerator.device)
>>> reduced_tensor = accelerator.reduce(process_tensor, reduction="sum")
>>> reduced_tensor
tensor([4, 6])

register_for_checkpointing

< source >

( *objects )

记录 objects，并在 save_state 或 load_state 期间保存或加载它们。

当在同一脚本中加载或保存状态时应使用此功能。它不适用于在不同脚本中使用。

每个 object 必须具有 load_state_dict 和 state_dict 函数才能被存储。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume `CustomObject` has a `state_dict` and `load_state_dict` function.
>>> obj = CustomObject()
>>> accelerator.register_for_checkpointing(obj)
>>> accelerator.save_state("checkpoint.pt")

register_load_state_pre_hook

< source >

( hook: Callable[..., None] ) → torch.utils.hooks.RemovableHandle

参数

hook (Callable) — 在 Accelerator.load_state() 中调用 load_checkpoint 之前要调用的函数。

torch.utils.hooks.RemovableHandle

一个句柄，可以通过调用 `handle.remove()` 来移除添加的钩子

注册一个预处理钩子，在 Accelerator.load_state() 中调用 load_checkpoint 之前运行。

钩子应具有以下签名：

hook(models: list[torch.nn.Module], input_dir: str) -> None

models 参数是保存在加速器状态 accelerator._models 下的模型，input_dir 参数是传递给 Accelerator.load_state() 的 input_dir 参数。

应仅与 Accelerator.register_save_state_pre_hook() 结合使用。这对于加载模型权重之外的配置很有用。也可用于使用自定义方法覆盖模型加载。在这种情况下，请确保从模型列表中移除已加载的模型。

register_save_state_pre_hook

< source >

( hook: Callable[..., None] ) → torch.utils.hooks.RemovableHandle

参数

hook (Callable) — 在 Accelerator.save_state() 中调用 save_checkpoint 之前要调用的函数。

torch.utils.hooks.RemovableHandle

一个句柄，可以通过调用 `handle.remove()` 来移除添加的钩子

注册一个预处理钩子，在 Accelerator.save_state() 中调用 save_checkpoint 之前运行。

钩子应具有以下签名：

hook(models: list[torch.nn.Module], weights: list[dict[str, torch.Tensor]], input_dir: str) -> None

models 参数是保存在加速器状态 accelerator._models 下的模型，weights 参数是 models 的状态字典，而 input_dir 参数是传递给 Accelerator.load_state() 的 input_dir 参数。

应仅与 Accelerator.register_load_state_pre_hook() 结合使用。这对于保存模型权重之外的配置很有用。也可用于使用自定义方法覆盖模型保存。在这种情况下，请确保从权重列表中移除已加载的权重。

save

< source >

( obj f safe_serialization = False )

参数

obj (object) — 要保存的对象。
f (str or os.PathLike) — 保存 obj 内容的位置。
safe_serialization (bool, optional, 默认为 False) — 是否使用 safetensors 保存 obj

每台机器只将传递的对象保存到磁盘一次。用于替代 torch.save。

注意：如果在 ProjectConfiguration 中传入了 save_on_each_node，将在每个节点上保存一次对象，而不仅仅是在主节点上保存一次。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> arr = [0, 1, 2, 3]
>>> accelerator.save(arr, "array.pkl")

save_model

< source >

( model: torch.nn.Module save_directory: Union[str, os.PathLike] max_shard_size: Union[int, str] = '10GB' safe_serialization: bool = True )

参数

model — (torch.nn.Module)：要保存的模型。模型可以是包装过的或未包装的。
save_directory (str 或 os.PathLike) — 要保存到的目录。如果不存在，将会创建。
max_shard_size (int or str, optional, 默认为 "10GB") — 检查点在分片前的最大大小。检查点分片后的大小将小于此值。如果以字符串表示，需要是数字后跟单位（如 "5MB"）。

如果模型的单个权重比 max_shard_size 大，它将被放在自己的检查点分片中，该分片将大于 max_shard_size。
safe_serialization (bool, optional, 默认为 True) — 是使用 safetensors 还是传统的 PyTorch 方式（使用 pickle）来保存模型。

保存模型，以便可以使用 load_checkpoint_in_model 重新加载

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model = ...
>>> accelerator.save_model(model, save_directory)

save_state

< source >

( output_dir: str = None safe_serialization: bool = True **save_model_func_kwargs )

参数

output_dir (str 或 os.PathLike) — 用于保存所有相关权重和状态的文件夹名称。
safe_serialization (bool, optional, 默认为 True) — 是使用 safetensors 还是传统的 PyTorch 方式（使用 pickle）来保存模型。
save_model_func_kwargs (dict, optional) — 用于保存模型的额外关键字参数，可以传递给底层的保存函数，例如 DeepSpeed 的 save_checkpoint 函数的可选参数。

将模型、优化器、缩放器、RNG 生成器和已注册对象的当前状态保存到一个文件夹中。

如果向 Accelerator 对象传递了启用了 automatic_checkpoint_naming 的 ProjectConfiguration，则检查点将保存到 self.project_dir/checkpoints。如果当前保存的数量大于 total_limit，则会删除最旧的保存。每个检查点保存在名为 checkpoint_<iteration> 的独立文件夹中。

否则，它们只会被保存到 output_dir。

仅当希望在训练期间保存检查点并在相同环境中恢复状态时使用。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer, lr_scheduler = ...
>>> model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
>>> accelerator.save_state(output_dir="my_checkpoint")

set_trigger

< source >

( )

在当前进程上将内部触发张量设置为 1。后续应使用此张量进行检查，该检查将在所有进程间进行。

注意：不需要 wait_for_everyone()

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume later in the training script
>>> # `should_do_breakpoint` is a custom function to monitor when to break,
>>> # e.g. when the loss is NaN
>>> if should_do_breakpoint(loss):
...     accelerator.set_trigger()
>>> # Assume later in the training script
>>> if accelerator.check_breakpoint():
...     break

skip_first_batches

< source >

( dataloader num_batches: int = 0 )

参数

dataloader (torch.utils.data.DataLoader) — 要跳过批次的数据加载器。
num_batches (int, optional, 默认为 0) — 要跳过的批次数

创建一个新的 torch.utils.data.DataLoader，它将高效地跳过前 num_batches 个批次。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)
>>> skipped_dataloader = accelerator.skip_first_batches(dataloader, num_batches=2)
>>> # for the first epoch only
>>> for input, target in skipped_dataloader:
...     optimizer.zero_grad()
...     output = model(input)
...     loss = loss_func(output, target)
...     accelerator.backward(loss)
...     optimizer.step()

>>> # subsequent epochs
>>> for input, target in dataloader:
...     optimizer.zero_grad()
...     ...

split_between_processes

< source >

( inputs: list | tuple | dict | torch.Tensor apply_padding: bool = False )

参数

inputs (list、tuple、torch.Tensor 或 list/tuple/torch.Tensor 的 dict) — 要在进程间分割的输入。
apply_padding (bool, optional, 默认为 False) — 是否通过重复输入的最后一个元素来应用填充，以使所有进程具有相同数量的元素。这在尝试对输出执行 Accelerator.gather() 等操作或传入的输入少于进程数时非常有用。如果使用，只需记得之后丢弃填充的元素。

在 self.num_processes 之间快速分割 input，然后可以在该进程上使用。在进行分布式推理（例如使用不同提示）时非常有用。

注意，当使用 dict 时，所有键都需要有相同数量的元素。

示例

# Assume there are two processes
from accelerate import Accelerator

accelerator = Accelerator()
with accelerator.split_between_processes(["A", "B", "C"]) as inputs:
    print(inputs)
# Process 0
["A", "B"]
# Process 1
["C"]

with accelerator.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs:
    print(inputs)
# Process 0
["A", "B"]
# Process 1
["C", "C"]

trigger_sync_in_backward

< source >

( model )

参数

model (torch.nn.Module) — 需要触发梯度同步的模型。

在 `Accelerator.no_sync` 下多次前向传播后，在模型的下一次反向传播中触发梯度同步（仅适用于多 GPU 场景）。

如果脚本不是在分布式模式下启动，此上下文管理器不执行任何操作。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer)

>>> with accelerator.no_sync():
...     loss_a = loss_func(model(input_a))  # first forward pass
...     loss_b = loss_func(model(input_b))  # second forward pass
>>> accelerator.backward(loss_a)  # No synchronization across processes, only accumulate gradients
>>> with accelerator.trigger_sync_in_backward(model):
...     accelerator.backward(loss_b)  # Synchronization across all processes
>>> optimizer.step()
>>> optimizer.zero_grad()

unscale_gradients

< source >

( optimizer = None )

参数

optimizer (torch.optim.Optimizer 或 list[torch.optim.Optimizer], optional) — 需要取消梯度缩放的优化器。如果未设置，将对所有传递给 prepare() 的优化器取消梯度缩放。

在 AMP 混合精度训练中取消梯度缩放。在所有其他设置中，这是一个空操作。

可能应通过 Accelerator.clip_grad_norm_() 或 Accelerator.clip_grad_value_() 调用。

示例

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer = accelerator.prepare(model, optimizer)
>>> outputs = model(inputs)
>>> loss = loss_fn(outputs, labels)
>>> accelerator.backward(loss)
>>> accelerator.unscale_gradients(optimizer=optimizer)

unwrap_model

< source >

( model keep_fp32_wrapper: bool = True keep_torch_compile: bool = True ) → torch.nn.Module

参数

model (torch.nn.Module) — 要解包的模型。
keep_fp32_wrapper (bool, optional, 默认为 True) — 如果添加了混合精度钩子，是否不移除它。
keep_torch_compile (bool, optional, 默认为 True) — 如果模型已编译，是否不解包编译后的模型。

torch.nn.Module

解包后的模型。

从 prepare() 可能添加的额外层中解包 model。在保存模型之前很有用。

示例

>>> # Assuming two GPU processes
>>> from torch.nn.parallel import DistributedDataParallel
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model = accelerator.prepare(MyModel())
>>> print(model.__class__.__name__)
DistributedDataParallel

>>> model = accelerator.unwrap_model(model)
>>> print(model.__class__.__name__)
MyModel

verify_device_map

< source >

( model: torch.nn.Module )

验证 model 未使用类似于 auto 的设备映射进行大模型推理的准备。

wait_for_everyone

< source >

( )

将停止当前进程的执行，直到所有其他进程都到达该点（因此当脚本仅在一个进程中运行时，此操作无效）。在保存模型之前很有用。

示例

>>> # Assuming two GPU processes
>>> import time
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> if accelerator.is_main_process:
...     time.sleep(2)
>>> else:
...     print("I'm waiting for the main process to finish its sleep...")
>>> accelerator.wait_for_everyone()
>>> # Should print on every process at the same time
>>> print("Everyone is here")

实用工具

accelerate.utils.gather_object

< source >

( object: typing.Any )

参数

object (可序列化对象的嵌套列表/元组/字典) — 要收集的数据。

从所有设备递归地收集嵌套列表/元组/字典中的对象。

< > 在 GitHub 上更新