torchtitan 训练指标采集与可视化指南MFU/内存/计时指标 TensorBoard 与 WB 日志接入【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitantorchtitan 内置了一整套自动化训练指标采集与导出系统训练过程中会自动统计 MFU模型浮点利用率、平均/最大损失、每设备每秒 token 数等系统级指标同时监控显存峰值与 OOM 次数并量化数据加载耗时瓶颈。本文以 docs/metrics.md 为主线结合 torchtitan/observability/metrics.py 与 torchtitan/trainer.py 的源码实现完整讲解指标采集的配置项、指标含义与底层计算方式并给出 TensorBoard 与 Weights BiasesWB两种可视化的端到端接入步骤。读完本文你将能在一台远程训练服务器上快速开启指标记录并通过本地浏览器或 WB 云端面板实时观测训练进度。一、指标采集总览三类自动化监控指标torchtitan 的MetricsProcessor位于 torchtitan/observability/metrics.py负责收集、聚合并导出训练指标。其设计目标是开箱即用无需在训练循环中手动埋点训练器会在每个合适的时间点自动调用它见 torchtitan/trainer.py 中的指标收集与上报逻辑。采集的指标可以分为三类系统级指标High level system metricsMFU、平均损失global_avg_loss、最大损失global_max_loss、每设备每秒 token 数tokens per second, tps、TFLOPS、梯度范数grad_norm等用于评估整体训练效率与收敛质量。内存指标Memory metrics显存峰值占用max active / max reserved同时给出 GiB 与百分比、显存分配重试次数num_alloc_retries与 OOM 次数num_ooms用于发现显存瓶颈。计时指标Timing metrics端到端每步耗时end_to_end、数据加载平均耗时data_loading及其占总时长百分比用于定位数据加载是否成为训练瓶颈。这些指标会同时输出到 STDOUT控制台彩色日志与可视化后端。控制台输出格式如下摘自 metrics.py 的实现step: 10 loss: 1.23456 grad_norm: 0.1234 memory: 42.15GiB(78.32%) tps: 12,345 tflops: 456.78 mfu: 42.10%若启用了量化FP8/MX 等MFU 会显示为N/A——因为 MFU 基于 BF16 峰值浮点算力计算在量化场景下会失真源码注释明确说明此点见 metrics.py。二、指标配置项MetricsProcessor.Config 全参数说明MetricsProcessor.Config是一个Configurable配置类详见 metrics.py可以通过配置文件、config_registry 函数或 CLI 命令行参数进行设置。全部配置项如下配置项类型默认值含义log_freqint10每隔多少个 iteration 向 TensorBoard / WB 写一次指标必须大于 0enable_tensorboardboolFalse是否将指标写入 TensorBoardenable_wandbboolFalse是否将指标写入 Weights Biasessave_tb_folderstrtb存放 TensorBoard / WB 状态文件的目录名位于 dump_folder 之下save_for_all_ranksboolFalse是否让所有 rank 都保存指标为False时仅由指定的单个 rank 保存disable_color_printingboolFalse是否禁用控制台彩色日志输出几个关键配置的源码级说明log_freq校验__post_init__中会强制校验log_freq 0否则抛出ValueError(metrics.log_freq must be greater than 0.)。该行为有单元测试覆盖见 test_metrics_config.py。save_for_all_ranks与流水线并行当该选项为False且pipeline_parallel_degree 1时指标组件只会使用最后一个流水线 stage 的 0 号 rank 保存指标——因为只有该 stage 才计算损失指标见 metrics.py 的文档注释。若设为True则每个 rank 的日志会写入各自独立的子目录目录名带rank_{rank}后缀。save_tb_folder与日志目录结构日志实际输出路径为dump_folder/save_tb_folder/YYYYmmdd-HHMM/。若启用了容错fault tolerance会在其后追加replica_{ft_replica_id}/子目录若save_for_all_ranksTrue再追加rank_{rank}/子目录。该路径拼接逻辑见 metrics.py。在配置文件中启用 TensorBoard 的最小写法如下metrics: enable_tensorboard: true log_freq: 50三、TensorBoard 可视化从远程服务器到本地浏览器TensorBoard 的接入是本文档的核心实操内容共分四个步骤完整命令见 docs/metrics.md。步骤 1在训练配置中开启 TensorBoard确保训练时metrics.enable_tensorboard被设为true可以通过 config_registry 函数设置也可以通过 CLI 传入。仓库中 Llama 3 的配置示例即采用此方式见 torchtitan/models/llama3/config_registry.pymetricsMetricsProcessor.Config( enable_tensorboardTrue, ),通过 CLI 开启的方式以run_train.sh为入口./run_train.sh --config llama3_8b metrics.enable_tensorboardtrue步骤 2建立 SSH 隧道在本地命令行中执行以下命令将远程服务器的 6006 端口转发到本地ssh -L 6006:127.0.0.1:6006 [username][hostname]步骤 3在远程服务器上启动 TensorBoard 后端在 SSH 隧道已登录的远程会话中进入 torchtitan 仓库目录启动 TensorBoard日志目录指向默认的./outputs/tbtensorboard --logdir./outputs/tb结合上一节源码可知实际的 event 文件会写入./outputs/tb/YYYYmmdd-HHMM/下因此指定--logdir./outputs/tb即可被 TensorBoard 自动递归发现。步骤 4本地浏览器访问在本地浏览器打开 SSH 隧道输出的 URL或直接访问http://localhost:6006/即可看到实时的损失曲线、MFU、内存占用等指标。补充流水线并行下保证损失可见使用流水线并行pipeline parallel训练时损失只在最后一个流水线 stage 上可见。MetricsProcessor内部通过_get_metrics_rank计算损失可见 rank即最后一个 stage 的首个 rank并配合ensure_pp_loss_visible检查该 rank 是否包含在LOG_RANK环境变量中若缺失会给出告警见 metrics.py。LOG_RANK在 run_train.sh 中默认设为0例如多机训练时可显式指定LOG_RANK0,16 ./run_train.sh来保证各阶段关键 rank 的日志可见。四、Weights BiasesWB可视化登录即可自动上报WB 是另一种零配置的可视化方案只要在环境中执行过wandb login完成登录WB 就会自动把指标发送到远程服务器见 docs/metrics.md。你需要做的仅仅是开启配置项./run_train.sh --config llama3_8b metrics.enable_wandbtrue或在 config_registry 中设置metricsMetricsProcessor.Config( enable_wandbTrue, ),从源码看WandBLogger还支持通过一系列环境变量定制 run 信息见 metrics.py环境变量作用默认值WANDB_TEAM所属团队/实体entity无WANDB_PROJECT项目名torchtitanWANDB_RUN_NAME本次运行的名称无WANDB_RUN_ID运行 ID可恢复断点无WANDB_RUN_NOTES运行备注无WANDB_RUN_TAGS运行标签无WANDB_RUN_GROUP运行分组无WANDB_RUN_JOB_TYPE任务类型无WANDB_RESUME_FROM/WANDB_FORK_FROM从既有 run 续跑/派生无需要说明的是wandb是延迟导入的import 语句位于__init__内部若未安装且开启了 WB_build_metric_logger会捕获异常并提示pip install wandb不会导致训练崩溃见 metrics.py。WB 的 run 配置还会带上完整 job configconfig_dict方便在面板上回看超参。另外TensorBoard 与 WB 可以同时启用——LoggerContainer会把所有已启用的 logger 聚合起来每次log()都会依次写入每个后端见 metrics.py。五、指标字典详解每个键的含义与计算方式MetricsProcessor.log()在每次上报时组装一个指标字典并交给后端见 metrics.py。训练指标的完整键列表如下指标键含义计算方式/来源loss_metrics/global_avg_loss全局平均损失global_loss_sum / global_valid_tokens对所有 rank 求和loss_metrics/global_max_loss各 rank 本地平均损失的最大值max(local_loss_sum / local_valid_tokens)grad_norm梯度裁剪后的梯度范数训练循环直接传入throughput(tps)每设备每秒处理的 token 数ntokens_since_last_log / (time_delta * non_data_parallel_size)tflops每设备浮点吞吐num_flops_per_token * tps / 1e12mfu(%)模型浮点利用率相对峰值100 * num_flops_per_token * tps / gpu_peak_flops量化训练时为N/Atime_metrics/end_to_end(s)平均每步端到端耗时time_delta / (step - step_last_log)time_metrics/data_loading(s)平均数据加载耗时采样区间内data_loading_times的平均值time_metrics/data_loading(%)数据加载占总时长百分比100 * sum(data_loading_times) / time_deltamemory/max_active(GiB)峰值激活显存GiBPyTorchactive_bytes.all.peak统计memory/max_active(%)峰值激活显存占卡容量百分比上述值 / 设备总显存memory/max_reserved(GiB)峰值保留显存GiBPyTorchreserved_bytes.all.peak统计memory/max_reserved(%)峰值保留显存占卡容量百分比上述值 / 设备总显存memory/num_alloc_retries显存分配重试次数PyTorchnum_alloc_retries统计memory/num_oomsOOM 抛出次数PyTorchnum_ooms统计此外还会合并以下额外指标见 trainer.pyn_tokens_seen全局已见 token 总数学习率调度器指标lr_metrics如lr、effective_lr等辅助损失指标collect_aux_loss_metrics如 MoE 的 auxiliary loss 相关统计。验证阶段则通过log_validation上报validation_metrics/loss、validation_metrics/throughput(tps)及对应内存指标见 metrics.py便于对比训练与验证曲线。内存指标的底层实现依赖DeviceMemoryMonitor它在构造时通过torch.cuda.reset_peak_memory_stats()与empty_cache()重置峰值统计并在每次log()后再次重置从而保证每个采样窗口内的数值都是本窗口峰值见 metrics.py。六、指标的上报节奏与 rank 选择MetricsProcessor.should_log()决定了何时上报第 1 步step 1以及step % log_freq 0的步都会触发日志见 metrics.py确保训练一开始就有基准数据。在多卡训练下默认save_for_all_ranksFalse只有一个rank 实际写 TensorBoard/WB避免重复数据与写放大。选择规则如下见_get_metrics_rank非流水线并行配置由 rank 0 负责流水线并行且使用ZBVZeroBubble调度由 rank 0 负责该调度在 rank 0 上返回 loss其他流水线并行调度由最后一个 pipeline stage 的首个 rank 负责即(world_size // pp_size) * (pp_size - 1)。这个设计保证了记录到面板上的 loss 始终是能看到损失的那张卡上报的。七、小结torchtitan 的指标体系将训练观测沉淀为开箱即用的标准能力MetricsProcessor.Config一行配置即可开启 TensorBoard 或 WBlog_freq控制上报频率save_tb_folder决定本地落盘位置其底层自动完成 loss 全局聚合、tps/TFLOPS/MFU 计算、显存峰值与 OOM 统计、数据加载耗时占比分析并通过单 rank 上报机制避免多卡日志重复。对于使用流水线并行的训练任务注意将最后一个 stage 的首个 rank 加入LOG_RANK以保证损失曲线在控制台与面板上完整可见。【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考