Accelerate 文档

实用函数和类

Accelerate

加入 Hugging Face 社区

并获得增强的文档体验

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

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

切换文档主题

开始使用

实用函数和类

下面是 🤗 Accelerate 提供的各种实用函数，按用例分类。

常量

在 🤗 Accelerate 中使用的常量，供参考

以下是使用 Accelerator.save_state() 时使用的常量

utils.MODEL_NAME: "pytorch_model" utils.OPTIMIZER_NAME: "optimizer" utils.RNG_STATE_NAME: "random_states" utils.SCALER_NAME: "scaler.pt" utils.SCHEDULER_NAME: "scheduler"

以下是使用 Accelerator.save_model() 时使用的常量

utils.WEIGHTS_NAME: "pytorch_model.bin" utils.SAFE_WEIGHTS_NAME: "model.safetensors" utils.WEIGHTS_INDEX_NAME: "pytorch_model.bin.index.json" utils.SAFE_WEIGHTS_INDEX_NAME: "model.safetensors.index.json"

数据类

这些是在 🤗 Accelerate 中使用的基本数据类，它们可以作为参数传递。

独立

这些是用于检查的独立数据类，例如正在使用的分布式系统的类型

class accelerate.utils.ComputeEnvironment

( value names = None module = None qualname = None type = None start = 1 )

表示计算环境的类型。

数值

LOCAL_MACHINE — 私有/自定义集群硬件。
AMAZON_SAGEMAKER — 作为计算环境的 Amazon SageMaker。

class accelerate.DistributedType

( value names = None module = None qualname = None type = None start = 1 )

表示分布式环境的类型。

数值

NO — 不是分布式环境，只是单个进程。
MULTI_CPU — 在多个 CPU 节点上分布式运行。
MULTI_GPU — 在多个 GPU 上分布式运行。
MULTI_MLU — 在多个 MLU 上分布式运行。
MULTI_SDAA — 在多个 SDAA 上分布式运行。
MULTI_MUSA — 在多个 MUSA 上分布式运行。
MULTI_NPU — 在多个 NPU 上分布式运行。
MULTI_XPU — 在多个 XPU 上分布式运行。
MULTI_HPU — 在多个 HPU 上分布式运行。
DEEPSPEED — 使用 DeepSpeed。
XLA — 使用 TorchXLA。

class accelerate.utils.DynamoBackend

( value names = None module = None qualname = None type = None start = 1 )

表示 dynamo 后端（参见 https://pytorch.ac.cn/docs/stable/torch.compiler.html）。

数值

NO — 不使用 torch dynamo。
EAGER — 使用 PyTorch 运行提取的 GraphModule。这在调试 TorchDynamo 问题时非常有用。
AOT_EAGER — 使用 AotAutograd，但没有编译器，即只使用 PyTorch eager 来处理 AotAutograd 提取的前向和后向图。这对于调试很有用，但不太可能提供加速。
INDUCTOR — 使用 TorchInductor 后端，结合 AotAutograd 和 cudagraphs，利用代码生成的 Triton 内核。了解更多
AOT_TS_NVFUSER — 结合 AotAutograd/TorchScript 的 nvFuser。了解更多
NVPRIMS_NVFUSER — 结合 PrimTorch 的 nvFuser。了解更多
CUDAGRAPHS — 结合 AotAutograd 的 cudagraphs。了解更多
OFI — 使用 Torchscript optimize_for_inference。仅用于推理。了解更多
FX2TRT — 使用 Nvidia TensorRT 进行推理优化。仅用于推理。了解更多
ONNXRT — 使用 ONNXRT 在 CPU/GPU 上进行推理。仅用于推理。了解更多
TENSORRT — 使用 ONNXRT 运行 TensorRT 进行推理优化。了解更多
AOT_TORCHXLA_TRACE_ONCE — 使用 Pytorch/XLA 和 TorchDynamo 优化进行训练。了解更多
TORCHXLA_TRACE_ONCE — 使用 Pytorch/XLA 和 TorchDynamo 优化进行推理。了解更多
IPEX — 使用 IPEX 在 CPU 上进行推理。仅用于推理。了解更多。
TVM — 使用 Apache TVM 进行推理优化。了解更多
HPU_BACKEND — 使用 HPU 后端进行推理优化。

class accelerate.utils.LoggerType

( value names = None module = None qualname = None type = None start = 1 )

表示支持的实验追踪器的类型

数值

ALL — 环境中所有支持的可用追踪器
TENSORBOARD — TensorBoard 作为实验追踪器
WANDB — wandb 作为实验追踪器
TRACKIO — trackio 作为实验追踪器
COMETML — comet_ml 作为实验追踪器
MLFLOW — mlflow 作为实验追踪器
CLEARML — clearml 作为实验追踪器
DVCLIVE — dvclive 作为实验追踪器
SWANLAB — swanlab 作为实验追踪器

class accelerate.utils.PrecisionType

( value names = None module = None qualname = None type = None start = 1 )

表示用于浮点值的精度类型

数值

NO — 使用全精度 (FP32)
FP16 — 使用半精度
BF16 — 使用脑浮点精度

class accelerate.utils.RNGType

( value names = None module = None qualname = None type = None start = 1 )

一个枚举。

class accelerate.utils.SageMakerDistributedType

( value names = None module = None qualname = None type = None start = 1 )

表示分布式环境的类型。

数值

NO — 不是分布式环境，只是单个进程。
DATA_PARALLEL — 使用 SageMaker 分布式数据并行。
MODEL_PARALLEL — 使用 SageMaker 分布式模型并行。

关键字参数

这些是用于特定交互的可配置参数，这些交互在 PyTorch 生态系统中由 Accelerate 在幕后处理。

class accelerate.AutocastKwargs

( enabled: bool = True cache_enabled: bool = None )

在您的 Accelerator 中使用此对象来自定义 torch.autocast 的行为。有关每个参数的更多信息，请参阅此上下文管理器的文档。

示例

from accelerate import Accelerator
from accelerate.utils import AutocastKwargs

kwargs = AutocastKwargs(cache_enabled=True)
accelerator = Accelerator(kwargs_handlers=[kwargs])

class accelerate.DistributedDataParallelKwargs

( dim: int = 0 broadcast_buffers: bool = True bucket_cap_mb: int = 25 find_unused_parameters: bool = False check_reduction: bool = False gradient_as_bucket_view: bool = False static_graph: bool = False comm_hook: DDPCommunicationHookType = <DDPCommunicationHookType.NO: 'no'> comm_wrapper: typing.Literal[<DDPCommunicationHookType.NO: 'no'>, <DDPCommunicationHookType.FP16: 'fp16'>, <DDPCommunicationHookType.BF16: 'bf16'>] = <DDPCommunicationHookType.NO: 'no'> comm_state_option: dict = <factory> )

在您的 Accelerator 中使用此对象，来自定义您的模型如何被包装在 torch.nn.parallel.DistributedDataParallel 中。有关每个参数的更多信息，请参阅此包装器的文档。

gradient_as_bucket_view 仅在 PyTorch 1.7.0 及更高版本中可用。

static_graph 仅在 PyTorch 1.11.0 及更高版本中可用。

示例

from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs

kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
accelerator = Accelerator(kwargs_handlers=[kwargs])

class accelerate.utils.FP8RecipeKwargs

( opt_level: typing.Literal['O1', 'O2'] = None use_autocast_during_eval: bool = None margin: int = None interval: int = None fp8_format: typing.Literal['HYBRID', 'E4M3', 'E5M2'] = None amax_history_len: int = None amax_compute_algo: typing.Literal['max', 'most_recent'] = None override_linear_precision: tuple = None backend: typing.Literal['MSAMP', 'TE'] = None )

已弃用。请改用适当的 FP8 配方 kwargs 类，例如 `TERecipeKwargs` 或 `MSAMPRecipeKwargs`。

