Hugging Face's logo

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

加速器

Accelerator 是用于在任何类型的训练配置中启用分布式训练的主要类。阅读 在代码中添加 Accelerator 教程，了解如何将 Accelerator 添加到您的脚本中。

加速器

类 accelerate.Accelerator

< >

( device_placement: bool = Truesplit_batches: bool = <object object at 0x7f8c18e61580>mixed_precision: PrecisionType | str | None = Nonegradient_accumulation_steps: int = 1cpu: bool = Falsedataloader_config: DataLoaderConfiguration | None = Nonedeepspeed_plugin: DeepSpeedPlugin | dict[str, DeepSpeedPlugin] | None = Nonefsdp_plugin: FullyShardedDataParallelPlugin | None = Nonetorch_tp_plugin: TorchTensorParallelPlugin | None = Nonemegatron_lm_plugin: MegatronLMPlugin | None = Nonerng_types: list[str | RNGType] | None = Nonelog_with: str | LoggerType | GeneralTracker | list[str | LoggerType | GeneralTracker] | None = Noneproject_dir: str | os.PathLike | None = Noneproject_config: ProjectConfiguration | None = Nonegradient_accumulation_plugin: GradientAccumulationPlugin | None = Nonestep_scheduler_with_optimizer: bool = Truekwargs_handlers: list[KwargsHandler] | None = Nonedynamo_backend: DynamoBackend | str | None = Nonedynamo_plugin: TorchDynamoPlugin | None = Nonedeepspeed_plugins: DeepSpeedPlugin | dict[str, DeepSpeedPlugin] | None = Noneparallelism_config: ParallelismConfig | None = None )

参数

  • device_placement (bool，可选，默认为 True) — 加速器是否应将对象放置到设备上（例如数据加载器生成的张量、模型等）。
  • 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，可选）— 已弃用：请改用 parallelism_config 和 tp_size。
  • megatron_lm_plugin (MegatronLMPlugin, 可选) — 使用此参数调整与 MegatronLM 相关的参数。该参数是可选的，也可以直接使用 accelerate config 进行配置
  • rng_types (的列表 str 或 RNGType) — 在准备好的数据加载器中，每次迭代开始时需要同步的随机数生成器列表。应为以下一项或多项：
    • "torch"：基础的 torch 随机数生成器
    • "cuda"：CUDA 随机数生成器（仅限 GPU）
    • "xla"：XLA 随机数生成器（仅限 TPU）
    • "generator"：如果底层数据集属于该类型，则为采样器（如果数据加载器中没有采样器，则为批采样器）或可迭代数据集（如果存在）的 torch.Generator。

    对于 PyTorch 版本 <=1.5.1，默认值为 ["torch"]；对于 PyTorch 版本 >= 1.6，默认值为 ["generator"]。

  • log_with (str 列表 str, LoggerType 或 GeneralTracker, 可选) — 用于实验跟踪的日志记录器列表。应为以下一项或多项：
    • "all"
    • "tensorboard"
    • "wandb"
    • "trackio"
    • "aim"
    • "comet_ml"
    • "mlflow"
    • "dvclive"
    • "swanlab" 如果选择了 "all"，将获取环境中所有可用的 tracker 并初始化它们。也可以接受 GeneralTracker 的实现以用于自定义 tracker，并且可以与 "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 无法满足需求、需要进一步调整时使用。

创建一个用于分布式训练或混合精度训练的加速器实例。

可用属性：

  • 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()

反向传播

< >

