torchtune.training

Checkpointing

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

FullModelHFCheckpointer	在HF格式中读取和写入检查点的检查点程序。
FullModelMetaCheckpointer	检查点保存器，用于以 Meta 的格式读取和写入检查点文件。
FullModelTorchTuneCheckpointer	Checkpointer 用于读取和写入与 torchtune 兼容格式的检查点。
ModelType	ModelType 由检查点保存器用于区分不同的模型架构。
FormattedCheckpointFiles	这个类提供了一种更简洁的方式来表示格式为 file_{i}_of_{n_files}.pth 的文件名列表。
update_state_dict_for_classifier	验证用于分类器模型检查点加载的状态字典。

降低精度

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

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

分布式

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

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 训练模型的默认策略。

内存管理

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

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

指标日志记录

各种日志工具。

metric_logging.CometLogger	用于Comet的记录器 (https://www.comet.com/site/)。
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	记录到磁盘。

性能与分析

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

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

其他

get_unmasked_sequence_lengths	返回每个批次元素的序列长度，排除被掩码的标记。
set_seed	设置常用库中伪随机数生成器种子的函数。