class accelerate.GradScalerKwargs

( init_scale: float = 65536.0 growth_factor: float = 2.0 backoff_factor: float = 0.5 growth_interval: int = 2000 enabled: bool = True )

在您的 Accelerator 中使用此对象来自定义混合精度的行为，特别是如何创建 torch.amp.GradScaler 或 torch.cuda.amp.GradScaler。有关每个参数的更多信息，请参阅此scaler的文档。

torch.cuda.amp.GradScaler 仅在 PyTorch 1.5.0 及更高版本中可用，而 torch.amp.GradScaler 仅在 PyTorch 2.4.0 及更高版本中可用。

示例

from accelerate import Accelerator
from accelerate.utils import GradScalerKwargs

kwargs = GradScalerKwargs(backoff_factor=0.25)
accelerator = Accelerator(kwargs_handlers=[kwargs])

class accelerate.utils.InitProcessGroupKwargs

( backend: typing.Optional[str] = 'nccl' init_method: typing.Optional[str] = None timeout: typing.Optional[datetime.timedelta] = None )

在您的 Accelerator 中使用此对象来自定义分布式进程的初始化。有关每个参数的更多信息，请参阅此方法的文档。

注意：如果 timeout 设置为 None，默认值将根据 backend 的设置而定。

from datetime import timedelta
from accelerate import Accelerator
from accelerate.utils import InitProcessGroupKwargs

kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=800))
accelerator = Accelerator(kwargs_handlers=[kwargs])

class accelerate.utils.KwargsHandler

( )

内部混入类，为数据类实现一个 to_kwargs() 方法。

to_kwargs

( )

返回一个字典，其中包含与此类默认值不同的属性及其值。

插件

这些是可以传递给 Accelerator 对象的插件。虽然它们在文档的其他地方有定义，但为了方便起见，这里列出了所有插件。

class accelerate.DeepSpeedPlugin

( hf_ds_config: typing.Any = None gradient_accumulation_steps: int = None gradient_clipping: float = None zero_stage: int = None is_train_batch_min: bool = True offload_optimizer_device: str = None offload_param_device: str = None offload_optimizer_nvme_path: str = None offload_param_nvme_path: str = None zero3_init_flag: bool = None zero3_save_16bit_model: bool = None transformer_moe_cls_names: str = None enable_msamp: bool = None msamp_opt_level: typing.Optional[typing.Literal['O1', 'O2']] = None )

参数

hf_ds_config (Any, 默认为 None) — DeepSpeed 配置文件路径、字典或 accelerate.utils.deepspeed.HfDeepSpeedConfig 类的对象。
gradient_accumulation_steps (int, 默认为 None) — 在更新优化器状态之前累积梯度的步数。如果未设置，将直接使用 Accelerator 中的值。
gradient_clipping (float, 默认为 None) — 启用梯度裁剪并设置其值。
zero_stage (int, 默认为 None) — 可选项为 0, 1, 2, 3。默认值将从环境变量中获取。
is_train_batch_min (bool, 默认为 True) — 如果同时指定了训练和评估数据加载器，此参数将决定 train_batch_size。
offload_optimizer_device (str, 默认为 None) — 可选项为 none|cpu|nvme。仅适用于 ZeRO Stages 2 和 3。
offload_param_device (str, 默认为 None) — 可选项为 none|cpu|nvme。仅适用于 ZeRO Stage 3。
offload_optimizer_nvme_path (str, 默认为 None) — 可选项为 /nvme|/local_nvme。仅适用于 ZeRO Stage 3。
offload_param_nvme_path (str, 默认为 None) — 可选项为 /nvme|/local_nvme。仅适用于 ZeRO Stage 3。
zero3_init_flag (bool, 默认为 None) — 指示是否保存 16 位模型的标志。仅适用于 ZeRO Stage-3。
zero3_save_16bit_model (bool, 默认为 None) — 指示是否保存 16 位模型的标志。仅适用于 ZeRO Stage-3。
transformer_moe_cls_names (str, 默认为 None) — 以逗号分隔的 Transformers MoE 层类名列表（区分大小写）。例如，MixtralSparseMoeBlock、Qwen2MoeSparseMoeBlock、JetMoEAttention、JetMoEBlock 等。
enable_msamp (bool, 默认为 None) — 指示是否为 FP8 训练启用 MS-AMP 后端的标志。
msasmp_opt_level (Optional[Literal["O1", "O2"]], 默认为 None) — MS-AMP 的优化级别（默认为 'O1'）。仅当 enable_msamp 为 True 时适用。应为 ['O1' 或 'O2'] 之一。

此插件用于集成 DeepSpeed。

deepspeed_config_process

( prefix = '' mismatches = None config = None must_match = True **kwargs )

使用 kwargs 中的值处理 DeepSpeed 配置。

选择

( _from_accelerator_state: bool = False )

设置 HfDeepSpeedWeakref 以使用当前的 DeepSpeed 插件配置

class accelerate.FullyShardedDataParallelPlugin

( fsdp_version: int = None sharding_strategy: typing.Union[str, ForwardRef('torch.distributed.fsdp.ShardingStrategy')] = None reshard_after_forward: typing.Union[str, ForwardRef('torch.distributed.fsdp.ShardingStrategy'), bool] = None backward_prefetch: typing.Union[str, ForwardRef('torch.distributed.fsdp.BackwardPrefetch'), NoneType] = None mixed_precision_policy: typing.Union[dict, ForwardRef('torch.distributed.fsdp.MixedPrecision'), ForwardRef('torch.distributed.fsdp.MixedPrecisionPolicy'), NoneType] = None auto_wrap_policy: typing.Union[typing.Callable, typing.Literal['transformer_based_wrap', 'size_based_wrap', 'no_wrap'], NoneType] = None cpu_offload: typing.Union[bool, ForwardRef('torch.distributed.fsdp.CPUOffload'), ForwardRef('torch.distributed.fsdp.CPUOffloadPolicy')] = None ignored_modules: typing.Union[collections.abc.Iterable[torch.nn.modules.module.Module], str, NoneType] = None state_dict_type: typing.Union[str, ForwardRef('torch.distributed.fsdp.StateDictType')] = None state_dict_config: typing.Union[ForwardRef('torch.distributed.fsdp.FullStateDictConfig'), ForwardRef('torch.distributed.fsdp.ShardedStateDictConfig'), NoneType] = None optim_state_dict_config: typing.Union[ForwardRef('torch.distributed.fsdp.FullOptimStateDictConfig'), ForwardRef('torch.distributed.fsdp.ShardedOptimStateDictConfig'), NoneType] = None limit_all_gathers: bool = True use_orig_params: typing.Optional[bool] = None param_init_fn: typing.Optional[typing.Callable[[torch.nn.modules.module.Module], NoneType]] = None sync_module_states: typing.Optional[bool] = None forward_prefetch: bool = None activation_checkpointing: bool = None cpu_ram_efficient_loading: bool = None transformer_cls_names_to_wrap: typing.Optional[list[str]] = None min_num_params: typing.Optional[int] = None )

参数