( 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

清除

< >

( *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_

< >

( parametersmax_normnorm_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_

< >

( parametersclip_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()

deepspeed_ulysses_dl_adapter

< >

( dlmodel )

这通常作为 prepare 的一部分被调用，但当数据加载器与模型分开准备时（用于外部的 accelerator.prepare 调用），需要在 prepare(model) 之后额外调用此方法（参见 HF Trainer 作为使用场景）

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, or a nested tuple/list/dictionary of torch.Tensor

参数

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

返回值

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

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

收集 tensor 在所有进程中的值，并在第一个维度上进行拼接。在评估时，可用于重新组合来自所有进程的预测结果。

注意： 此 gather 操作在所有进程中执行。

示例：

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

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

gather_for_metrics

< >

( input_datause_gather_object = False )

参数

  • input_data (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

< >

( modelunwrap = True ) → dict

参数

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

返回值

dict

模型的状态字典，可能不包含完整精度。

返回通过 Accelerator.prepare() 发送的模型的状态字典，可能不包含完整精度。

示例：

>>> 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: strunwrap: 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

< >

( joinableseven_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 = Noneload_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。

加载模型、优化器、缩放器、随机数生成器以及已注册对象的当前状态。

应仅与 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.Tensorlearning_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 = Nonebuffer_seq_dims: list[int] | None = Noneno_restore_buffers: set[torch.Tensor] | None = None )

参数

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

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

context_parallel 目前支持 FSDP2，并且要求 parallelism_config.cp_size >

  1. 如果上述任一条件未满足，该上下文管理器将不产生任何效果，但为了减少代码改动，它不会抛出异常。

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

示例：

>>> 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 = 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 = Nonelocal_process_index: int | None = 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 = 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 = Noneprocess_index: int | None = None )

参数

  • function (Callable, optional) — 要装饰的函数。
  • process_index (int, optional) — 用于运行该函数的进程的索引。

一个装饰器，用于仅在指定的进程索引上运行被装饰的函数。也可以通过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

< >

( tensordim = 0pad_index = 0pad_first = False ) → torch.Tensor, or a nested tuple/list/dictionary of 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

< >

( *argsdevice_placement = None )

参数

  • *args (对象列表 — 以下任意类型的对象：
    • torch.utils.data.DataLoader：PyTorch 数据加载器
    • torch.nn.Module：PyTorch 模块
    • torch.optim.Optimizer：PyTorch 优化器
    • torch.optim.lr_scheduler.LRScheduler：PyTorch 学习率调度器
  • 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.DataLoaderdevice_placement = Noneslice_fn_for_dispatch = None )

参数

  • data_loader (torch.utils.data.DataLoader) — 一个用于准备的普通 PyTorch DataLoader
  • device_placement (bool, 可选) — 是否将批次放置在准备好的数据加载器中的正确设备上。默认值为 self.device_placement。
  • slice_fn_for_dispatch (Callable, optional) -- If passed, this function will be used to slice tensors across num_processes. Will default to [slice_tensors()](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.slice_tensors). This argument is used only when dispatch_batchesis set toTrue` and will be ignored otherwise.

为任意分布式环境下的训练准备 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.Moduledevice_placement: bool | None = Noneevaluation_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.Optimizerdevice_placement = None )

参数

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

为任意分布式环境下的训练准备一个 PyTorch Optimizer。建议改用 Accelerator.prepare()。

示例：

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

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

prepare_scheduler

< >

( 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

< >

( *args**kwargs )

print() 的直接替代方案，用于在每台服务器上仅打印一次。

示例：

>>> from accelerate import Accelerator

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

配置文件

< >

( profile_handler: ProfileKwargs | None = None )

参数

  • profile_handler (ProfileKwargs, 可选) — 用于此上下文管理器的 profile handler。如果未传入，将使用 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

< >

( tensorreduction = 'sum'scale = 1.0 ) → torch.Tensor, or a nested tuple/list/dictionary of torch.Tensor

参数

  • tensor (torch.Tensor，或嵌套的 torch.Tensor 元组/列表/字典) — 要在所有进程间进行归约的张量。
  • reduction (str，可选，默认为 “sum”) — 归约类型，可以是 ‘sum’、‘mean’、‘max’ 或 ‘none’ 之一。如果为 ‘none’，则不执行任何操作。
  • scale (float，可选，默认为 1.0) — 在 reduce 之后应用的默认缩放值，仅在 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

< >

( *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

< >

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

参数

返回值

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() 配合使用。可用于在加载模型权重之外加载配置。也可用于通过自定义方法覆盖模型加载。在这种情况下，请确保从 models 列表中移除已加载的模型。

register_save_state_pre_hook

< >

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

参数

返回值

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() 配合使用。可用于在保存模型权重的同时保存配置。也可用于通过自定义方法覆盖模型保存。在这种情况下，请确保从权重列表中移除已加载的权重。

保存

< >

( objfsafe_serialization = False )

参数

  • obj (object) — 要保存的对象。
  • f (str 或 os.PathLike) — 用于保存 obj 内容的位置。
  • safe_serialization（bool，可选，默认为 False）——是否使用 safetensors 保存 obj

将传入的对象每台机器仅保存到磁盘一次。使用它来替代 torch.save。

注意： 如果 save_on_each_node 作为 ProjectConfiguration 传入，则会在每个节点上各保存一次该对象， 而不是仅在主节点上保存一次。

示例：

>>> from accelerate import Accelerator

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

save_model

< >

( model: torch.nn.Modulesave_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 或 str, 可选, 默认为 "10GB") — 分片前检查点的最大大小。分片后的检查点每个的大小都将小于此值。如果以字符串形式表示，则必须是数字后跟单位（如 "5MB").

    如果模型中的单个权重大于 max_shard_size，它将被放置在独立的检查点分片中，该分片的大小将超过 max_shard_size。

  • safe_serialization（bool，可选，默认为 True）— 是否使用 safetensors 保存模型，还是使用传统的 PyTorch 方式（即使用 pickle）。

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

示例：

>>> from accelerate import Accelerator

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

save_state

< >

( output_dir: str | None = Nonesafe_serialization: bool = True**save_model_func_kwargs )

参数

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

将模型、优化器、缩放器、随机数生成器以及已注册对象的当前状态保存到文件夹中。

如果向 Accelerator 对象传递了 ProjectConfiguration 并启用了 automatic_checkpoint_naming， 则检查点将保存到 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

< >

( )

在当前进程中将内部触发张量设置为 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

< >

( dataloadernum_batches: int = 0 )

参数

  • dataloader (torch.utils.data.DataLoader) — 用于跳过批次的加载器。
  • num_batches（int，可选，默认为 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

< >

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

参数

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

将 input 快速拆分到 self.num_processes 个进程之间，随后可在对应进程上使用。适用于分布式推理场景，例如处理不同的提示词。

请注意，当使用 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

< >

( 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

< >

( optimizer = None )

参数

  • optimizer (torch.optim.Optimizer 或 list[torch.optim.Optimizer]，可选) — 用于取消梯度缩放的优化器。如果未设置，则会对传递给 prepare() 的所有优化器取消梯度缩放。

在混合精度训练中使用 AMP 对梯度进行反缩放。在其他所有设置中，此操作为空操作。

可能应通过 Accelerator.clipgrad_norm() 或 Accelerator.clipgrad_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

< >

( modelkeep_fp32_wrapper: bool = Truekeep_torch_compile: bool = True ) → torch.nn.Module

参数

  • model (torch.nn.Module) — 要解包的模型。
  • keep_fp32_wrapper (bool，可选，默认为 True) — 是否不移除已添加的混合精度钩子。
  • keep_torch_compile (bool，可选，默认为 True) — 如果模型已编译，是否不取消编译模型的包装。

返回值

torch.nn.Module

未包装的模型。

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

示例：

>>> # 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

< >

( model: torch.nn.Module )

验证 model 是否未使用类似 auto 的 device-map 进行大模型推理准备。

wait_for_everyone

< >

( )

将在当前进程的执行暂停，直到所有其他进程都到达该点（因此，当脚本仅在单个进程中运行时，此操作不会执行任何操作）。在保存模型之前使用此操作非常有用。

示例：

>>> # 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

< >

( object: Any )

参数

  • object（可 pickle 对象的嵌套列表/元组/字典）— 要收集的数据。

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

Update on GitHub

←在 TPU 上进行训练 有状态类→