ML Goodput Measurement

概述

ML Goodput Measurement是一个库，旨在与Cloud TPU一起使用，以记录必要的信息并查询作业的Goodput和Badput分解。它可以通过pip安装来导入其模块，并检索关于训练作业总体生产性Goodput和Badput来源的信息。该软件包公开API接口，用于从用户应用程序记录有用的信息并查询作业运行的Goodput，以深入了解ML工作负载的生产力和计算资源的利用率。

该软件包还公开了Goodput监控API，允许异步查询和将作业的Goodput导出到Tensorboard，具有可配置的上传间隔。

组件

ML Goodput Measurement库由以下主要组件组成

• GoodputRecorder

• GoodputCalculator

• GoodputMonitor

GoodputRecorder公开API供客户端导出训练作业进展过程中的关键时间戳，即允许记录生产步骤时间和总作业运行时间的API。库将序列化和存储这些数据在Google Cloud Logging中。

GoodputCalculator公开API根据记录的数据计算Goodput。Cloud Logging异步处理其内部操作。计算Goodput的推荐方法是运行一个独立于训练应用程序的分析程序，要么在CPU实例上，要么在用户的开发机器上。

GoodputMonitor公开API异步查询和上传Goodput数据到Tensorboard。它通过在内部实例化一个GoodputCaluclator来实现这一点。

安装

要安装ML Goodput Measurement软件包，请在TPU VM上运行以下命令

pip install ml-goodput-measurement

使用方法

使用此包需要设置一个已启用计费的Google Cloud项目以正确使用Google Cloud Logging。如果您没有Google Cloud项目，或者您的Google Cloud项目没有启用计费，请执行以下操作

1. 在Google Cloud控制台的项目选择页面上，选择或创建一个Google Cloud项目。

2. 确保您的Google Cloud项目已启用计费。说明可在此处找到这里

3. 启用 Cloud Logging API。

要在Cloud TPU上运行您的训练，请按照以下说明设置Cloud TPU环境这里。

要了解更多关于Google Cloud Logging的信息，请访问此页面。

导入

要使用此包，导入goodput模块

from ml_goodput_measurement import goodput
from ml_goodput_measurement import monitoring

定义Google Cloud Logging记录器的名称。

创建一个运行特定的记录器名称，Cloud Logging条目可以写入并从中读取。

例如

goodput_logger_name = f'goodput_{config.run_name}'

创建一个GoodputRecorder对象

接下来，创建具有以下参数的记录器对象

1. job_name：作业的完整运行名称。
2. logger_name：Cloud Logging记录器对象的名称（在上一步骤中创建）。
3. logging_enabled：此进程是否启用了Cloud Logging。

注意：对于多工作员设置，请确保只有一个工作员写入日志以避免重复。例如，在JAX中，检查可以是if jax.process_index() == 0

注意： logging_enabled默认为False，如果没有写入任何日志，则无法完成Goodput计算。

例如

goodput_recorder = goodput.GoodputRecorder(job_name=config.run_name, logger_name=goodput_logger_name, logging_enabled=(jax.process_index() == 0))

使用GoodputRecorder记录数据

记录作业的开始和结束时间

使用记录器对象记录作业的整体开始和结束时间。

例如

def main(argv: Sequence[str]) -> None:
# Initialize configs…
goodput_recorder.record_job_start_time(datetime.datetime.now())
# TPU Initialization and device scanning…
# Set up other things for the main training loop…
# Main training loop
train_loop(config)
goodput_recorder.record_job_end_time(datetime.datetime.now())

记录步骤时间

使用record_step_start_time(step_count)使用记录器对象记录步骤的开始时间。

例如

def train_loop(config, state=None):
# Set up mesh, model, state, checkpoint manager…

# Initialize functional train arguments and model parameters…

# Define the compilation

for step in np.arange(start_step, config.steps):
  goodput_recorder.record_step_start_time(step)
  # Training step…

return state

记录TPU初始化、训练准备和数据加载时间

• 使用record_tpu_init_start_time和record_tpu_init_end_time使用记录器对象记录TPU初始化时间。
• 使用record_training_preparation_start_time和record_training_preparation_end_time使用记录器对象记录训练准备时间。
• 使用record_data_loading_start_time和record_data_loading_end_time使用记录器对象记录数据加载时间。

例如

def train_loop(config, state=None):
goodput_recorder.record_tpu_init_start_time()
# Set up mesh, model, state, checkpoint manager…
goodput_recorder.record_tpu_init_end_time()
goodput_recorder.record_training_preparation_start_time()
# Set up training set, initialize functional train arguments and model parameters…
# Define the compilation
# Set up any metrics collectors
goodput_recorder.record_training_preparation_end_time()

