实用函数和类

以下是 🤗 Accelerate 提供的各种实用函数，按使用场景分类。

常量

在 🤗 Accelerate 中用于参考的常量

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 中广泛使用的基础 dataclass，可以作为参数传入。

独立

这些是用于检查的独立 dataclass，例如所使用的分布式系统的类型

类 accelerate.utils.ComputeEnvironment

< source >

( *values )

表示计算环境的类型。

值：

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

类 accelerate.DistributedType

< source >

( *values )

表示一种分布式环境类型。

值：

否 — 不是分布式环境，仅为单进程。
MULTI_CPU — 在多个 CPU 节点上进行分布式。
MULTI_GPU — 在多个 GPU 上分布式运行。
MULTI_MLU — 在多个 MLU 上进行分布式。
MULTI_SDAA — 分布在多个 SDAAs 上。
MULTI_MUSA — 在多个 MUSA 上分布式运行。
MULTI_NPU — 在多个 NPU 上进行分布式。
MULTI_XPU — 在多个 XPU 上进行分布式。
MULTI_HPU — 在多个 HPU 上进行分布式。
MULTI_NEURON — 分布在多个 Neuron 核心上。
DEEPSPEED — 使用 DeepSpeed。
FSDP — 使用完全分片数据并行（FSDP）。
XLA — 使用 TorchXLA。
MEGATRON_LM — 使用 Megatron-LM。

类 accelerate.utils.DynamoBackend

< source >

( *values )

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

值：

否 — 请勿使用 torch dynamo。
EAGER — 使用 PyTorch 运行提取的 GraphModule。这在调试 TorchDynamo 问题时非常有用。
AOT_EAGER — 使用 AotAutograd 但不使用编译器，即仅使用 PyTorch eager 来执行 AotAutograd 提取的前向和反向图。这有助于调试，但通常不会带来性能提升。
INDUCTOR — 通过利用代码生成的 Triton 内核，使用带有 AotAutograd 和 cudagraphs 的 TorchInductor 后端。阅读更多
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 优化，用于推理。阅读更多
TVM — 使用 Apache TVM 进行推理优化。阅读更多
HPU_BACKEND — 使用 HPU 后端进行推理优化。
NEURON — 使用 AWS Neuron 后端来支持 Trainium/Inferentia。

类 accelerate.utils.LoggerType

< source >

( *values )

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

值：

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

类 accelerate.utils.PrecisionType

< source >

( *values )

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

值：

否 — 使用全精度（FP32）
FP16 — 使用半精度
BF16 — 使用脑浮点精度

类 accelerate.utils.RNGType

< source >

( *values )

class accelerate.utils.SageMakerDistributedType

< source >

( *values )

表示一种分布式环境类型。

值：

否 — 不是分布式环境，仅为单进程。
DATA_PARALLEL — 使用 sagemaker 分布式数据并行。
MODEL_PARALLEL — 使用 sagemaker 分布式模型并行。

关键字参数

这些是 PyTorch 生态系统中特定交互的可配置参数，由 Accelerate 在底层进行处理。

类 accelerate.AutocastKwargs

< source >

( enabled: bool = Truecache_enabled: bool | None = None )

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

示例：

from accelerate import Accelerator
from accelerate.utils import AutocastKwargs

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

类 accelerate.DistributedDataParallelKwargs

< source >

( dim: int = 0broadcast_buffers: bool = Truebucket_cap_mb: int = 25find_unused_parameters: bool = Falsecheck_reduction: bool = Falsegradient_as_bucket_view: bool = Falsestatic_graph: bool = Falsecomm_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 包装。有关每个参数的更多信息，请参阅此 wrapper 的文档。

示例：

from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs

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

类 accelerate.GradScalerKwargs

< source >

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

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

示例：

from accelerate import Accelerator
from accelerate.utils import GradScalerKwargs

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

类 accelerate.InitProcessGroupKwargs

< source >

( backend: str | None = 'nccl'init_method: str | None = Nonetimeout: datetime.timedelta | None = 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])

类 accelerate.utils.KwargsHandler

< source >

( )

实现 to_kwargs() 方法的内部混入类，用于 dataclass。

to_kwargs

< source >

( )

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

插件

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

类 accelerate.DeepSpeedPlugin

< source >

