使用 ML Goodput Measurement 库监控有效吞吐量

ML Goodput Measurement 库 (ml-goodput-measurement) 是一个 Python 软件包，可帮助您衡量在 Cloud TPU 虚拟机上运行的机器学习训练工作负载的效率。该库提供的指标可用于衡量工作负载有效吞吐量，即用于取得有意义的训练进展的 TPU 使用时间所占的比例。相反，低效时间是指花费在非生产性活动（如启动开销、I/O 停滞和中断恢复）上的总时间所占的比例。

您可以使用 Cloud Monitoring 信息中心和 TensorBoard 实时直观呈现有效吞吐量指标，从而能够找出瓶颈、优化资源利用率，并最终降低训练成本。

如需了解详情，请参阅 ML Goodput Measurement GitHub 代码库。

有效吞吐量指标

ML Goodput Measurement 库提供以下指标，这些指标也可在 Cloud Monitoring 和 TensorBoard 中查看。此表中的指标必须以 compute.googleapis.com/workload/ 为前缀。例如，goodput_time 的完整指标名称为 compute.googleapis.com/workload/goodput_time。

  • goodput_time：有效训练总时长（以秒为单位）。这可以解释为累计有效吞吐量。

  • badput_time：以秒为单位的低效训练总时间（启动、停滞、恢复）。这可以解释为累积的坏吐量。

  • total_elapsed_time：工作负载的总运行时间（以秒为单位）。经过的时间是从应用启动时开始计算，一直到当前时间或作业完成时间。

  • interval_goodput：指定时间段（例如过去 24 小时）内的有效吞吐率。该指标提供吞吐量的滚动窗口。

  • interval_badput：指定时间段内的无效吞吐率。该指标提供了一个用于计算无效吞吐量的滚动窗口，有助于识别 I/O 操作峰值等暂时性问题。

  • disruptions：作业中断的累计次数。中断是指导致训练过程意外停止并需要重新启动的事件。例如，硬件故障和维护事件。

  • step_time_deviation：因训练步长变化而导致非生产性时间。这也称为“抖动”。该指标用于衡量对突发情况敏感的稳定性，即最近“最糟糕”步数与历史基准之间的偏差（以秒为单位）。步是训练循环的单次迭代。

  • performance：以秒为单位的估计最快稳定步频（基准）。理想步进时间是指在最佳条件下（不受瞬态噪声或抖动影响），单个训练步进可花费的最短时间。

  • max_productive_steps：成功保留的最高步数。

安装

按照以下步骤为 TPU 工作负载设置 ML Goodput Measurement 库：

  1. 启用 Cloud Logging API 和 Cloud Monitoring API。
  2. 如果要在 Google Kubernetes Engine (GKE) 上部署，请为所有节点池配置 cloud-platform 访问权限范围。

在训练主机和分析机器上安装 ml-goodput-measurement 软件包：

pip install ml-goodput-measurement

使用 Goodput 库进行监控

如需使用 ML Goodput Measurement 库，请初始化 GoodputRecorder 实例，通过使用 record_event 上下文管理器封装关键部分来对训练代码进行插桩，并使用 GoodputMonitor 实时监控指标。监控器会运行一个后台进程，定期根据记录的事件计算有效吞吐量指标，并将其上传到 Cloud Monitoring 和 TensorBoard 以进行实时分析和可视化。

初始化 Goodput 记录器

初始化 GoodputRecorder，它是机器学习有效吞吐量测量库的核心组件。

import datetime
import jax
from ml_goodput_measurement import measurement

# Define a unique logger name for this specific run
logger_name = f'goodput_{config.run_name}'

# Instantiate the recorder
goodput_recorder = measurement.GoodputRecorder(
    job_name=config.run_name,
    logger_name=logger_name,
    logging_enabled=(jax.process_index() == 0)
)

记录活动

使用 record_event 上下文管理器封装训练代码。

def train_loop(config):

  # 1. Wrap the entire Job (Start/End)
  with goodput_recorder.record_event(measurement.Event.JOB):

    # 2. Record Hardware Initialization
    with goodput_recorder.record_event(measurement.Event.ACCELERATOR_INIT):
      # ... perform device mesh setup ...
      initialize_tpu(config)

    # 3. Record Training Prep
    with goodput_recorder.record_event(measurement.Event.TRAINING_PREP):
      # ... create checkpoint managers, setup model ...
      model = training_prep(config)

    # 4. Main Training Loop
    for step in range(config.steps):

      # Record Data Loading
      with goodput_recorder.record_event(measurement.Event.DATA_LOADING):
        batch = get_next_batch()

      # Record Step Start (CRITICAL: Pass the step number!)
      with goodput_recorder.record_event(measurement.Event.STEP, step):
        output = train_step(model, batch)

      # 5. Record Custom Events (e.g., Evaluation)
      if step % eval_interval == 0:
        with goodput_recorder.record_event(measurement.Event.CUSTOM, "eval_step"):
          run_evaluation()

如需查看如何使用 MaxText 将训练代码与 Goodput 库集成，请参阅 goodput.py。

监控活动

使用