for step in np.arange(start_step, config.steps):
  goodput_recorder.record_data_loading_start_time()
  example_batch = load_next_batch(data_iterator, example_batch, config)
  goodput_recorder.record_data_loading_end_time()
  goodput_recorder.record_step_start_time(step)
  # Training step…

return state

使用GoodputCalculator检索Goodput

为了检索作业运行的Goodput，您需要做的是使用作业运行的名称和用于记录该作业运行数据的Cloud Logging记录器名称实例化一个GoodputCalculator对象。然后调用get_job_goodput API以获取作业运行的计算Goodput。

建议从运行在训练机器之外的实例中对作业运行进行get_job_goodput调用。

创建一个GoodputCalculator对象

创建计算器对象

goodput_logger_name = f'goodput_{config.run_name}' # You can choose your own logger name.
goodput_calculator = goodput.GoodputCalculator(job_name=config.run_name, logger_name=goodput_logger_name)

检索Goodput

最后，调用get_job_goodput API以检索整个作业运行的Goodput。此API接受一个可选参数include_badput_breakdown，默认为False。

返回的结果是查询时作业的Goodput的元组，一个映射各种Badput来源及其相应百分比的字典以及最后记录的步骤。如果没有设置include_badput_breakdown，则返回一个空的Badput字典。

如果您只对Goodput感兴趣

total_goodput, _, _ = goodput_calculator.get_job_goodput()
print(f"Total job goodput: {total_goodput:.2f}%")

检索Badput分解

Badput分解是各种Badput来源及其对应值的字典表示。Badput是指工作在非训练任务上所花费的时间占总工作寿命的百分比。这包括进行TPU初始化、训练准备、检查点加载、编译或重新编译、数据加载、检查点保存以及因中断而损失的时间。

当前库支持以下Badput分解分类

# Supported Badput Types
class BadputType(enum.Enum):
 """The type of Badput."""
 TPU_INITIALIZATION = 1
 TRAINING_PREP = 2
 PROGRAM_STARTUP = 3
 DATA_LOADING = 4
 UNPRODUCTIVE_CHECKPOINTING = 5
 WASTED_PROGRESS_FROM_DISRUPTION = 6
 OTHER = 7

如果您想同时获取Badput分解和Goodput

goodput, badput_breakdown, last_step = goodput_calculator.get_job_goodput(include_badput_breakdown=True)
print(f"Last step recorded: {last_step}")
print(f"Goodput: {goodput:.2f}%")
print(f"Badput due to TPU initialization: {badput_breakdown[goodput.BadputType.TPU_INITIALIZATION]:.2f}%")
print(f"Badput due to training preparation: {badput_breakdown[goodput.BadputType.TRAINING_PREP]:.2f}%")
print(f"Badput due to program startup: {badput_breakdown[goodput.BadputType.PROGRAM_STARTUP]:.2f}%")
print(f"Badput due to data loading: {badput_breakdown[goodput.BadputType.DATA_LOADING]:.2f}%")
print(f"Badput due to disruption and wasted progress: {badput_breakdown[goodput.BadputType.WASTED_PROGRESS_FROM_DISRUPTION]:.2f}%")

使用GoodputMonitor监控Goodput

为了在Tensorboard上监控作业的Goodput，您只需使用作业运行名称、云日志名称和Goodput监控配置（如下所述）创建一个GoodputMonitor对象。然后调用start_goodput_uploader API以异步查询并将测量的Goodput上传到指定的Tensorboard目录。

创建一个GoodputMonitor对象

使用以下参数创建一个GoodputMonitor对象

1. job_name：作业的完整运行名称。
2. logger_name：Cloud Logging记录器对象的名称（在上一步骤中创建）。
3. tensorboard_dir：写入TensorBoard数据的目录。
4. upload_interval：查询和上传数据到TensorBoard的时间间隔。
5. monitoring_enabled：是否启用监控。如果应用程序需要监控Goodput，则应将此值设置为True。只有一个工作者应该启用监控。
6. include_badput_breakdown：是否查询和上传Badput分解数据到Tensorboard。

注意：请确保只有一个工作者启用Goodput监控。例如，在JAX中，检查可以是if jax.process_index() == 0

例如

goodput_logger_name = f'goodput_{config.run_name}' # You can choose your own logger name.
goodput_monitoring_enabled = config.monitor_goodput and jax.process_index() == 0 # Check for configs whether or not the enable monitoring.

goodput_monitor = monitoring.GoodputMonitor(
      job_name=config.run_name,
      logger_name=logger_name,
      tensorboard_dir=config.tensorboard_dir,
      upload_interval=config.goodput_upload_interval_seconds,
      monitoring_enabled=True,
      include_badput_breakdown=True,
    )

开始异步“查询和上传”Goodput

调用start_goodput_uploader API以启动一个线程，该线程持续查询和上传Goodput。

goodput_monitor.start_goodput_uploader()