( hf_ds_config: Any = Nonegradient_accumulation_steps: int = Nonegradient_clipping: float = Nonezero_stage: int = Noneis_train_batch_min: bool = Trueoffload_optimizer_device: str = Noneoffload_param_device: str = Noneoffload_optimizer_nvme_path: str = Noneoffload_param_nvme_path: str = Nonezero3_init_flag: bool = Nonezero3_save_16bit_model: bool = Nonetransformer_moe_cls_names: str = 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 等。

此插件用于集成 DeepSpeed。

deepspeed_config_process

< source >

( prefix = ''mismatches = Noneconfig = Nonemust_match = True**kwargs )

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

选择

< source >

( _from_accelerator_state: bool = False )

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

类 accelerate.FullyShardedDataParallelPlugin

< source >

( fsdp_version: int = Nonesharding_strategy: str | ForwardRef('torch.distributed.fsdp.ShardingStrategy') = Nonereshard_after_forward: str | ForwardRef('torch.distributed.fsdp.ShardingStrategy') | bool = Nonebackward_prefetch: str | ForwardRef('torch.distributed.fsdp.BackwardPrefetch') | None = Nonemixed_precision_policy: dict | str | ForwardRef('torch.distributed.fsdp.MixedPrecision') | ForwardRef('torch.distributed.fsdp.MixedPrecisionPolicy') | None = Noneauto_wrap_policy: typing.Callable | typing.Literal['transformer_based_wrap', 'size_based_wrap', 'no_wrap'] | None = Nonecpu_offload: bool | ForwardRef('torch.distributed.fsdp.CPUOffload') | ForwardRef('torch.distributed.fsdp.CPUOffloadPolicy') = Noneignored_modules: collections.abc.Iterable[torch.nn.Module] | str | None = Nonestate_dict_type: str | ForwardRef('torch.distributed.fsdp.StateDictType') = Nonestate_dict_config: ForwardRef('torch.distributed.fsdp.FullStateDictConfig') | ForwardRef('torch.distributed.fsdp.ShardedStateDictConfig') | None = Noneoptim_state_dict_config: ForwardRef('torch.distributed.fsdp.FullOptimStateDictConfig') | ForwardRef('torch.distributed.fsdp.ShardedOptimStateDictConfig') | None = Nonelimit_all_gathers: bool = Trueuse_orig_params: bool | None = Noneparam_init_fn: typing.Callable[[torch.nn.Module], NoneType] | None = Nonesync_module_states: bool | None = Noneforward_prefetch: bool = Noneactivation_checkpointing: bool = Noneactivation_checkpointing_offload: bool = Nonecpu_ram_efficient_loading: bool = Nonetransformer_cls_names_to_wrap: list[str] | None = Nonemin_num_params: int | None = 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，则应为 bool 类型；否则应为 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, str, torch.distributed.fsdp.MixedPrecision, torch.distributed.fsdp.MixedPrecisionPolicy]]，默认为 None) — 用于启用 FullyShardedDataParallel 混合精度训练的配置。如果传入 dict，它应包含以下键：param_dtype、reduce_dtype 和 buffer_dtype；如果 fsdp_version 设置为 2，则可以是 torch.distributed.fsdp.MixedPrecisionPolicy 的实例。如果传入 str，它应为以下值之一：fp8、fp16、bf16、fp32，用于设置 param_dtype、reduce_dtype 和 buffer_dtype。
auto_wrap_policy (Optional(Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]), defaults to NO_WRAP) -- A callable or string specifying a policy to recursively wrap layers with FSDP. If a string, it must be one of transformer_based_wrap, size_based_wrap, or no_wrap. See torch.distributed.fsdp.wrap.size_based_wrap_policy` for a direction on what it should look like.
cpu_offload (Union[bool, torch.distributed.fsdp.CPUOffload, torch.distributed.fsdp.CPUOffloadPolicy]，默认为 False) — 是否将参数卸载到 CPU。应为 bool 类型，或者当 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 按名称匹配模块。如果 fsdp_version 设置为 2，这些模块将被转换为参数并使用。
state_dict_type (Union[str, torch.distributed.fsdp.StateDictType]，默认为 'FULL_STATE_DICT') — 要使用的 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 配置。如果未传入，则根据 state_dict_type 确定。
optim_state_dict_config (Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig]，默认为 None) — 要使用的 Optim state dict 配置。如果未传入，则根据 state_dict_type 确定。
limit_all_gathers (bool，默认为 True) — 是否让 FSDP 显式同步 CPU 线程，以防止过多的 in-flight all-gathers。此 bool 仅影响调度 all-gathers 的分片策略。启用此选项有助于减少 CUDA malloc 重试次数。
use_orig_params (bool，默认为 False) — 是否使用原始参数作为优化器的参数。
param_init_fn (Optional[Callable[[torch.nn.Module], None]，默认为 None) — 一个 Callable[torch.nn.Module] -> None，用于指定当前位于 meta 设备上的模块应如何初始化到实际设备上。仅当 sync_module_states 为 True 时适用。默认是一个 lambda，它会对模块调用 to_empty。
sync_module_states (bool，默认为 False) — 是否应让每个单独包装的 FSDP 单元从 rank 0 广播模块参数，以确保初始化后所有 rank 上的参数一致。除非 cpu_ram_efficient_loading 为 True，否则默认为 False；若为 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

< source >

( model )

给定 model，根据传入的策略创建 auto_wrap_policy，并判断是否可以使用 transformer_cls_to_wrap

set_mixed_precision

< source >

( mixed_precisionbuffer_autocast = Falseoverride = False )

设置 FSDP 的混合精度策略

set_state_dict_type

< source >

( state_dict_type = None )

根据 StateDictType 设置状态字典配置。

validate_mixed_precision_policy

< source >

( )

验证混合精度策略，将其抽象化，以便在不需要时避免引入相关导入。

类 accelerate.utils.GradientAccumulationPlugin

< source >

( num_steps: int = Noneadjust_scheduler: bool = Truesync_with_dataloader: bool = Truesync_each_batch: bool = False )

参数

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

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

示例：

from accelerate.utils import GradientAccumulationPlugin

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

类 accelerate.utils.MegatronLMPlugin

< source >

( tp_degree: int = Nonepp_degree: int = Noneuse_custom_fsdp: bool = Noneoverlap_cpu_optimizer_d2h_h2d: bool = Noneno_load_optim: bool = Noneeod_mask_loss: bool = Noneno_save_optim: bool = Noneoptimizer_cpu_offload: bool = Noneuse_precision_aware_optimizer: bool = Nonedecoder_last_pipeline_num_layers: int = Nonerecompute_granularity: str = Nonerecompute_method: str = Nonerecompute_num_layers: int = Noneattention_backend: bool = Noneexpert_model_parallel_size: int = Nonecontext_parallel_size: int = Noneattention_dropout: float = Nonehidden_dropout: float = Noneattention_softmax_in_fp32: bool = Noneexpert_tensor_parallel_size: int = Nonecalculate_per_token_loss: bool = Noneuse_rotary_position_embeddings: bool = Nonenum_micro_batches: int = Nonegradient_clipping: float = Nonesequence_parallelism: bool = Nonerecompute_activations: bool = Noneuse_distributed_optimizer: bool = Nonepipeline_model_parallel_split_rank: int = Nonenum_layers_per_virtual_pipeline_stage: int = Noneis_train_batch_min: str = Truetrain_iters: int = Nonetrain_samples: int = Noneweight_decay_incr_style: str = 'constant'start_weight_decay: float = Noneend_weight_decay: float = Nonelr_decay_style: str = 'linear'lr_decay_iters: int = Nonelr_decay_samples: int = Nonelr_warmup_iters: int = Nonelr_warmup_samples: int = Nonelr_warmup_fraction: float = Nonemin_lr: float = 0consumed_samples: list = Noneno_wd_decay_cond: typing.Callable | None = Nonescale_lr_cond: typing.Callable | None = Nonelr_mult: float = 1.0megatron_dataset_flag: bool = Falseseq_length: int = Noneencoder_seq_length: int = Nonedecoder_seq_length: int = Nonetensorboard_dir: str = Noneset_all_logging_options: bool = Falseeval_iters: int = 100eval_interval: int = 1000return_logits: bool = Falsecustom_train_step_class: typing.Any | None = Nonecustom_train_step_kwargs: dict[str, typing.Any] | None = Nonecustom_model_provider_function: typing.Callable | None = Nonecustom_prepare_model_function: typing.Callable | None = Nonecustom_megatron_datasets_provider_function: typing.Callable | None = Nonecustom_get_batch_function: typing.Callable | None = Nonecustom_loss_function: typing.Callable | None = Noneother_megatron_args: dict[str, typing.Any] | None = 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_optimizer (bool，默认为 None) — 启用分布式优化器。
pipeline_model_parallel_split_rank (int，默认为 None) — 编码器与解码器应拆分所在的 Rank。
num_layers_per_virtual_pipeline_stage (int，默认为 None) — 每个虚拟流水线阶段的层数。
is_train_batch_min (str，默认为 True) — 如果同时指定了训练和评估数据加载器，此参数将决定 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') — 权重衰减增量函数。choices=[“constant”, “linear”, “cosine”]。
start_weight_decay (float，默认为 None) — L2 正则化的初始权重衰减系数。
end_weight_decay (float，默认为 None) — L2 正则化中运行结束时的权重衰减系数。
lr_decay_style (str，默认为 'linear') — 学习率衰减函数。choices=[‘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-(iters/samples) 比例。
min_lr (float，默认为 0) — 学习率的最小值。调度器会将低于此阈值的值裁剪掉。
consumed_samples (List，默认为 None) — 与传入 accelerator.prepare 的 dataloader 顺序一致，已消耗的样本数量。
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 train_valid_test 数据集提供函数。
custom_get_batch_function (Optional，默认为 None) — 自定义的获取批次函数。
custom_loss_function (Optional，默认为 None) — 自定义损失函数。
other_megatron_args (Optional，默认为 None) — 其他 Megatron-LM 参数。请参阅 Megatron-LM。

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

类 accelerate.utils.TorchDynamoPlugin

< source >

( backend: DynamoBackend = Nonemode: str = Nonefullgraph: bool = Nonedynamic: bool = Noneoptions: Any = Nonedisable: bool = Falseuse_regional_compilation: bool = None )

参数

backend (DynamoBackend，默认为 None) — 一个有效的 Dynamo 后端。有关更多详细信息，请参阅 https://pytorch.org/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

< source >

( load_in_8bit: bool = Falsellm_int8_threshold: float = 6.0load_in_4bit: bool = Falsebnb_4bit_quant_type: str = 'fp4'bnb_4bit_use_double_quant: bool = Falsebnb_4bit_compute_dtype: str = 'fp16'torch_dtype: dtype = Noneskip_modules: list = Nonekeep_in_fp32_modules: list = None )

参数

load_in_8bit (bool，默认为 False) — 启用 8 位量化。
llm_int8_threshold (float，默认为 6.0) — 外轮廓阈值。仅在 load_in_8bit=True 时有效。
load_in_4bit (bool，默认为 False) — 启用 4bit 量化。
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) — 明确列出我们不予量化的模块。这些模块的数据类型将为 torch_dtype。
keep_in_fp32_modules (List，默认为 None) — 明确列出我们不予量化的模块。这些模块将保持为 torch.float32。

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

类 accelerate.DataLoaderConfiguration

< source >

( split_batches: bool = Falsedispatch_batches: bool = Noneeven_batches: bool = Trueuse_seedable_sampler: bool = Falsedata_seed: int = Nonenon_blocking: bool = Falseuse_stateful_dataloader: bool = False )

参数

split_batches (bool，默认为 False) — 是否让 accelerator 将 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 时，用于配置数据加载器相关项的设置。

类 accelerate.utils.ProjectConfiguration

< source >

( project_dir: str = Nonelogging_dir: str = Noneautomatic_checkpoint_naming: bool = Falsetotal_limit: int = Noneiteration: int = 0save_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

< source >

( project_dir: str | None = None )

将 self.project_dir 和 self.logging_dir 设置为相应的值。

环境变量

这些是可根据不同使用场景启用的环境变量

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

数据操作与运算

这些包括模拟相同 torch 操作的数据操作，但可用于分布式进程。

accelerate.utils.broadcast

< source >

( tensorfrom_process: int = 0 )

参数

tensor（嵌套列表/元组/字典，元素为 torch.Tensor）— 要收集的数据。
from_process (int，可选，默认为 0) — 发送数据的进程

在嵌套的张量列表/元组/字典中，递归地将张量广播到所有设备。

accelerate.utils.broadcast_object_list

< source >

( object_listfrom_process: int = 0 )

参数

object_list（可序列化对象列表）— 要广播的对象列表。此列表将被就地修改。
from_process (int，可选，默认为 0) — 发送数据的进程。

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

accelerate.utils.concatenate

< source >

( datadim = 0 )

参数

data（嵌套列表/元组/字典，其中包含张量列表 torch.Tensor）— 要拼接的数据。
dim (int，可选，默认为 0) — 用于拼接的维度。

递归地拼接嵌套列表/元组/字典中形状相同的张量列表中的张量。如果只有一个数据批次，则按原样返回。

accelerate.utils.convert_outputs_to_fp32

< source >

( model_forward )

accelerate.utils.convert_to_fp32

< source >

( tensor )

参数

tensor（嵌套列表/元组/字典，元素为 torch.Tensor）— 要从 FP16/BF16 转换为 FP32 的数据。

递归地将嵌套列表/元组/字典中 FP16/BF16 精度的张量元素转换为 FP32。

accelerate.utils.gather

< source >

( tensor )

参数

tensor（嵌套列表/元组/字典，元素为 torch.Tensor）— 要收集的数据。

递归地从所有设备上的嵌套列表/元组/字典形式的张量中收集张量。

accelerate.utils.gather_object

< source >

( object: Any )

参数

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

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

accelerate.utils.get_grad_scaler

< source >

( distributed_type: DistributedType = None**kwargs )

参数

distributed_type (DistributedType，可选，默认为 None) — 分布式环境的类型。
kwargs — 用于所使用 GradScaler 构造函数的额外参数。

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

accelerate.utils.get_mixed_precision_context_manager

< source >

( native_amp: bool = Falseautocast_kwargs: AutocastKwargs = None )

参数

native_amp (bool，可选，默认为 False) — 是否实际启用混合精度。

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

accelerate.utils.listify

< source >

( data )

参数

data（嵌套列表/元组/字典，元素为 torch.Tensor）— 用于转换为常规数字的数据。

递归地在嵌套的列表/元组/字典中查找张量，并将其转换为数字列表。

accelerate.utils.pad_across_processes

< source >

( tensordim = 0pad_index = 0pad_first = False )

参数

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

递归地对来自所有设备的嵌套列表/元组/字典中的张量进行填充，使其具有相同的大小，以便安全地收集。

accelerate.utils.recursively_apply

< source >

( funcdata*argstest_type = <function is_torch_tensor at 0x7f8c167ca8d0>error_on_other_type = False**kwargs )

参数

func (callable) — 要递归应用的函数。
data（嵌套列表/元组/字典，包含 test_type 接受的对象）— 用于应用 func 的数据
*args — 在应用于解包后的数据时，将传递给 func 的位置参数。
test_type (Callable[[Any], bool]，可选，默认为 is_torch_tensor) — 用于判断是否应对给定的叶对象应用 func 的谓词。
error_on_other_type (bool，可选，默认为 False) — 在解包 data 后，如果得到的对象使 test_type 返回 False，是否返回错误。如果为 False，该函数将保持此类对象不变。
**kwargs（额外的关键字参数，可选） — 在应用于解包后的数据时，将传递给 func 的关键字参数。

递归地对给定基础类型的嵌套列表/元组/字典数据结构应用函数。

accelerate.utils.reduce

< source >

( tensorreduction = 'mean'scale = 1.0 )

参数

tensor（嵌套列表/元组/字典，元素为 torch.Tensor）— 要归约的数据。
reduction (str，可选，默认为 "mean") — 归约方法。可以是 “mean”、“sum”、“max” 或 “none”
scale (float, 可选) — 在 reduce 之后应用的默认缩放值，仅在 XLA 上有效。

通过给定操作的均值，在所有进程上递归地归约嵌套列表/元组/字典中的张量列表中的张量。

accelerate.utils.send_to_device

< source >

( tensordevicenon_blocking = Falseskip_keys = None )

参数

tensor（嵌套列表/元组/字典，元素为 torch.Tensor）— 要发送到指定设备的数据。
device (torch.device) — 要发送数据的目标设备。
non_blocking (bool，可选，默认为 False) — 如果为 True，则向设备的数据传输将异步执行，从而可以与计算重叠进行。仅当设备支持此功能时（例如 CUDA）才有效。
skip_keys (str 或 List[str]，可选) — 字典 tensor 中不应发送到指定 device 的键或键列表。具有这些键的条目将保留在其原始设备上。

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

accelerate.utils.slice_tensors

< source >

( datatensor_sliceprocess_index = Nonenum_processes = None )

参数

data（嵌套列表/元组/字典，元素为 torch.Tensor）— 要切片的数据。
tensor_slice (slice) — 要提取的切片。

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

环境检查

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

accelerate.utils.is_bf16_available

< source >

( ignore_tpu = False )

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

accelerate.utils.is_mps_available

< source >

( min_version = '1.12' )

检查 MPS 设备是否可用。所需的最低版本为 1.12。

accelerate.utils.is_npu_available

( check_device = False )

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

accelerate.utils.is_torch_version

< source >

( operation: strversion: str )

参数

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

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

accelerate.utils.is_torch_xla_available

( check_is_tpu = Falsecheck_is_gpu = False )

检查 torch_xla 是否可用。要在已安装 torch xla 的环境中训练原生 pytorch 任务，请将 USE_TORCH_XLA 设置为 false。

accelerate.utils.is_xpu_available

( check_device = False )

检查是否可通过原生 PyTorch（>=2.7）使用 XPU 加速，以及环境中是否存在 XPU

环境操作

accelerate.utils.patch_environment

< source >

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

< source >

( )

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

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

示例：

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

< source >

( mixed_precision = 'no'save_location: str = '/home/runner/.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 时，除了运行 accelerate config，还可以使用 [~utils.write_basic_config] 作为快速配置的替代方案。

accelerate.utils.set_numa_affinity

( local_process_index: intverbose: bool | None = 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

< source >

( local_process_index: intverbose: bool | None = None )

参数

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

覆盖当前进程设置的任何 NUMA 亲和性。此操作开销极大，需要重新计算要设置的亲和性，理想情况下应改用 utils.environment.set_numa_affinity。

accelerate.utils.purge_accelerate_environment

< source >

( func_or_cls )

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

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

TrainingArguments(fp16=True, …)

将设置以下环境变量：

ACCELERATE_MIXED_PRECISION=fp16

这可能会影响后续代码，因为该环境变量会覆盖 TrainingArguments(fp16=False)。这在单元测试中尤为相关，因为我们希望避免各个测试之间产生相互影响。使用此装饰器修饰单元测试函数或整个类，以确保每个测试结束后，环境变量都会被清理。该装饰器同时适用于 unittest.TestCase 和普通类（pytest）；当用于修饰父类时同样有效。

内存

accelerate.find_executable_batch_size

< source >

( function: callable | None = Nonestarting_batch_size: int = 128reduce_batch_size_fn: callable | None = None )

参数

function (callable, optional) — 要包装的函数
starting_batch_size (int, 可选) — 尝试装入内存的批大小
reduce_batch_size_fn (callable, 可选) — 用于确定发生内存溢出错误后新批大小的函数。如果未提供，则每次失败时批大小会乘以 0.9。该函数不接受任何参数，并应返回新的（缩减后的）批大小，类型为 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

< source >

( model: Module )

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

accelerate.utils.compute_module_sizes

< source >

( model: Moduledtype: str | torch.device | None = Nonespecial_dtypes: dict[str, str | torch.device] | None = Nonebuffers_only: bool = False )

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

accelerate.utils.extract_model_from_parallel

< source >

( modelkeep_fp32_wrapper: bool = Truekeep_torch_compile: bool = Truerecursive: 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

< source >

参数

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 用于其他操作（例如 Transformers 的 generate 函数）时，这一特性非常实用。

计算一个 max_memory 字典，供 infer_auto_device_map() 使用，以平衡每个可用 GPU 的使用情况。

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

accelerate.utils.get_max_layer_size

< source >

( modules: listmodule_sizes: dictno_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

< source >

参数

model (torch.nn.Module) — 要分析的模型。
max_memory (Dict, 可选) — 一个从设备标识符到最大内存的字典。如果未设置，则默认为可用的最大内存。示例：max_memory={0: "1GB"}。
no_split_module_classes (List[str], 可选) — 不应跨设备拆分的层类名列表（例如任何具有残差连接的层）。
dtype (str 或 torch.dtype，可选) — 如果提供，加载时权重将被转换为该类型。
special_dtypes (Dict[str, Union[str, torch.device]], 可选) — 如果提供，则指定某些特定权重需要考虑的特殊数据类型（将覆盖用于所有权重的默认 dtype）。
verbose (bool，可选，默认为 False) — 在函数构建 device_map 时，是否提供调试语句。
clean_result (bool，可选，默认为 True) — 通过将所有位于同一设备上的子模块归为一组，来清理生成的 device_map。
offload_buffers (bool，可选，默认为 False) — 对于卸载到 CPU 或硬盘上的层，是否同时卸载缓冲区以及参数。
fallback_allocation (bool，可选，默认为 False) — 当常规分配失败时，尝试使用 BFS 分配一个符合大小限制的模块。

为给定模型计算设备映射，优先分配给 GPU，然后卸载到 CPU，最后卸载到磁盘，

我们不会超出任何 GPU 的可用内存。
如果需要将层卸载到 CPU，那么 GPU 0 上始终会预留足够的空间，以便将 CPU 上占用空间最大的层重新加载回 GPU。
如果需要将数据卸载到 CPU，我们不会超出 CPU 上可用的 RAM。
如果需要将层卸载到磁盘，CPU 上始终会保留足够的空间，以便将已卸载到磁盘的层中体积最大的那一层重新加载回 CPU。

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

accelerate.load_checkpoint_in_model

< source >

参数

model (torch.nn.Module) — 要加载检查点的模型。
checkpoint (str 或 os.PathLike) — 要加载的文件夹检查点。它可以是：
- 指向包含完整模型状态字典的文件的路径
- 指向包含分片检查点索引的 .json 文件的路径
- 指向一个文件夹的路径，该文件夹包含一个唯一的 .index.json 文件和检查点的分片。
- 指向包含唯一的 pytorch_model.bin 或 model.safetensors 文件的文件夹的路径。
device_map (Dict[str, Union[int, str, torch.device]], 可选) — 一个映射，用于指定每个子模块应放置的位置。它无需细化到每个参数/缓冲区名称，只要某个模块名称包含在内，其所有子模块都会被发送到同一设备。
offload_folder (str 或 os.PathLike，可选) — 如果 device_map 包含任何值为 "disk" 的项，则指定用于卸载权重的文件夹。
dtype (str 或 torch.dtype，可选) — 如果提供，加载时权重将被转换为该类型。
offload_state_dict (bool，可选，默认为 False) — 如果为 True，则会将 CPU 状态字典临时卸载到硬盘上，以避免当 CPU 状态字典的权重加上最大分片无法容纳时耗尽 CPU 内存。
offload_buffers (bool，可选，默认为 False) — 是否将缓冲区包含在卸载到磁盘的权重中。
keep_in_fp32_modules(List[str], 可选) — 一个模块列表，这些模块将保持为 torch.float32 数据类型。
offload_8bit_bnb (bool, 可选) — 是否启用将 8 位模块卸载到 cpu/disk。
strict (bool，可选，默认为 False) — 是否严格强制检查点 state_dict 中的键与模型的 state_dict 的键相匹配。
full_state_dict (bool，可选，默认为 True) — 如果设置为 True，则加载的 state_dict 中的所有张量都会被收集。加载的 state_dict 中不会包含 ShardedTensor 和 DTensor。
broadcast_from_rank0 (False，可选，默认为 False) — 当该选项为 True 时，必须初始化一个分布式 ProcessGroup。rank0 应接收完整的 state_dict，并将 state_dict 中的张量逐个广播到其他 rank。其他 rank 将接收这些张量，并根据模型中的本地分片进行分片（如适用）。

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

Once loaded across devices, you still need to call dispatch_model() on your model to make it able to run. To group the checkpoint loading and dispatch in one single call, use load_checkpoint_and_dispatch().

accelerate.utils.load_offloaded_weights

< source >

( modelindexoffload_folder )

参数

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

从 offload 文件夹加载权重到模型中。

accelerate.utils.load_state_dict

< source >

( checkpoint_filedevice_map = None )

参数

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

从指定文件加载检查点。如果检查点采用 safetensors 格式并传入了设备映射，则权重可以直接快速加载到 GPU 上。

accelerate.utils.offload_state_dict

< source >

( save_dir: str | os.PathLikestate_dict: dict )

参数

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

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

accelerate.utils.retie_parameters

< source >

( modeltied_params )

参数

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

如果链接已断开（例如在添加钩子时），则重新绑定给定模型中的绑定参数。

accelerate.utils.set_module_tensor_to_device

< source >

参数

module (torch.nn.Module) — 我们要移动的张量所在的模块。
tensor_name (str) — 参数/缓冲区的完整名称。
device (int、str 或 torch.device) — 用于设置张量的设备。
value (torch.Tensor, 可选) — 张量的值（在从 meta 设备迁移到其他设备时很有用）。
dtype (torch.dtype，可选) — 如果传入该值，参数的值将被转换为该 dtype。否则，value 将被转换为模型中现有参数的 dtype。
fp16_statistics (torch.HalfTensor, 可选) — 用于 8 位模型序列化的模块 fp16 统计信息列表。
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

< source >

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

参数

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

返回值

list[torch.nn.Module]

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

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

并行

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

accelerate.utils.extract_model_from_parallel

< source >

( modelkeep_fp32_wrapper: bool = Truekeep_torch_compile: bool = Truerecursive: 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

< source >

( objfsave_on_each_node: bool = Falsesafe_serialization: bool = False )

参数

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

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

accelerate.utils.load

< source >

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

< source >

( )

在脚本中引入一个阻塞点，确保所有进程都到达此点后再继续执行。

请确保所有进程都会执行到此指令，否则其中一个进程将永久挂起。

随机

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

accelerate.utils.set_seed

< source >

( seed: intdevice_specific: bool = Falsedeterministic: bool = False )

参数

seed (int) — 要设置的种子。
device_specific (bool，可选，默认为 False) — 是否根据 self.process_index 在每个设备上对种子进行略微不同的设置。
deterministic (bool，可选，默认为 False) — 是否在有可用确定性算法时使用它们。这可能会降低训练速度。

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

accelerate.utils.synchronize_rng_state

< source >

( rng_type: accelerate.utils.dataclasses.RNGType | None = Nonegenerator: torch.Generator | None = None )

accelerate.synchronize_rng_states

< source >

( rng_types: listgenerator: torch.Generator | None = None )

PyTorch XLA

其中包括在使用 PyTorch 与 XLA 时的一些实用工具。

accelerate.utils.install_xla

< source >

( upgrade: bool = False )

参数

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

一个辅助函数，用于在 Google Colaboratory 中根据 torch 版本安装相应的 xla 轮子。

示例：

>>> from accelerate.utils import install_xla

>>> install_xla(upgrade=True)

加载模型权重

其中包括一些用于加载检查点的实用工具。

accelerate.load_checkpoint_in_model

< source >

参数

model (torch.nn.Module) — 要加载检查点的模型。
checkpoint (str 或 os.PathLike) — 要加载的文件夹检查点。它可以是：
- 指向包含完整模型状态字典的文件的路径
- 指向包含分片检查点索引的 .json 文件的路径
- 指向一个文件夹的路径，该文件夹包含一个唯一的 .index.json 文件和检查点的分片。
- 指向包含唯一的 pytorch_model.bin 或 model.safetensors 文件的文件夹的路径。
device_map (Dict[str, Union[int, str, torch.device]], 可选) — 一个映射，用于指定每个子模块应放置的位置。它无需细化到每个参数/缓冲区名称，只要某个模块名称包含在内，其所有子模块都会被发送到同一设备。
offload_folder (str 或 os.PathLike，可选) — 如果 device_map 包含任何值为 "disk" 的项，则指定用于卸载权重的文件夹。
dtype (str 或 torch.dtype，可选) — 如果提供，加载时权重将被转换为该类型。
offload_state_dict (bool，可选，默认为 False) — 如果为 True，则会将 CPU 状态字典临时卸载到硬盘上，以避免当 CPU 状态字典的权重加上最大分片无法容纳时耗尽 CPU 内存。
offload_buffers (bool，可选，默认为 False) — 是否将缓冲区包含在卸载到磁盘的权重中。
keep_in_fp32_modules(List[str], 可选) — 一个模块列表，这些模块将保持为 torch.float32 数据类型。
offload_8bit_bnb (bool, 可选) — 是否启用将 8 位模块卸载到 cpu/disk。
strict (bool，可选，默认为 False) — 是否严格强制检查点 state_dict 中的键与模型的 state_dict 的键相匹配。
full_state_dict (bool，可选，默认为 True) — 如果设置为 True，则加载的 state_dict 中的所有张量都会被收集。加载的 state_dict 中不会包含 ShardedTensor 和 DTensor。
broadcast_from_rank0 (False，可选，默认为 False) — 当该选项为 True 时，必须初始化一个分布式 ProcessGroup。rank0 应接收完整的 state_dict，并将 state_dict 中的张量逐个广播到其他 rank。其他 rank 将接收这些张量，并根据模型中的本地分片进行分片（如适用）。

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

Once loaded across devices, you still need to call dispatch_model() on your model to make it able to run. To group the checkpoint loading and dispatch in one single call, use load_checkpoint_and_dispatch().

量化

其中包括用于量化模型的实用工具。

accelerate.utils.load_and_quantize_model

< source >

参数

model (torch.nn.Module) — 输入模型。该模型可以已经加载，也可以位于 meta 设备上
bnb_quantization_config (BnbQuantizationConfig) — bitsandbytes 量化参数
weights_location (str 或 os.PathLike) — 要加载的 weights_location 文件夹。它可以是：
- 指向包含完整模型状态字典的文件的路径
- 指向包含分片检查点索引的 .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 或 os.PathLike，可选) — 如果 device_map 包含任何值为 "disk" 的项，则指定用于卸载权重的文件夹。
offload_state_dict (bool，可选，默认为 False) — 如果为 True，则会将 CPU 状态字典临时卸载到硬盘上，以避免当 CPU 状态字典的权重加上最大分片无法容纳时耗尽 CPU 内存。

返回值

torch.nn.Module

量化后的模型

此函数将使用通过 bnb_quantization_config 传入的关联配置对输入模型进行量化。如果模型位于 meta 设备上，我们将根据传入的 device_map 加载并分发权重。如果模型已经加载，我们将量化该模型并将模型放到 GPU 上，

Update on GitHub