torchtune.utils

Checkpointing

torchtune 提供了检查点程序，允许在训练和与其他生态系统组件互操作之间无缝转换检查点格式。有关检查点的全面概述，请参阅 检查点深度解析。

FullModelHFCheckpointer	在HF格式中读取和写入检查点的检查点程序。
FullModelMetaCheckpointer	检查点保存器，用于以 Meta 的格式读取和写入检查点文件。
FullModelTorchTuneCheckpointer	Checkpointer 用于读取和写入与 torchtune 兼容格式的检查点。
ModelType	ModelType 由检查点保存器用于区分不同的模型架构。

分布式

用于启用分布式训练并支持其工作的工具。

FSDPPolicyType	一种可用作 FSDP 包装策略的函数数据类型。
init_distributed	初始化进程组，这是 torch.distributed 所必需的。
is_distributed	检查是否已设置初始化 torch.distributed 所需的所有环境变量，并确认 distributed 已正确安装。
get_world_size_and_rank	获取默认进程组中当前进程的当前世界大小（即总秩数）和秩编号的函数。
get_full_finetune_fsdp_wrap_policy	根据指定的标志memory_efficient_fsdp_wrap和modules_to_wrap检索一个FSDP包装策略。
lora_fsdp_wrap_policy	使用 FSDP 包装通过 LoRA 训练模型的默认策略。

降低精度

用于在低精度环境下工作的工具。

get_dtype	获取与给定精度字符串对应的 torch.dtype。
set_default_dtype	用于设置 torch 默认数据类型的上下文管理器。
validate_expected_param_dtype	验证所有输入参数是否具有预期的数据类型。
get_quantizer_mode	给定一个量化器对象，返回指定量化类型的字符串。

内存管理

在训练期间减少内存消耗的实用工具。

set_activation_checkpointing	用于对传入模型应用激活检查点的工具。
OptimizerInBackwardWrapper	一个极简类，旨在为反向传播中运行的优化器提供检查点保存与加载功能。
create_optim_in_bwd_wrapper	为在反向传播中运行的优化器步骤创建包装器。
register_optim_in_bwd_hooks	为反向传播中运行的优化器步骤注册钩子。

性能与分析

torchtune 提供了用于分析和调试微调任务的内存和性能的工具。

get_memory_stats	计算传入设备的内存摘要。
log_memory_stats	将包含内存统计信息的字典记录到日志中。
setup_torch_profiler	设置 profile 并返回带有后设置更新的配置器配置。

指标日志记录

各种日志工具。

metric_logging.WandBLogger	用于Weights and Biases应用程序的记录器 (https://wandb.ai/)。
metric_logging.TensorBoardLogger	用于 PyTorch 实现的 TensorBoard 的日志记录器 (https://pytorch.org/docs/stable/tensorboard.html)。
metric_logging.StdoutLogger	记录到标准输出。
metric_logging.DiskLogger	记录到磁盘。

数据

用于处理数据和数据集的工具。

padded_collate	将批次中的序列填充至该批次中最长序列的长度，并将整数列表转换为张量。
padded_collate_dpo	为直接偏好优化（DPO）填充一批序列。

其他

get_logger	获取一个带有流处理器的日志记录器。
get_device	接受可选设备字符串的函数，验证其在给定机器和分布式设置下是否正确且可用，并返回一个 torch.device。
set_seed	设置常用库中伪随机数生成器种子的函数。
generate	根据提示生成模型的令牌。
torch_version_ge	检查 torch 版本是否大于或等于指定版本。
TuneRecipeArgumentParser	一个有用的工具子类，它为argparse.ArgumentParser添加了一个内置参数“config”。