Accelerate 文档
实用工具函数和类
并获得增强的文档体验
开始使用
实用工具函数和类
以下是 🤗 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 — 通过利用代码生成的 Triton 内核,将 TorchInductor 后端与 AotAutograd 和 cudagraphs 一起使用。了解更多
- AOT_TS_NVFUSER — nvFuser 与 AotAutograd/TorchScript。了解更多
- NVPRIMS_NVFUSER — nvFuser 与 PrimTorch。了解更多
- CUDAGRAPHS — cudagraphs 与 AotAutograd。了解更多
- OFI — 使用 Torchscript optimize_for_inference。仅用于推理。了解更多
- FX2TRT — 使用 Nvidia TensorRT 进行推理优化。仅用于推理。了解更多
- ONNXRT — 在 CPU/GPU 上使用 ONNXRT 进行推理。仅用于推理。了解更多
- TENSORRT — 使用 ONNXRT 运行 TensorRT 进行推理优化。了解更多
- AOT_TORCHXLA_TRACE_ONCE — 将 Pytorch/XLA 与 TorchDynamo 优化一起使用,用于训练。了解更多
- TORCHXLA_TRACE_ONCE — 将 Pytorch/XLA 与 TorchDynamo 优化一起使用,用于推理。了解更多
- IPEX — 在 CPU 上使用 IPEX 进行推理。仅用于推理。了解更多。
- TVM — 使用 Apach TVM 进行推理优化。了解更多
- HPU_BACKEND — 使用 HPU 后端进行推理优化。
class accelerate.utils.LoggerType
< 源代码 >( value names = None module = None qualname = None type = None start = 1 )
表示受支持的实验跟踪器的类型
值
- ALL — 环境中所有可用的受支持跟踪器
- TENSORBOARD — TensorBoard 作为实验跟踪器
- WANDB — wandb 作为实验跟踪器
- COMETML — comet_ml 作为实验跟踪器
- DVCLIVE — dvclive 作为实验跟踪器
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 分布式模型并行。
Kwargs
这些是 PyTorch 生态系统中特定交互的可配置参数,Accelerate 在后台处理这些参数。
在你的 Accelerator 中使用此对象来自定义 torch.autocast
的行为方式。有关每个参数的更多信息,请参阅此 上下文管理器 的文档。
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 及更高版本中可用。
class accelerate.utils.FP8RecipeKwargs
< source >( opt_level: typing.Literal['O1', 'O2'] = None use_autocast_during_eval: bool = None margin: int = None interval: int = None fp8_format: typing.Literal['E4M3', 'HYBRID'] = 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 recipe kwargs 类之一,例如 TERecipeKwargs
或 MSAMPRecipeKwargs
。
class accelerate.GradScalerKwargs
< source >( init_scale: float = 65536.0 growth_factor: float = 2.0 backoff_factor: float = 0.5 growth_interval: int = 2000 enabled: bool = True )
在您的 Accelerator 中使用此对象来自定义混合精度行为,特别是如何创建使用的 torch.cuda.amp.GradScaler
。请参阅此 scaler 的文档,以获取有关每个参数的更多信息。
GradScaler
仅在 PyTorch 1.5.0 及更高版本中可用。
class accelerate.InitProcessGroupKwargs
< source >( backend: typing.Optional[str] = 'nccl' init_method: typing.Optional[str] = None timeout: typing.Optional[datetime.timedelta] = None )
在您的 Accelerator 中使用此对象来自定义分布式进程的初始化。请参阅此 方法 的文档,以获取有关每个参数的更多信息。
注意:如果 timeout
设置为 None
,则默认值将基于 backend
的设置方式。
实现数据类 to_kwargs()
方法的内部混入类。
返回一个字典,其中包含与此类默认值不同的属性值。
插件
这些是可以传递给 Accelerator 对象的插件。虽然它们在文档的其他地方定义,但为了方便起见,所有插件都可以在此处查看
class accelerate.DeepSpeedPlugin
< source >( 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 和 eval 数据加载器,这将决定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
< source >( prefix = '' mismatches = None config = None must_match = True **kwargs )
使用 kwargs 中的值处理 DeepSpeed 配置。
设置 HfDeepSpeedWeakref 以使用当前的 deepspeed 插件配置
class accelerate.FullyShardedDataParallelPlugin
< source >( 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.Optional[collections.abc.Iterable[torch.nn.modules.module.Module]] = 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,则应为bool
或torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload
或torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy
的实例。 - ignored_modules (
Optional[Iterable[torch.nn.Module]]
,默认为None
) — 在使用 FSDP 包装时要忽略的模块列表。 - 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
时适用。默认情况下是一个lambda
,它在模块上调用to_empty
。 - sync_module_states (
bool
,默认为False
) — 每个单独包装的 FSDP 单元是否应从 rank 0 广播模块参数,以确保它们在初始化后在所有 rank 中都相同。除非cpu_ram_efficient_loading
为True
,否则默认为False
,然后将强制启用。 - 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
时适用。
此插件用于启用完全分片数据并行。
给定 model
,根据传入的策略以及我们是否可以使用 transformer_cls_to_wrap
创建一个 auto_wrap_policy
为 FSDP 设置混合精度策略
根据 StateDictType
设置状态字典配置。
验证混合精度策略,为了在不需要时避免导入,此验证被抽象出来。
class accelerate.utils.GradientAccumulationPlugin
< source >( num_steps: int = None adjust_scheduler: bool = True sync_with_dataloader: bool = True sync_each_batch: bool = False )
用于配置梯度累积行为的插件。您只能将 gradient_accumulation_plugin
或 gradient_accumulation_steps
之一传递给 Accelerator。同时传递两者会引发错误。
class accelerate.utils.MegatronLMPlugin
< source >( 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
) — 如果同时指定了训练和评估数据加载器,这将决定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
, defaults toNone
) — 学习率衰减迭代次数。如果为 None,则默认为train_iters
。 - lr_decay_samples (
int
, defaults toNone
) — 学习率衰减样本数。如果为 None,则默认为train_samples
。 - lr_warmup_iters (
int
, defaults toNone
) — 线性预热学习率的迭代次数。 - lr_warmup_samples (
int
, defaults toNone
) — 线性预热学习率的样本数。 - lr_warmup_fraction (
float
, defaults toNone
) — 用于线性预热学习率的 lr-warmup-(iters/samples) 的比例。 - min_lr (
float
, defaults to0
) — 学习率的最小值。调度器将低于此阈值的值裁剪掉。 - consumed_samples (
List
, defaults toNone
) — 与accelerator.prepare
调用中的数据加载器顺序相同的已消耗样本数列表。 - no_wd_decay_cond (
Optional
, defaults toNone
) — 禁用权重衰减的条件。 - scale_lr_cond (
Optional
, defaults toNone
) — 缩放学习率的条件。 - lr_mult (
float
, defaults to1.0
) — 学习率乘数。 - megatron_dataset_flag (
bool
, defaults toFalse
) — 指示数据集格式是否遵循 Megatron-LM 的 Indexed/Cached/MemoryMapped 格式。 - seq_length (
int
, defaults toNone
) — 要处理的最大序列长度。 - encoder_seq_length (
int
, defaults toNone
) — 编码器要处理的最大序列长度。 - decoder_seq_length (
int
, defaults toNone
) — 解码器要处理的最大序列长度。 - tensorboard_dir (
str
, defaults toNone
) — 用于保存 tensorboard 日志的路径。 - set_all_logging_options (
bool
, defaults toFalse
) — 是否设置所有日志记录选项。 - eval_iters (
int
, defaults to100
) — 用于评估验证/测试的迭代次数。 - eval_interval (
int
, defaults to1000
) — 在验证集上运行评估的间隔。 - return_logits (
bool
, defaults toFalse
) — 是否从模型返回 logits。 - custom_train_step_class (
Optional
, defaults toNone
) — 自定义训练步骤类。 - custom_train_step_kwargs (
Optional
, defaults toNone
) — 自定义训练步骤 kwargs。 - custom_model_provider_function (
Optional
, defaults toNone
) — 自定义模型提供函数。 - custom_prepare_model_function (
Optional
, defaults toNone
) — 自定义准备模型函数。 - custom_megatron_datasets_provider_function (
Optional
, defaults toNone
) — 自定义 megatron train_valid_test 数据集提供函数。 - custom_get_batch_function (
Optional
, defaults toNone
) — 自定义获取批次函数。 - custom_loss_function (
Optional
, defaults toNone
) — 自定义损失函数。 - other_megatron_args (
Optional
, defaults toNone
) — 其他 Megatron-LM 参数。请参考 Megatron-LM。
用于 Megatron-LM 的插件,以启用张量、流水线、序列和数据并行。还可启用选择性激活重计算和优化的融合内核。
class accelerate.utils.TorchDynamoPlugin
< source >( backend: DynamoBackend = None mode: str = None fullgraph: bool = None dynamic: bool = None options: typing.Any = None disable: bool = False )
参数
- backend (
DynamoBackend
, defaults toNone
) — 有效的 Dynamo 后端。有关更多详细信息,请参阅 https://pytorch.ac.cn/docs/stable/torch.compiler.html。 - mode (
str
, defaults toNone
) — 可选值为 'default'、'reduce-overhead' 或 'max-autotune'。 - fullgraph (
bool
, 默认为None
) — 是否可以将模型拆分为多个子图。 - dynamic (
bool
, 默认为None
) — 是否使用动态形状进行追踪。 - options (
Any
, 默认为None
) — 传递给后端的选项字典。 - disable (
bool
, 默认为False
) — 将 torch.compile() 变为无操作,用于测试
此插件用于使用 PyTorch 2.0 编译模型
配置
这些是可以配置并传递到适当集成的类
class accelerate.utils.BnbQuantizationConfig
< source >( 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
< source >( 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
) — 加速器是否应跨设备拆分数据加载器产生批次。如果为True
,则在任何类型的分布式进程上使用的实际批量大小都将相同,但它必须是您正在使用的num_processes
的整数倍。如果为False
,则使用的实际批量大小将是脚本中设置的批量大小乘以进程数。 - dispatch_batches (
bool
, 默认为None
) — 如果设置为True
,则 Accelerator 准备的数据加载器仅在主进程上迭代,然后将批次拆分并广播到每个进程。对于底层数据集为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 准备的数据加载器将利用非阻塞的主机到设备传输,从而更好地重叠数据加载器通信和计算。建议准备好的数据加载器将pin_memory
设置为True
以正常工作。 - use_stateful_dataloader (
bool
, 默认为False
) — 如果设置为True
,则 Accelerator 准备的数据加载器将由 torchdata.StatefulDataLoader 支持。这需要安装支持 StatefulDataLoader 的torchdata
版本 0.8.0 或更高版本。
调用 accelerator.prepare
时,用于数据加载器相关项的配置。
class accelerate.utils.ProjectConfiguration
< source >( 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 对象的配置。
将 self.project_dir
和 self.logging_dir
设置为适当的值。
环境变量
这些是可以为不同用例启用的环境变量
ACCELERATE_DEBUG_MODE
(str
): 是否在调试模式下运行 accelerate。 更多信息请参考这里。
数据操作与运算
这些包括模仿相同 torch
操作的数据运算,但可用于分布式进程。
accelerate.utils.broadcast
< source >( tensor from_process: int = 0 )
递归地将张量广播到所有设备的张量嵌套列表/元组/字典中。
accelerate.utils.broadcast_object_list
< source >( object_list from_process: int = 0 )
将可pickle对象列表从一个进程广播到其他进程。
accelerate.utils.concatenate
< source >( data dim = 0 )
递归地拼接形状相同的张量嵌套列表/元组/字典中的张量。
accelerate.utils.convert_to_fp32
< source >( tensor )
递归地将 FP16/BF16 精度中的张量嵌套列表/元组/字典中的元素转换为 FP32。
递归地从所有设备收集张量嵌套列表/元组/字典中的张量。
accelerate.utils.gather_object
< source >( object: typing.Any )
递归地从所有设备收集对象嵌套列表/元组/字典中的对象。
accelerate.utils.get_grad_scaler
< source >( distributed_type: DistributedType = None **kwargs )
一个通用助手,它将根据环境初始化正确的 GradScaler
实现并返回它。
accelerate.utils.get_mixed_precision_context_manager
< source >( native_amp: bool = False autocast_kwargs: AutocastKwargs = None )
返回用于自动混合精度转换的上下文管理器
accelerate.utils.listify
< source >( data )
递归地在嵌套列表/元组/字典中查找张量,并将它们转换为数字列表。
accelerate.utils.pad_across_processes
< source >( tensor dim = 0 pad_index = 0 pad_first = False )
递归地填充来自所有设备的张量嵌套列表/元组/字典中的张量,使其大小相同,以便可以安全地收集它们。
accelerate.utils.recursively_apply
< source >( func data *args test_type = <function is_torch_tensor at 0x7eff010bf010> error_on_other_type = False **kwargs )
参数
- func (
callable
) — 要递归应用的函数。 - data (nested list/tuple/dictionary of
main_type
) — 要在其上应用func
的数据。 - *args — 当应用于解包的数据时,将传递给
func
的位置参数。 - main_type (
type
, optional, defaults totorch.Tensor
) — 应用func
的对象的基本类型。 - error_on_other_type (
bool
, optional, defaults toFalse
) — 如果在解包data
后,我们遇到不是main_type
类型的对象,是否返回错误。如果为False
,则该函数将保持与main_type
类型不同的对象不变。 - **kwargs (additional keyword arguments, optional) — 当应用于解包的数据时,将传递给
func
的关键字参数。
递归地对给定基本类型的数据结构(即嵌套列表/元组/字典)应用函数。
accelerate.utils.reduce
< source >( tensor reduction = 'mean' scale = 1.0 )
递归地通过给定操作的平均值来归约跨所有进程的张量嵌套列表/元组/字典中的张量列表。
accelerate.utils.send_to_device
< source >( tensor device non_blocking = False skip_keys = None )
递归地将张量嵌套列表/元组/字典中的元素发送到给定设备。
accelerate.utils.slice_tensors
< source >( data tensor_slice process_index = None num_processes = None )
递归地在张量的嵌套列表/元组/字典中进行切片。
环境检查
这些功能检查当前工作环境的状态,包括操作系统本身的信息、它能支持的内容以及是否安装了特定的依赖项。
检查是否支持 bf16,可以选择忽略 TPU
检查是否安装了 ipex。
检查 MPS 设备是否可用。要求的最低版本为 1.12。
accelerate.utils.is_torch_version
< source >( operation: str version: str )
将当前的 PyTorch 版本与给定参考版本进行比较,并使用运算符。
检查 torch_xla
是否可用。要在安装了 torch xla 的环境中训练原生 pytorch 作业,请将 USE_TORCH_XLA 设置为 false。
检查是否通过 intel_extension_for_pytorch
或通过 stock PyTorch (>=2.4) 提供了 XPU 加速,并可能检查环境中是否存在 XPU
环境操作
一个上下文管理器,它会将传递给 os.environ
的每个关键字参数添加到环境中,并在退出时删除它们。
会将 kwargs
中的值转换为字符串,并将所有键转换为大写。
一个上下文管理器,它将临时清除环境变量。
当此上下文退出时,先前的环境变量将恢复。
accelerate.commands.config.default.write_basic_config
< source >( mixed_precision = 'no' save_location: str = '/github/home/.cache/huggingface/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 )
将当前进程分配给特定的 NUMA 节点。当每个节点至少有两个 cpu 时,效果最佳。
此结果在调用之间缓存。如果要覆盖它,请使用 accelerate.utils.environment.override_numa_afifnity
。
accelerate.utils.environment.override_numa_affinity
< source >( local_process_index: int verbose: typing.Optional[bool] = None )
覆盖为当前进程设置的任何 NUMA 亲和性。这非常繁重,需要重新计算要设置的亲和性,理想情况下,您应该改用 utils.environment.set_numa_affinity
。
装饰器,用于清理由装饰类或函数设置的 accelerate 环境变量。
在某些情况下,调用某些类或函数可能会导致设置 accelerate 环境变量,并且之后不会清理。例如,当调用
TrainingArguments(fp16=True, …)
将设置以下环境变量
ACCELERATE_MIXED_PRECISION=fp16
这可能会影响后续代码,因为环境变量优先于 TrainingArguments(fp16=False)。这对于单元测试尤其重要,在单元测试中,我们希望避免各个测试之间产生副作用。使用此装饰器装饰单元测试函数或整个类,以确保在每次测试后清理环境变量。这适用于 unittest.TestCase 和普通类 (pytest);它也适用于装饰父类时。
内存
accelerate.find_executable_batch_size
< source >( function: <built-in function callable> = None starting_batch_size: int = 128 )
一个基本的装饰器,它将尝试执行 function
。如果由于与内存不足或 CUDNN 相关的异常而失败,则批处理大小将减半并传递给 function
function
必须将 batch_size
参数作为其第一个参数。
建模
这些实用程序与 PyTorch 模型交互相关
计算模型的总大小及其最大层
accelerate.utils.compute_module_sizes
< source >( 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
< source >( model keep_fp32_wrapper: bool = True keep_torch_compile: bool = True recursive: bool = False ) → torch.nn.Module
从其分布式容器中提取模型。
accelerate.utils.get_balanced_memory
< source >( 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]]
, 可选) — 如果提供,则为某些特定权重考虑的特殊数据类型(将覆盖用于所有权重的默认数据类型)。 - 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
< source >( modules: list module_sizes: dict no_split_module_classes: list ) → Tuple[int, List[str]]
实用工具函数,它将扫描命名模块列表并返回一个完整层使用的最大大小。层的定义是
- 一个没有直接子模块(只有参数和缓冲区)的模块
- 其类名称在
no_split_module_classes
列表中的模块
accelerate.infer_auto_device_map
< source >( 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]]
, optional) — 如果提供,则为某些特定权重考虑的特殊数据类型(将覆盖用作所有权重默认数据类型的 dtype)。 - verbose (
bool
, optional, defaults toFalse
) — 是否在函数构建 device_map 时提供调试语句。 - clean_result (
bool
, optional, defaults toTrue
) — 通过将所有位于同一设备上的子模块分组在一起,来清理生成的 device_map。 - offload_buffers (
bool
, optional, defaults toFalse
) — 在卸载到 CPU 或硬盘的层中,是否也卸载缓冲区以及参数。 - fallback_allocation (
bool
, optional, defaults toFalse
) — 当常规分配失败时,尝试使用 BFS 分配一个适合大小限制的模块。
为给定模型计算设备映射,优先考虑 GPU,然后卸载到 CPU,最后卸载到磁盘,以便
- 我们不会超过任何 GPU 的可用内存。
- 如果需要卸载到 CPU,GPU 0 上始终有剩余空间,可以将卸载到 CPU 上的最大尺寸层放回 GPU 0。
- 如果需要卸载到 CPU,我们不会超过 CPU 上的可用 RAM。
- 如果需要卸载到磁盘,CPU 上始终有剩余空间,可以将卸载到磁盘上的最大尺寸层放回 CPU。
所有计算都通过分析模型参数的大小和数据类型来完成。因此,模型可以在元设备上(就像在 init_empty_weights
上下文管理器中初始化一样)。
accelerate.load_checkpoint_in_model
< source >( 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 )
参数
- model (
torch.nn.Module
) — 我们想要在其中加载检查点的模型。 - checkpoint (
str
oros.PathLike
) — 要加载的文件夹检查点。它可以是:- 包含整个模型状态字典的文件的路径
- 包含分片检查点索引的
.json
文件的路径 - 包含唯一的
.index.json
文件和检查点分片的文件夹的路径。 - 包含唯一的 pytorch_model.bin 或 model.safetensors 文件的文件夹的路径。
- device_map (
Dict[str, Union[int, str, torch.device]]
, optional) — 一个映射,指定每个子模块应放置的位置。它不需要细化到每个参数/缓冲区名称,一旦给定的模块名称在其中,它的每个子模块都将被发送到同一设备。 - offload_folder (
str
oros.PathLike
, optional) — 如果device_map
包含任何值"disk"
,则为我们将要卸载权重的文件夹。 - dtype (
str
ortorch.dtype
, optional) — 如果提供,权重将在加载时转换为该类型。 - offload_state_dict (
bool
, optional, defaults toFalse
) — 如果为True
,将临时将 CPU 状态字典卸载到硬盘驱动器上,以避免在 CPU 状态字典的权重 + 最大分片不适合时耗尽 CPU RAM。 - offload_buffers (
bool
, optional, defaults toFalse
) — 是否将缓冲区包含在卸载到磁盘的权重中。 - keep_in_fp32_modules(
List[str]
, optional) — 我们保持在torch.float32
dtype 中的模块列表。 - offload_8bit_bnb (
bool
, optional) — 是否启用在 cpu/磁盘上卸载 8 位模块。 - strict (
bool
, optional, defaults toFalse
) — 是否严格强制检查点 state_dict 中的键与模型 state_dict 的键匹配。
将(可能是分片的)检查点加载到模型中,可能会在加载时将权重发送到给定设备。
在跨设备加载后,您仍然需要在模型上调用 dispatch_model() 以使其能够运行。 要在一个调用中将检查点加载和分发分组,请使用 load_checkpoint_and_dispatch()。
accelerate.utils.load_offloaded_weights
< source >( model index offload_folder )
将权重从卸载文件夹加载到模型中。
accelerate.utils.load_state_dict
< source >( checkpoint_file device_map = None )
从给定文件加载检查点。 如果检查点是 safetensors 格式并且传递了设备映射,则可以直接在 GPU 上快速加载权重。
accelerate.utils.offload_state_dict
< source >( save_dir: typing.Union[str, os.PathLike] state_dict: dict )
在给定文件夹中卸载 state dict。
accelerate.utils.retie_parameters
< source >( model tied_params )
如果链接已断开(例如在添加 hooks 时),则在给定模型中重新绑定绑定参数。
accelerate.utils.set_module_tensor_to_device
< source >( 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 )
参数
- 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
, 可选) — 要在模块上设置的 fp16 统计信息列表,用于 8 位模型序列化。 - tied_params_map (Dict[int, Dict[torch.device, torch.Tensor]], 可选, 默认为
None
) — 当前数据指针到设备字典的映射,用于已分发的绑定权重。 对于给定的执行设备,此参数对于重用设备上共享权重的第一个可用指针以用于所有其他指针非常有用,而不是复制内存。
一个辅助函数,用于将模块的给定张量(参数或缓冲区)设置在特定设备上(请注意,执行 param.to(device)
会创建一个未链接到参数的新张量,这就是我们需要此函数的原因)。
accelerate.utils.get_module_children_bottom_up
< source >( model: Module ) → list[torch.nn.Module]
以自下而上的顺序遍历模型,并按该顺序返回子模块。
Parallel
这些包括在并行工作时应使用的一般实用程序。
accelerate.utils.extract_model_from_parallel
< source >( model keep_fp32_wrapper: bool = True keep_torch_compile: bool = True recursive: bool = False ) → torch.nn.Module
从其分布式容器中提取模型。
accelerate.utils.save
< source >( obj f save_on_each_node: bool = False safe_serialization: bool = False )
将数据保存到磁盘。 代替 torch.save()
使用。
accelerate.utils.load
< source >( f map_location = None **kwargs )
torch.load()
的兼容替代品,如果 torch
版本为 2.4.0 或更高版本,则允许使用 weights_only
。 否则将忽略 kwarg。
还将为 numpy 数组添加(然后删除)异常
在脚本中引入阻塞点,确保所有进程都到达此点后再继续。
确保所有进程都将到达此指令,否则您的一个进程将永远挂起。
Random
这些实用程序与设置和同步所有随机状态有关。
accelerate.utils.set_seed
< source >( seed: int device_specific: bool = False deterministic: bool = False )
用于在 random
, numpy
, torch
中设置 seed 以实现可复现行为的辅助函数。
accelerate.utils.synchronize_rng_state
< source >( rng_type: typing.Optional[accelerate.utils.dataclasses.RNGType] = None generator: typing.Optional[torch._C.Generator] = None )
accelerate.synchronize_rng_states
< source >( rng_types: list generator: typing.Optional[torch._C.Generator] = None )
PyTorch XLA
这些工具包含在使用 PyTorch 和 XLA 时非常有用的实用程序。
accelerate.utils.install_xla
< source >( upgrade: bool = False )
辅助函数,用于根据 Google Colaboratory 中的 torch
版本安装合适的 xla wheels。
加载模型权重
这些工具包含加载检查点时非常有用的实用程序。
accelerate.load_checkpoint_in_model
< source >( 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 )
参数
- model (
torch.nn.Module
) — 我们想要在其中加载检查点的模型。 - checkpoint (
str
oros.PathLike
) — 要加载的文件夹检查点。它可以是:- 包含整个模型状态字典的文件的路径
- 包含分片检查点索引的
.json
文件的路径 - 包含唯一的
.index.json
文件和检查点分片的文件夹的路径。 - 包含唯一的 pytorch_model.bin 或 model.safetensors 文件的文件夹的路径。
- device_map (
Dict[str, Union[int, str, torch.device]]
, optional) — 一个映射,指定每个子模块应放置的位置。它不需要细化到每个参数/缓冲区名称,一旦给定的模块名称在其中,它的每个子模块都将被发送到同一设备。 - offload_folder (
str
oros.PathLike
, optional) — 如果device_map
包含任何值"disk"
,则为我们将要卸载权重的文件夹。 - dtype (
str
ortorch.dtype
, optional) — 如果提供,权重将在加载时转换为该类型。 - offload_state_dict (
bool
, optional, defaults toFalse
) — 如果为True
,将临时将 CPU 状态字典卸载到硬盘驱动器上,以避免在 CPU 状态字典的权重 + 最大分片不适合时耗尽 CPU RAM。 - offload_buffers (
bool
, optional, defaults toFalse
) — 是否将缓冲区包含在卸载到磁盘的权重中。 - keep_in_fp32_modules(
List[str]
, optional) — 我们保持在torch.float32
dtype 中的模块列表。 - offload_8bit_bnb (
bool
, optional) — 是否启用在 cpu/磁盘上卸载 8 位模块。 - strict (
bool
, optional, defaults toFalse
) — 是否严格强制检查点 state_dict 中的键与模型 state_dict 的键匹配。
将(可能是分片的)检查点加载到模型中,可能会在加载时将权重发送到给定设备。
在跨设备加载后,您仍然需要在模型上调用 dispatch_model() 以使其能够运行。 要在一个调用中将检查点加载和分发分组,请使用 load_checkpoint_and_dispatch()。
量化
这些工具包含量化模型时非常有用的实用程序。
accelerate.utils.load_and_quantize_model
< source >( 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 设备上 - bnb_quantization_config (
BnbQuantizationConfig
) — bitsandbytes 量化参数 - weights_location (
str
或os.PathLike
) — 要加载的权重位置文件夹。它可以是:- 包含整个模型状态字典的文件的路径
- 包含分片检查点索引的
.json
文件的路径 - 包含唯一的
.index.json
文件和检查点分片的文件夹的路径。 - 包含唯一的 pytorch_model.bin 文件的文件夹的路径。
- device_map (
Dict[str, Union[int, str, torch.device]]
, optional) — 指定每个子模块应放置位置的映射。不需要细化到每个参数/缓冲区名称,一旦给定的模块名称在其中,它的每个子模块都将被发送到相同的设备。 - no_split_module_classes (
List[str]
, optional) — 不应跨设备拆分的层类名称列表(例如,任何具有残差连接的层)。 - max_memory (
Dict
, optional) — 设备标识符到最大内存的字典。如果未设置,则默认为最大可用内存。 - offload_folder (
str
或os.PathLike
, optional) — 如果device_map
包含任何值"disk"
,则为我们将要卸载权重的文件夹。 - offload_state_dict (
bool
, optional, defaults toFalse
) — 如果为True
,将临时将 CPU 状态字典卸载到硬盘驱动器上,以避免在 CPU 状态字典的权重 + 最大的分片不合适时耗尽 CPU RAM。
返回
torch.nn.Module
量化模型
此函数将使用传入 bnb_quantization_config
中的相关配置量化输入模型。如果模型在 meta 设备中,我们将根据传入的 device_map
加载和分发权重。如果模型已加载,我们将量化模型并将模型放在 GPU 上,