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	给定一个量化器对象，返回指定量化类型的字符串。

分布式

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

init_distributed	初始化进程组，这是 torch.distributed 所必需的。
is_distributed	检查是否已设置初始化 torch.distributed 所需的所有环境变量，并确认 distributed 已正确安装。
get_world_size_and_rank	获取默认进程组中当前进程的当前世界大小（即总秩数）和秩编号的函数。
gather_cpu_state_dict	在 CPU 上将分片状态字典转换为完整状态字典，仅在 rank0 返回非空结果以避免 CPU 内存峰值。

内存管理

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

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

调度器

用于在训练过程中控制学习率的工具。

get_cosine_schedule_with_warmup	创建一个学习率调度器，线性地将学习率从 0.0 增加到 lr，在 num_warmup_steps 个周期内完成，然后在剩余的 num_training_steps-num_warmup_steps 个周期内以余弦曲线下降到 0.0（假设 num_cycles = 0.5）。
get_lr	Full_finetune_distributed 和 full_finetune_single_device 假设所有优化器具有相同的学习率（LR），此处用于验证所有学习率是否一致，若一致则返回 True。

指标日志记录

各种日志工具。

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	设置常用库中伪随机数生成器种子的函数。