fsdp_version (int, 默认为 1) — 要使用的 FSDP 版本。默认为 1。如果设置为 2，启动器期望配置已转换为 FSDP2 格式。
sharding_strategy (Union[str, torch.distributed.fsdp.ShardingStrategy], 默认为 'FULL_SHARD') — 要使用的分片策略。应为 str 或 torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy 的实例。已弃用，建议使用 reshard_after_forward。
reshard_after_forward (Union[str, torch.distributed.fsdp.ShardingStrategy, bool], 默认为 fsdp_version=1 时为 'FULL_SHARD'，fsdp_version=2 时为 True) — 要使用的分片策略。如果 fsdp_version 设置为 2，则应为布尔值，否则应为 str 或 torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy 的实例。
backward_prefetch (Union[str, torch.distributed.fsdp.BackwardPrefetch], 默认为 'NO_PREFETCH') — 要使用的反向预取策略。应为 str 或 torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch 的实例。
mixed_precision_policy (Optional[Union[dict, torch.distributed.fsdp.MixedPrecision, torch.distributed.fsdp.MixedPrecisionPolicy]], 默认为 None) — 用于启用 FullyShardedDataParallel 混合精度训练的配置。如果传入 dict，它应包含以下键：param_dtype、reduce_dtype 和 buffer_dtype。如果 fsdp_version 设置为 2，则可以是 torch.distributed.fsdp.MixedPrecisionPolicy 的实例。
auto_wrap_policy (Optional(Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]), 默认为 NO_WRAP) — 指定递归地用 FSDP 包装层的策略，可以是可调用对象或字符串。如果为字符串，则必须是 transformer_based_wrap、size_based_wrap 或 no_wrap 之一。有关其具体形式，请参阅 torch.distributed.fsdp.wrap.size_based_wrap_policy。
cpu_offload (Union[bool, torch.distributed.fsdp.CPUOffload, torch.distributed.fsdp.CPUOffloadPolicy], 默认为 False) — 是否将参数卸载到 CPU。如果 fsdp_version 设置为 2，应为布尔值或 torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload 或 torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy 的实例。
ignored_modules (Optional[Union[Iterable[torch.nn.Module], str]], 默认为 None) — 在使用 FSDP 包装时要忽略的模块列表。如果传入字符串，将使用正则表达式的 fullmatch 按名称匹配模块。
state_dict_type (Union[str, torch.distributed.fsdp.StateDictType], 默认为 'FULL_STATE_DICT') — 要使用的状态字典类型。如果为字符串，则必须是 full_state_dict、local_state_dict 或 sharded_state_dict 之一。
state_dict_config (Optional[Union[torch.distributed.fsdp.FullStateDictConfig, torch.distributed.fsdp.ShardedStateDictConfig], 默认为 None) — 要使用的状态字典配置。如果未传入，则根据 state_dict_type 确定。
optim_state_dict_config (Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig], 默认为 None) — 要使用的优化器状态字典配置。如果未传入，则根据 state_dict_type 确定。
limit_all_gathers (bool, 默认为 True) — 是否让 FSDP 显式同步 CPU 线程，以防止过多的正在进行的 all-gather 操作。此布尔值仅影响调度 all-gather 的分片策略。启用此选项有助于减少 CUDA malloc 重试的次数。
use_orig_params (bool, 默认为 False) — 是否为优化器使用原始参数。
param_init_fn (Optional[Callable[[torch.nn.Module], None], 默认为 None) — 一个 Callable[torch.nn.Module] -> None，用于指定当前在元设备上的模块应如何初始化到实际设备上。仅当 sync_module_states 为 True 时适用。默认是一个在模块上调用 to_empty 的 lambda 函数。
sync_module_states (bool, 默认为 False) — 每个单独包装的 FSDP 单元是否应从 rank 0 广播模块参数，以确保它们在初始化后在所有 rank 上都相同。默认为 False，除非 cpu_ram_efficient_loading 为 True，此时将被强制启用。
forward_prefetch (bool, 默认为 False) — 是否让 FSDP 在前向传播执行时显式预取下一个即将到来的 all-gather 操作。仅与静态图一起使用。
activation_checkpointing (bool, 默认为 False) — 一种通过清除某些层的激活并在反向传播期间重新计算它们来减少内存使用的技术。实际上，这是用额外的计算时间换取减少的内存使用。
cpu_ram_efficient_loading (bool, 默认为 None) — 如果为 True，只有第一个进程加载预训练模型检查点，而所有其他进程的权重为空。仅适用于 Transformers。使用此选项时，sync_module_states 必须为 True。
transformer_cls_names_to_wrap (Optional[List[str]], 默认为 None) — 要包装的 Transformer 层类名列表。仅当 auto_wrap_policy 为 transformer_based_wrap 时适用。
min_num_params (Optional[int], 默认为 None) — 一个模块必须具有的最小参数数量才能被包装。仅当 auto_wrap_policy 为 size_based_wrap 时适用。

该插件用于启用完全分片数据并行。

set_auto_wrap_policy

( model )

给定 model，根据传入的策略以及是否可以使用 transformer_cls_to_wrap 创建一个 auto_wrap_policy。

set_mixed_precision

( mixed_precision buffer_autocast = False override = False )

为 FSDP 设置混合精度策略

set_state_dict_type

( state_dict_type = None )

根据 `StateDictType` 设置 state dict 配置。

validate_mixed_precision_policy

( )

验证混合精度策略，抽象出来以避免在不需要时引入导入。

class accelerate.utils.GradientAccumulationPlugin

( num_steps: int = None adjust_scheduler: bool = True sync_with_dataloader: bool = True sync_each_batch: bool = False )

参数

num_steps (int) — 累积梯度的步数。
adjust_scheduler (bool, 可选, 默认为 True) — 是否调整调度器的步数以考虑累积的步数。如果使用的调度器未针对梯度累积进行调整，则应为 True。
sync_with_dataloader (bool, 可选, 默认为 True) — 是否在 dataloader 结束时同步设置梯度。
sync_each_batch (bool, 可选) — 是否在每个数据批次同步设置梯度。设置为 True 可能会在分布式训练中使用梯度累积时减少内存需求，但会牺牲速度。

一个用于配置梯度累积行为的插件。您只能将 gradient_accumulation_plugin 或 gradient_accumulation_steps 中的一个传递给 Accelerator。同时传递两者会引发错误。

示例

from accelerate.utils import GradientAccumulationPlugin

gradient_accumulation_plugin = GradientAccumulationPlugin(num_steps=2)
accelerator = Accelerator(gradient_accumulation_plugin=gradient_accumulation_plugin)

class accelerate.utils.MegatronLMPlugin

( tp_degree: int = None pp_degree: int = None num_micro_batches: int = None gradient_clipping: float = None sequence_parallelism: bool = None recompute_activations: bool = None use_distributed_optimizer: bool = None pipeline_model_parallel_split_rank: int = None num_layers_per_virtual_pipeline_stage: int = None is_train_batch_min: str = True train_iters: int = None train_samples: int = None weight_decay_incr_style: str = 'constant' start_weight_decay: float = None end_weight_decay: float = None lr_decay_style: str = 'linear' lr_decay_iters: int = None lr_decay_samples: int = None lr_warmup_iters: int = None lr_warmup_samples: int = None lr_warmup_fraction: float = None min_lr: float = 0 consumed_samples: list = None no_wd_decay_cond: typing.Optional[typing.Callable] = None scale_lr_cond: typing.Optional[typing.Callable] = None lr_mult: float = 1.0 megatron_dataset_flag: bool = False seq_length: int = None encoder_seq_length: int = None decoder_seq_length: int = None tensorboard_dir: str = None set_all_logging_options: bool = False eval_iters: int = 100 eval_interval: int = 1000 return_logits: bool = False custom_train_step_class: typing.Optional[typing.Any] = None custom_train_step_kwargs: typing.Optional[dict[str, typing.Any]] = None custom_model_provider_function: typing.Optional[typing.Callable] = None custom_prepare_model_function: typing.Optional[typing.Callable] = None custom_megatron_datasets_provider_function: typing.Optional[typing.Callable] = None custom_get_batch_function: typing.Optional[typing.Callable] = None custom_loss_function: typing.Optional[typing.Callable] = None other_megatron_args: typing.Optional[dict[str, typing.Any]] = None )

参数

tp_degree (int, 默认为 None) — 张量并行度。
pp_degree (int, 默认为 None) — 流水线并行度。
num_micro_batches (int, 默认为 None) — 微批次数量。
gradient_clipping (float, 默认为 None) — 基于全局 L2 范数的梯度裁剪值（0 表示禁用）。
sequence_parallelism (bool, 默认为 None) — 启用序列并行。
recompute_activations (bool, 默认为 None) — 启用选择性激活重计算。
use_distributed_optimizr (bool, 默认为 None) — 启用分布式优化器。
pipeline_model_parallel_split_rank (int, 默认为 None) — 编码器和解码器应被拆分的 rank。
num_layers_per_virtual_pipeline_stage (int, 默认为 None) — 每个虚拟流水线阶段的层数。
is_train_batch_min (str, 默认为 True) — 如果同时指定了训练和评估 dataloader，这将决定 `micro_batch_size`。
train_iters (int, 默认为 None) — 所有训练运行中要训练的总样本数。注意，在使用 `MegatronLMDummyScheduler` 时，应提供 train-iters 或 train-samples。
train_samples (int, 默认为 None) — 所有训练运行中要训练的总样本数。注意，在使用 `MegatronLMDummyScheduler` 时，应提供 train-iters 或 train-samples。
weight_decay_incr_style (str, 默认为 'constant') — 权重衰减增量函数。可选值=[“constant”, “linear”, “cosine”]。
start_weight_decay (float, 默认为 None) — L2 正则化的初始权重衰减系数。
end_weight_decay (float, 默认为 None) — 运行结束时 L2 正则化的权重衰减系数。
lr_decay_style (str, 默认为 'linear') — 学习率衰减函数。可选值=[‘constant’, ‘linear’, ‘cosine’]。
lr_decay_iters (int, 默认为 None) — 学习率衰减的迭代次数。如果为 None，则默认为 `train_iters`。
lr_decay_samples (int, 默认为 None) — 学习率衰减的样本数。如果为 None，则默认为 `train_samples`。
lr_warmup_iters (int, 默认为 None) — 线性预热学习率的迭代次数。
lr_warmup_samples (int，默认为 None) — 用于线性预热学习率的样本数量。
lr_warmup_fraction (float，默认为 None) — 用于线性预热学习率的 lr-warmup-(迭代次数/样本数) 的比例。
min_lr (float，默认为 0) — 学习率的最小值。调度器会将低于此阈值的值进行裁剪。
consumed_samples (List，默认为 None) — 已消耗的样本数量，顺序与传递给 accelerator.prepare 调用的 dataloaders 相同。
no_wd_decay_cond (Optional，默认为 None) — 禁用权重衰减的条件。
scale_lr_cond (Optional，默认为 None) — 缩放学习率的条件。
lr_mult (float，默认为 1.0) — 学习率乘数。
megatron_dataset_flag (bool，默认为 False) — 数据集格式是否遵循 Megatron-LM 的 Indexed/Cached/MemoryMapped 格式。
seq_length (int，默认为 None) — 要处理的最大序列长度。
encoder_seq_length (int，默认为 None) — 编码器要处理的最大序列长度。
decoder_seq_length (int，默认为 None) — 解码器要处理的最大序列长度。
tensorboard_dir (str，默认为 None) — 保存 tensorboard 日志的路径。
set_all_logging_options (bool，默认为 False) — 是否设置所有日志记录选项。
eval_iters (int，默认为 100) — 评估验证/测试时运行的迭代次数。
eval_interval (int，默认为 1000) — 在验证集上运行评估的间隔。
return_logits (bool，默认为 False) — 是否从模型返回 logits。
custom_train_step_class (Optional，默认为 None) — 自定义训练步骤类。
custom_train_step_kwargs (Optional，默认为 None) — 自定义训练步骤的关键字参数。
custom_model_provider_function (Optional，默认为 None) — 自定义模型提供函数。
custom_prepare_model_function (Optional，默认为 None) — 自定义模型准备函数。
custom_megatron_datasets_provider_function (Optional，默认为 None) — 自定义 megatron 训练/验证/测试数据集提供函数。
custom_get_batch_function (Optional，默认为 None) — 自定义获取批次函数。
custom_loss_function (Optional，默认为 None) — 自定义损失函数。
other_megatron_args (Optional，默认为 None) — 其他 Megatron-LM 参数。请参考 Megatron-LM。

用于 Megatron-LM 的插件，以启用张量、流水线、序列和数据并行。还可启用选择性激活重计算和优化的融合内核。

class accelerate.utils.TorchDynamoPlugin

( backend: DynamoBackend = None mode: str = None fullgraph: bool = None dynamic: bool = None options: typing.Any = None disable: bool = False use_regional_compilation: bool = None )

参数

backend (DynamoBackend，默认为 None) — 一个有效的 Dynamo 后端。更多详情请参阅 https://pytorch.ac.cn/docs/stable/torch.compiler.html。
mode (str，默认为 None) — 可选值为 ‘default’、‘reduce-overhead’ 或 ‘max-autotune’。
fullgraph (bool，默认为 None) — 是否可以将模型分解为多个子图。
dynamic (bool，默认为 None) — 是否使用动态形状进行追踪。
options (Any，默认为 None) — 传递给后端的选项字典。
disable (bool，默认为 False) — 将 torch.compile() 变为空操作，用于测试。
use_regional_compilation (bool，默认为 None) — 使用此选项可减少 torch.compile() 的冷启动编译时间，通过针对相同类的重复块并顺序编译它们来命中编译器缓存。例如，在 GPT2LMHeadModel 中，重复的块/类是 GPT2Block，可以通过 model.transformer.h[0] 访问。模型的其余部分（例如 model.lm_head）将单独编译。

此插件用于使用 PyTorch 2.0 编译模型

配置

这些是可以配置并传递给相应集成的类

class accelerate.utils.BnbQuantizationConfig

( load_in_8bit: bool = False llm_int8_threshold: float = 6.0 load_in_4bit: bool = False bnb_4bit_quant_type: str = 'fp4' bnb_4bit_use_double_quant: bool = False bnb_4bit_compute_dtype: str = 'fp16' torch_dtype: dtype = None skip_modules: list = None keep_in_fp32_modules: list = None )

参数

load_in_8bit (bool，默认为 False) — 启用 8 位量化。
llm_int8_threshold (float，默认为 6.0) — 离群值阈值。仅在 load_in_8bit=True 时相关。
load_in_4_bit (bool，默认为 False) — 启用 4 位量化。
bnb_4bit_quant_type (str，默认为 fp4) — 设置 bnb.nn.Linear4Bit 层中的量化数据类型。选项为 {‘fp4’,‘np4’}。
bnb_4bit_use_double_quant (bool，默认为 False) — 启用嵌套量化，即再次量化第一次量化产生的量化常数。
bnb_4bit_compute_dtype (bool，默认为 fp16) — 这设置了计算类型，可能与输入类型不同。例如，输入可能是 fp32，但计算可以设置为 bf16 以加速。选项为 {‘fp32’,‘fp16’,‘bf16’}。
torch_dtype (torch.dtype，默认为 None) — 这设置了其余未量化层的 dtype。bitsandbytes 库建议对 8 位模型将值设置为 torch.float16，对 4 位模型使用与计算 dtype 相同的 dtype。
skip_modules (List[str]，默认为 None) — 我们不进行量化的模块的显式列表。这些模块的 dtype 将是 torch_dtype。
keep_in_fp32_modules (List，默认为 None) — 我们不进行量化的模块的显式列表。我们将它们保持为 torch.float32。

一个用于启用 BitsAndBytes 4 位和 8 位量化的插件

class accelerate.DataLoaderConfiguration

( split_batches: bool = False dispatch_batches: bool = None even_batches: bool = True use_seedable_sampler: bool = False data_seed: int = None non_blocking: bool = False use_stateful_dataloader: bool = False )

参数

split_batches (bool，默认为 False) — 加速器是否应将 dataloader 产生的批次在设备间进行拆分。如果为 True，实际使用的批次大小在任何类型的分布式进程上都将相同，但它必须是您使用的 num_processes 的整数倍。如果为 False，实际使用的批次大小将是您脚本中设置的大小乘以进程数。
dispatch_batches (bool，默认为 None) — 如果设置为 True，由 Accelerator 准备的 dataloader 仅在主进程上迭代，然后将批次拆分并广播到每个进程。对于底层数据集为 IterableDataset 的 DataLoader，默认为 True，否则为 False。
even_batches (bool，默认为 True) — 如果设置为 True，在所有进程的总批次大小不能完全整除数据集的情况下，将复制数据集开头的样本，以便批次可以平均分配给所有工作进程。
use_seedable_sampler (bool，默认为 False) — 是否使用完全可设定种子的随机采样器 (data_loader.SeedableRandomSampler)。确保使用不同的采样技术可以完全复现训练结果。虽然不同种子之间的结果可能不同，但使用多个不同种子进行比较时，平均差异可以忽略不计。为获得最佳效果，还应与 set_seed() 一起运行。
data_seed (int，默认为 None) — 当使用 use_seedable_sampler 时，用于底层生成器的种子。如果为 None，生成器将使用 torch 的当前默认种子。
non_blocking (bool，默认为 False) — 如果设置为 True，由 Accelerator 准备的 dataloader 将利用非阻塞的主机到设备传输，从而更好地实现 dataloader 通信和计算之间的重叠。建议准备好的 dataloader 将 pin_memory 设置为 True 以正常工作。
use_stateful_dataloader (bool，默认为 False) — 如果设置为 True，由 Accelerator 准备的 dataloader 将由 torchdata.StatefulDataLoader 支持。这需要安装支持 StatefulDataLoader 的 torchdata 0.8.0 或更高版本。

调用 accelerator.prepare 时与 dataloader 相关的项目的配置。

class accelerate.utils.ProjectConfiguration

( project_dir: str = None logging_dir: str = None automatic_checkpoint_naming: bool = False total_limit: int = None iteration: int = 0 save_on_each_node: bool = False )

参数

project_dir (str，默认为 None) — 用于存储数据的目录路径。
logging_dir (str, 默认为 None) — 用于存储本地兼容日志记录器日志的目录路径。如果为 None，则默认为 project_dir。
automatic_checkpoint_naming (bool, 默认为 False) — 保存的状态是否应自动进行迭代命名。
total_limit (int, 默认为 None) — 要保留的总保存状态的最大数量。
iteration (int, 默认为 0) — 当前的保存迭代次数。
save_on_each_node (bool, 默认为 False) — 在进行多节点分布式训练时，是在每个节点上保存模型和检查点，还是仅在主节点上保存。

基于项目内部需求的 Accelerator 对象配置。

set_directories

( project_dir: str = None )

将 self.project_dir 和 self.logging_dir 设置为适当的值。

环境变量

这些是可以为不同用例启用的环境变量

ACCELERATE_DEBUG_MODE (str): 是否以调试模式运行 accelerate。更多信息请参见此处。

数据操作与运算

这些数据操作与 torch 中的操作类似，但可以在分布式进程上使用。

accelerate.utils.broadcast

( tensor from_process: int = 0 )

参数

tensor (嵌套的 list/tuple/dictionary of torch.Tensor) — 要收集的数据。
from_process (int, 可选, 默认为 0) — 发送数据的进程。

递归地将嵌套的 list/tuple/dictionary 中的张量广播到所有设备。

accelerate.utils.broadcast_object_list

( object_list from_process: int = 0 )

参数

object_list (list of picklable objects) — 要广播的对象列表。该列表将被原地修改。
from_process (int, 可选, 默认为 0) — 发送数据的进程。

将一个可序列化对象列表从一个进程广播到其他进程。

accelerate.utils.concatenate

( data dim = 0 )

参数

data (嵌套的 list/tuple/dictionary of lists of tensors torch.Tensor) — 要连接的数据。
dim (int, 可选, 默认为 0) — 连接的维度。

递归地连接嵌套的 list/tuple/dictionary 中的张量列表，这些张量具有相同的形状。

accelerate.utils.convert_outputs_to_fp32

( model_forward )

accelerate.utils.convert_to_fp32

( tensor )

参数

tensor (嵌套的 list/tuple/dictionary of torch.Tensor) — 要从 FP16/BF16 转换为 FP32 的数据。

递归地将嵌套的 list/tuple/dictionary 中 FP16/BF16 精度的张量元素转换为 FP32。

accelerate.utils.gather

( tensor )

参数

tensor (嵌套的 list/tuple/dictionary of torch.Tensor) — 要收集的数据。

递归地从所有设备收集嵌套的 list/tuple/dictionary 中的张量。

accelerate.utils.gather_object

( object: typing.Any )

参数

object (嵌套的 list/tuple/dictionary of picklable object) — 要收集的数据。

递归地从所有设备收集嵌套的 list/tuple/dictionary 中的对象。

accelerate.utils.get_grad_scaler

( distributed_type: DistributedType = None **kwargs )

参数

distributed_type (DistributedType, 可选, 默认为 None) — 分布式环境的类型。
kwargs — 用于 GradScaler 构造函数的附加参数。

一个通用的辅助函数，它将根据环境初始化正确的 GradScaler 实现并返回它。

accelerate.utils.get_mixed_precision_context_manager

( native_amp: bool = False autocast_kwargs: AutocastKwargs = None )

参数

native_amp (bool, 可选, 默认为 False) — 是否实际启用混合精度。
cache_enabled (bool, 可选, 默认为 True) — 是否应启用 autocast 内部的权重缓存。

返回一个用于自动转换混合精度的上下文管理器

accelerate.utils.listify

( data )

参数

data (嵌套的 list/tuple/dictionary of torch.Tensor) — 要转换为常规数字的数据。

递归地在嵌套的 list/tuple/dictionary 中查找张量，并将它们转换为数字列表。

accelerate.utils.pad_across_processes

( tensor dim = 0 pad_index = 0 pad_first = False )

参数

tensor (嵌套的 list/tuple/dictionary of torch.Tensor) — 要收集的数据。
dim (int, 可选, 默认为 0) — 填充的维度。
pad_index (int, 可选, 默认为 0) — 用于填充的值。
pad_first (bool, 可选, 默认为 False) — 是在开头还是在末尾填充。

递归地将来自所有设备的嵌套 list/tuple/dictionary 中的张量填充到相同的大小，以便它们可以安全地被收集。

accelerate.utils.recursively_apply

( func data *args test_type = <function is_torch_tensor at 0x7f9bc05c5360> error_on_other_type = False **kwargs )

参数

func (callable) — 要递归应用的函数。
data (嵌套的 list/tuple/dictionary of main_type) — 要应用 func 的数据
*args — 将传递给 func 的位置参数，应用于解包后的数据。
main_type (type, 可选, 默认为 torch.Tensor) — 应用 func 的对象的基础类型。
error_on_other_type (bool, 可选, 默认为 False) — 解包 data 后，如果遇到非 main_type 类型的对象，是否返回错误。如果为 False，函数将保持非 main_type 类型的对象不变。
**kwargs (附加关键字参数, 可选) — 将传递给 func 的关键字参数，应用于解包后的数据。

在嵌套的 list/tuple/dictionary 数据结构上递归地应用一个函数，该数据结构具有给定的基础类型。

accelerate.utils.reduce

( tensor reduction = 'mean' scale = 1.0 )

参数

tensor (嵌套的 list/tuple/dictionary of torch.Tensor) — 要归约的数据。
reduction (str, 可选, 默认为 "mean") — 一种归约方法。可以是“mean”、“sum”或“none”。
scale (float, 可选) — 归约后应用的默认缩放值，仅在 XLA 上有效。

通过给定操作的均值，递归地归约所有进程中嵌套的 list/tuple/dictionary 中的张量列表。

accelerate.utils.send_to_device

( tensor device non_blocking = False skip_keys = None )

参数

tensor (torch.Tensor 的嵌套列表/元组/字典) — 要发送到指定设备的数据。
device (torch.device) — 要将数据发送到的设备。

以递归方式将嵌套列表/元组/字典中的张量元素发送到指定设备。

accelerate.utils.slice_tensors

( data tensor_slice process_index = None num_processes = None )

参数

data (torch.Tensor 的嵌套列表/元组/字典) — 要切片的数据。
tensor_slice (slice) — 要进行的切片操作。

递归地对嵌套列表/元组/字典中的张量进行切片。

环境检查

这些功能用于检查当前工作环境的状态，包括操作系统本身的信息、其支持的功能以及是否安装了特定的依赖项。

accelerate.utils.is_bf16_available

( ignore_tpu = False )

检查是否支持bf16，可选择忽略TPU

accelerate.utils.is_ipex_available

( )

检查ipex是否已安装。

accelerate.utils.is_mps_available

( min_version = '1.12' )

检查MPS设备是否可用。最低版本要求为1.12。

accelerate.utils.is_npu_available

( check_device = False )

检查 `torch_npu` 是否已安装，并可能检查环境中是否存在 NPU。

accelerate.utils.is_torch_version

( operation: str version: str )

参数

operation (str) — 运算符的字符串表示，例如 `">"` 或 `"<="`
version (str) — PyTorch 的字符串版本号

使用给定的操作符，将当前 PyTorch 版本与参考版本进行比较。

accelerate.utils.is_torch_xla_available

( check_is_tpu = False check_is_gpu = False )

检查 `torch_xla` 是否可用。要在已安装 torch_xla 的环境中训练原生的 PyTorch 作业，请将 USE_TORCH_XLA 设置为 false。

accelerate.utils.is_xpu_available

( check_device = False )

检查 XPU 加速是否可用，可以是 `intel_extension_for_pytorch` 或原生 PyTorch (>=2.4)，并可能检查环境中是否存在 XPU。

环境操作

accelerate.utils.patch_environment

( **kwargs )

一个上下文管理器，它会将每个传入的关键字参数添加到 `os.environ` 中，并在退出时将其移除。

会将 `kwargs` 中的值转换为字符串，并将所有键转换为大写。

示例

>>> import os
>>> from accelerate.utils import patch_environment

>>> with patch_environment(FOO="bar"):
...     print(os.environ["FOO"])  # prints "bar"
>>> print(os.environ["FOO"])  # raises KeyError

accelerate.utils.clear_environment

( )

一个上下文管理器，它会临时清除环境变量。

当此上下文退出时，之前的环境变量将被恢复。

示例

>>> import os
>>> from accelerate.utils import clear_environment

>>> os.environ["FOO"] = "bar"
>>> with clear_environment():
...     print(os.environ)
...     os.environ["FOO"] = "new_bar"
...     print(os.environ["FOO"])
{}
new_bar

>>> print(os.environ["FOO"])
bar

accelerate.commands.config.default.write_basic_config

( mixed_precision = 'no' save_location: str = '/github/home/.cache/huggingface/accelerate/default_config.yaml' )

参数

mixed_precision (str, 可选, 默认为 “no”) — 要使用的混合精度。应为 “no”、“fp16” 或 “bf16” 之一。
save_location (str, 可选, 默认为 default_json_config_file) — 可选的自定义保存位置。在使用 `accelerate launch` 时应传递给 `--config_file`。默认位置在 huggingface 缓存文件夹内（`~/.cache/huggingface`），但可以通过设置 `HF_HOME` 环境变量来覆盖，其后跟 `accelerate/default_config.yaml`。

创建并保存一个基础集群配置，用于可能具有多个 GPU 的本地机器。如果机器只有 CPU，也会设置 CPU。

首次设置 🤗 Accelerate 时，可以使用 [~utils.write_basic_config] 作为快速配置的替代方案，而不是运行 `accelerate config`。

accelerate.utils.set_numa_affinity

( local_process_index: int verbose: typing.Optional[bool] = None )

参数

local_process_index (int) — 当前服务器上当前进程的索引。
verbose (bool, 可选) — 是否为每个进程打印新的 CPU 核心分配。如果启用了 `ACCELERATE_DEBUG_MODE`，则默认为 True。

将当前进程分配给特定的 NUMA 节点。理想情况下，每个节点至少有 2 个 CPU 时效率最高。

此结果在调用之间会被缓存。如果要覆盖它，请使用 `accelerate.utils.environment.override_numa_afifnity`。

accelerate.utils.environment.override_numa_affinity

( local_process_index: int verbose: typing.Optional[bool] = None )

参数

local_process_index (int) — 当前服务器上当前进程的索引。
verbose (bool, 可选) — 是否打印每个 CPU 的分配情况。如果启用了 `ACCELERATE_DEBUG_MODE`，则默认为 True。

覆盖为当前进程设置的任何 NUMA 亲和性。这个操作非常耗费资源，需要重新计算亲和性来设置，理想情况下应使用 `utils.environment.set_numa_affinity` 代替。

accelerate.utils.purge_accelerate_environment

( func_or_cls )

用于清理由被装饰的类或函数设置的 accelerate 环境变量的装饰器。

在某些情况下，调用特定的类或函数可能会导致 accelerate 环境变量被设置，并且之后没有被清理。例如，当调用

TrainingArguments(fp16=True, …)

将设置以下环境变量

ACCELERATE_MIXED_PRECISION=fp16

这可能会影响后续代码，因为环境变量的优先级高于 TrainingArguments(fp16=False)。这对于单元测试尤其重要，我们希望避免各个测试之间产生副作用。用此装饰器装饰单元测试函数或整个类，以确保每次测试后环境变量都被清理。这适用于 unittest.TestCase 和普通类 (pytest)；当装饰父类时也有效。

内存

accelerate.find_executable_batch_size

( function: <built-in function callable> = None starting_batch_size: int = 128 reduce_batch_size_fn: <built-in function callable> = None )

参数

function (callable, 可选) — 要包装的函数
starting_batch_size (int, 可选) — 尝试放入内存的批处理大小

一个基本的装饰器，它会尝试执行 `function`。如果因内存不足或 CUDNN 相关的异常而失败，批处理大小将乘以 0.9 并传递给 `function`。

`function` 必须接受一个 `batch_size` 参数作为其第一个参数。

示例

>>> from accelerate.utils import find_executable_batch_size


>>> @find_executable_batch_size(starting_batch_size=128)
... def train(batch_size, model, optimizer):
...     ...


>>> train(model, optimizer)

模型

这些工具与 PyTorch 模型的交互有关

accelerate.utils.calculate_maximum_sizes

( model: Module )

计算模型的总大小及其最大层的尺寸

accelerate.utils.compute_module_sizes

( model: Module dtype: typing.Union[torch.device, str, NoneType] = None special_dtypes: typing.Optional[dict[str, typing.Union[str, torch.device]]] = None buffers_only: bool = False )

计算给定模型的每个子模块的大小。

accelerate.utils.extract_model_from_parallel

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

参数

model (torch.nn.Module) — 要提取的模型。
keep_fp32_wrapper (bool, 可选) — 是否从模型中移除混合精度钩子。
keep_torch_compile (bool, 可选) — 是否解包已编译的模型。
recursive (bool, 可选, 默认为 False) — 是否递归提取 `model` 中所有的 `module.module` 实例，以及递归解包子层，而不仅仅是顶层的分布式容器。

返回

torch.nn.Module

提取出的模型。

从其分布式容器中提取模型。

accelerate.utils.get_balanced_memory

( model: Module max_memory: typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None no_split_module_classes: typing.Optional[list[str]] = None dtype: typing.Union[str, torch.dtype, NoneType] = None special_dtypes: typing.Optional[dict[str, typing.Union[str, torch.device]]] = None low_zero: 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）。
low_zero (bool, 可选) — 最小化 GPU 0 上的权重数量，这在 GPU 0 用于其他操作（如 Transformers 的 generate 函数）时很方便。

为 infer_auto_device_map() 计算一个 `max_memory` 字典，该字典将平衡每个可用 GPU 的使用。

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

accelerate.utils.get_max_layer_size

( modules: list module_sizes: dict no_split_module_classes: list ) → Tuple[int, List[str]]

参数

modules (List[Tuple[str, torch.nn.Module]]) — 我们要确定最大层大小的命名模块列表。
module_sizes (Dict[str, int]) — 一个将每个层名称映射到其大小的字典（由 `compute_module_sizes` 生成）。
no_split_module_classes (List[str]) — 我们不希望被分割的层的类名列表。

返回

Tuple[int, List[str]]

层的最大大小以及实现该最大大小的层名称列表。

一个工具函数，它将扫描命名模块列表并返回一个完整层所使用的最大大小。层的定义为

没有直接子模块（只有参数和缓冲区）的模块
其类名在 `no_split_module_classes` 列表中的模块

accelerate.infer_auto_device_map

( model: Module max_memory: typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None no_split_module_classes: typing.Optional[list[str]] = None dtype: typing.Union[str, torch.dtype, NoneType] = None special_dtypes: typing.Optional[dict[str, typing.Union[str, torch.dtype]]] = None verbose: bool = False clean_result: bool = True offload_buffers: bool = False fallback_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 的最大尺寸的层。
如果需要卸载到 CPU，我们不会超过 CPU 上的可用 RAM。
如果需要卸载到磁盘，CPU 上始终留有空间来放回卸载到磁盘的最大尺寸的层。

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

accelerate.load_checkpoint_in_model

( model: Module checkpoint: typing.Union[str, os.PathLike] device_map: typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None offload_folder: typing.Union[str, os.PathLike, NoneType] = None dtype: typing.Union[str, torch.dtype, NoneType] = None offload_state_dict: bool = False offload_buffers: bool = False keep_in_fp32_modules: list = None offload_8bit_bnb: bool = False strict: bool = False full_state_dict: bool = True broadcast_from_rank0: bool = False )

参数

model (torch.nn.Module) — 我们希望加载检查点的模型。
checkpoint (str or os.PathLike) — 要加载的检查点文件夹。它可以是：
- 包含整个模型状态字典的文件的路径
- 包含分片检查点索引的 .json 文件的路径
- 包含一个唯一的 .index.json 文件和检查点分片的文件夹路径
- 包含一个唯一的 pytorch_model.bin 或 model.safetensors 文件的文件夹路径。
device_map (Dict[str, Union[int, str, torch.device]], 可选) — 一个指定每个子模块应放置位置的映射。它不需要细化到每个参数/缓冲区名称，一旦给定的模块名称在其中，其所有子模块都将被发送到同一设备。
offload_folder (str or os.PathLike, 可选) — 如果 device_map 包含任何值 "disk"，则为我们将卸载权重的文件夹。
dtype (str or torch.dtype, 可选) — 如果提供，权重在加载时将被转换为该类型。
offload_state_dict (bool, 可选, 默认为 False) — 如果为 True，将临时将 CPU 状态字典卸载到硬盘，以避免在 CPU 状态字典的权重 + 最大分片的权重不适合时耗尽 CPU RAM。
offload_buffers (bool, 可选, 默认为 False) — 是否将缓冲区包含在卸载到磁盘的权重中。
keep_in_fp32_modules(List[str], 可选) — 保持 torch.float32 dtype 的模块列表。
offload_8bit_bnb (bool, 可选) — 是否启用在 cpu/磁盘上卸载 8-bit 模块。
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，并将其中的张量逐个广播到其他 rank。其他 rank 将接收张量，并根据模型中的本地分片进行分片（如果适用）。

将一个（可能分片的）检查点加载到模型中，可能会在加载时将权重发送到给定的设备。

在跨设备加载后，您仍需要在模型上调用 dispatch_model() 以使其能够运行。要将检查点加载和调度组合成一个单一调用，请使用 load_checkpoint_and_dispatch()。

accelerate.utils.load_offloaded_weights

( model index offload_folder )

参数

model (torch.nn.Module) — 要加载权重的模型。
index (dict) — 一个字典，包含从模型卸载的每个参数的参数名称及其元数据。
offload_folder (str) — 存储卸载权重的文件夹。

将卸载文件夹中的权重加载到模型中。

accelerate.utils.load_state_dict

( checkpoint_file device_map = None )

参数

checkpoint_file (str) — 要加载的检查点路径。
device_map (Dict[str, Union[int, str, torch.device]], 可选) — 一个指定每个子模块应放置位置的映射。它不需要细化到每个参数/缓冲区名称，一旦给定的模块名称在其中，其所有子模块都将被发送到同一设备。

从给定文件加载检查点。如果检查点是 safetensors 格式并且传递了设备映射，权重可以快速直接加载到 GPU 上。

accelerate.utils.offload_state_dict

( save_dir: typing.Union[str, os.PathLike] state_dict: dict )

参数

save_dir (str or os.PathLike) — 卸载状态字典的目录。
state_dict (Dict[str, torch.Tensor]) — 要卸载的张量字典。

在给定文件夹中卸载状态字典。

accelerate.utils.retie_parameters

( model tied_params )

参数

model (torch.nn.Module) — 要重新绑定参数的模型。
tied_params (List[List[str]]) — 通过 find_tied_parameters 获得的参数名称到绑定参数名称的映射。

如果链接被破坏（例如在添加钩子时），则在给定模型中重新绑定参数。

accelerate.utils.set_module_tensor_to_device

( module: Module tensor_name: str device: typing.Union[int, str, torch.device] value: typing.Optional[torch.Tensor] = None dtype: typing.Union[str, torch.dtype, NoneType] = None fp16_statistics: typing.Optional[torch.HalfTensor] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None non_blocking: bool = False clear_cache: bool = True )

参数

module (torch.nn.Module) — 我们要移动的张量所在的模块。
tensor_name (str) — 参数/缓冲区的全名。
device (int, str or torch.device) — 设置张量的设备。
value (torch.Tensor, 可选) — 张量的值（在从元设备移动到任何其他设备时很有用）。
dtype (torch.dtype, 可选) — 如果与参数值一起传递，参数的值将被转换为此 dtype。否则，value 将被转换为模型中现有参数的 dtype。
fp16_statistics (torch.HalfTensor, 可选) — 要在模块上设置的 fp16 统计信息列表，用于 8 位模型序列化。
tied_params_map (Dict[int, Dict[torch.device, torch.Tensor]], 可选, 默认为 None) — 当前数据指针到已调度绑定权重的设备字典的映射。对于给定的执行设备，此参数有助于重用设备上共享权重的第一个可用指针，而不是复制内存。
non_blocking (bool, 可选, 默认为 False) — 如果为 True，设备传输将相对于主机异步进行（如果可能）。
clear_cache (bool, 可选, 默认为 True) — 在设备上设置张量后是否清除设备缓存。

一个辅助函数，用于将模块的给定张量（参数或缓冲区）设置到特定设备上（请注意，执行 param.to(device) 会创建一个与参数不关联的新张量，这就是为什么我们需要这个函数）。

accelerate.utils.get_module_children_bottom_up

( model: Module return_fqns: bool = False ) → list[torch.nn.Module]

参数

model (torch.nn.Module) — 要获取其子模块的模型

返回

list[torch.nn.Module]

一个按自下而上顺序排列的 model 的子模块列表。最后一个元素是 model 本身。

按自下而上的顺序遍历模型，并按该顺序返回子模块。

并行

这些包括在并行工作时应使用的通用工具。

accelerate.utils.extract_model_from_parallel

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

参数

model (torch.nn.Module) — 要提取的模型。
keep_fp32_wrapper (bool, 可选) — 是否从模型中移除混合精度钩子。
keep_torch_compile (bool, 可选) — 是否解包已编译的模型。
recursive (bool, 可选, 默认为 False) — 是否递归提取 `model` 中所有的 `module.module` 实例，以及递归解包子层，而不仅仅是顶层的分布式容器。

返回

torch.nn.Module

提取出的模型。

从其分布式容器中提取模型。

accelerate.utils.save

( obj f save_on_each_node: bool = False safe_serialization: bool = False )

参数

obj — 要保存的数据
f — 用于保存数据的文件（或类文件对象）
save_on_each_node (bool, 可选, 默认为 False) — 是否仅在全局主进程上保存
safe_serialization (bool, 可选, 默认为 False) — 是否使用 safetensors 或传统的 PyTorch 方式（使用 pickle）保存 obj。

将数据保存到磁盘。代替 torch.save() 使用。

accelerate.utils.load

( f map_location = None **kwargs )

参数

f — 用于加载数据的文件（或类文件对象）
map_location — 一个函数、torch.device、字符串或指定如何重新映射存储位置的字典
**kwargs — 传递给 torch.load() 的额外关键字参数。

torch.load() 的兼容性替代品，如果 torch 版本为 2.4.0 或更高，则允许使用 weights_only。否则将忽略该关键字参数。

还会为 numpy 数组添加（然后移除）一个异常处理

accelerate.utils.wait_for_everyone

( )

在脚本中引入一个阻塞点，确保所有进程在继续之前都已达到此点。

请确保所有进程都会到达此指令，否则您的某个进程将永远挂起。

随机

这些工具与设置和同步所有随机状态有关。

accelerate.utils.set_seed

( seed: int device_specific: bool = False deterministic: bool = False )

参数

seed (int) — 要设置的种子。
device_specific (bool, 可选, 默认为 False) — 是否根据 self.process_index 在每个设备上略微区分种子。
deterministic (bool, 可选, 默认为 False) — 在可用时是否使用确定性算法。可能会减慢训练速度。

用于设置 random、numpy、torch 中种子的辅助函数，以实现可复现的行为。

accelerate.utils.synchronize_rng_state

( rng_type: typing.Optional[accelerate.utils.dataclasses.RNGType] = None generator: typing.Optional[torch._C.Generator] = None )

accelerate.synchronize_rng_states

( rng_types: list generator: typing.Optional[torch._C.Generator] = None )

PyTorch XLA

这些是使用 PyTorch with XLA 时有用的实用工具。

accelerate.utils.install_xla

( upgrade: bool = False )

参数

upgrade (bool, 可选, 默认为 False) — 是否升级 torch 并安装最新的 torch_xla wheels。

根据 Google Colaboratory 中的 torch 版本安装合适的 xla wheels 的辅助函数。

示例

>>> from accelerate.utils import install_xla

>>> install_xla(upgrade=True)

加载模型权重

这些包括加载检查点时有用的实用工具。

accelerate.load_checkpoint_in_model

( model: Module checkpoint: typing.Union[str, os.PathLike] device_map: typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None offload_folder: typing.Union[str, os.PathLike, NoneType] = None dtype: typing.Union[str, torch.dtype, NoneType] = None offload_state_dict: bool = False offload_buffers: bool = False keep_in_fp32_modules: list = None offload_8bit_bnb: bool = False strict: bool = False full_state_dict: bool = True broadcast_from_rank0: bool = False )

参数

model (torch.nn.Module) — 我们希望加载检查点的模型。
checkpoint (str or os.PathLike) — 要加载的检查点文件夹。它可以是：
- 包含整个模型状态字典的文件的路径
- 包含分片检查点索引的 .json 文件的路径
- 包含一个唯一的 .index.json 文件和检查点分片的文件夹路径
- 包含一个唯一的 pytorch_model.bin 或 model.safetensors 文件的文件夹路径。
device_map (Dict[str, Union[int, str, torch.device]], 可选) — 一个指定每个子模块应放置位置的映射。它不需要细化到每个参数/缓冲区名称，一旦给定的模块名称在其中，其所有子模块都将被发送到同一设备。
offload_folder (str or os.PathLike, 可选) — 如果 device_map 包含任何值 "disk"，则为我们将卸载权重的文件夹。
dtype (str or torch.dtype, 可选) — 如果提供，权重在加载时将被转换为该类型。
offload_state_dict (bool, 可选, 默认为 False) — 如果为 True，将临时将 CPU 状态字典卸载到硬盘，以避免在 CPU 状态字典的权重 + 最大分片的权重不适合时耗尽 CPU RAM。
offload_buffers (bool, 可选, 默认为 False) — 是否将缓冲区包含在卸载到磁盘的权重中。
keep_in_fp32_modules(List[str], 可选) — 保持 torch.float32 dtype 的模块列表。
offload_8bit_bnb (bool, 可选) — 是否启用在 cpu/磁盘上卸载 8-bit 模块。
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，并将其中的张量逐个广播到其他 rank。其他 rank 将接收张量，并根据模型中的本地分片进行分片（如果适用）。

将一个（可能分片的）检查点加载到模型中，可能会在加载时将权重发送到给定的设备。

在跨设备加载后，您仍需要在模型上调用 dispatch_model() 以使其能够运行。要将检查点加载和调度组合成一个单一调用，请使用 load_checkpoint_and_dispatch()。

量化

这些包括用于量化模型的实用工具。

accelerate.utils.load_and_quantize_model

( model: Module bnb_quantization_config: BnbQuantizationConfig weights_location: typing.Union[str, os.PathLike] = None device_map: typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None no_split_module_classes: typing.Optional[list[str]] = None max_memory: typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None offload_folder: typing.Union[str, os.PathLike, NoneType] = None offload_state_dict: bool = False ) → torch.nn.Module

参数

model (torch.nn.Module) — 输入模型。模型可以已加载或在元设备（meta device）上。
bnb_quantization_config (BnbQuantizationConfig) — bitsandbytes 量化参数。
weights_location (str or os.PathLike) — 要加载的权重文件夹。它可以是：
- 包含整个模型状态字典的文件的路径
- 指向分片检查点索引的 .json 文件的路径
- 包含唯一 .index.json 文件和检查点分片的文件夹的路径
- 包含唯一 pytorch_model.bin 文件的文件夹的路径
device_map (Dict[str, Union[int, str, torch.device]], 可选) — 指定每个子模块应放置位置的映射。无需细化到每个参数/缓冲区名称，一旦给定的模块名称在其中，其所有子模块都将被发送到同一设备。
no_split_module_classes (List[str], 可选) — 不应跨设备分割的层类名称列表（例如任何具有残差连接的层）。
max_memory (Dict, 可选) — 设备标识符到最大内存的字典。如果未设置，则默认为可用的最大内存。
offload_folder (str or os.PathLike, 可选) — 如果 device_map 包含任何 "disk" 值，则为我们将卸载权重的文件夹。
offload_state_dict (bool, 可选, 默认为 False) — 如果为 True，则会临时将 CPU 状态字典卸载到硬盘上，以避免在 CPU 状态字典的权重 + 最大分片的权重不适合时耗尽 CPU RAM。

返回

torch.nn.Module

量化后的模型

此函数将使用在 bnb_quantization_config 中传递的关联配置对输入模型进行量化。如果模型在元设备中，我们将根据传递的 device_map 加载和分派权重。如果模型已加载，我们将对模型进行量化并将模型放在 GPU 上。

< > 在 GitHub 上更新

←FP8 Megatron-LM 实用工具→