SageMaker Training Compiler 最佳实践和注意事项

一般最佳实践

• 请务必使用支持的实例类型或经过测试的模型。

• 在训练脚本中使用 Hugging Face Transformers 库为 NLP 模型创建分词器时，请务必通过指定 padding='max_length' 来使用静态输入张量形状。请勿使用 padding='longest'，因为填充批处理中的最长序列可能会改变每个训练批处理的张量形状。动态输入形状可以启动对模型的重新编译，并可能会增加总训练时间。有关 Transformers 分词器的填充选项的更多信息，请参阅 Hugging Face Transformers 文档中的填充和截断。

• 测量 GPU 内存利用率，确保使用可容纳 GPU 内存的最大批处理大小。Amazon SageMaker Training Compiler 可在训练期间减少模型的内存占用，这通常可让您在 GPU 内存中容纳更大的 batch_size。使用更大的 batch_size 会提高 GPU 利用率并减少总训练时间。

  在调整批处理大小时，还必须适当地调整 learning_rate。例如，如果您将批处理大小增加了 k 倍，则需要线性调整 learning_rate（只需乘以 k）或乘以 k 的平方根。这是为了在缩短的训练时间内实现相同或相似的收敛行为。有关针对常用模型测试的 batch_size 的参考，请参阅经过测试的模型。

• 要调试编译器加速的训练作业，请在 compiler_config 参数中启用 debug 标志。这使 SageMaker AI 能够将调试日志放入 SageMaker 训练作业日志中。

  huggingface_estimator=HuggingFace(
      ...
      compiler_config=TrainingCompilerConfig(debug=True)
  )

  请注意，如果您使用编译器启用训练作业的完整调试，这可能会增加一些开销。

PyTorch 的最佳实践

• 如果您引入一个 PyTorch 模型并需要为其设置检查点，请确保使用 PyTorch/XLA 的模型保存函数来为模型设置正确的检查点。有关该函数的更多信息，请参阅 XLA 设备上的 PyTorch 文档中的 torch_xla.core.xla_model.save。

  要了解如何将修改添加到 PyTorch 脚本中，请参阅直接使用 PyTorch 的大型语言模型（没有 Hugging Face Transformers Trainer API）。

  有关使用模型保存函数的实际应用的更多信息，请参阅 PyTorch/XLA TPU 上的 Hugging Face：更快且成本更低的训练博客中的检查点写入和加载。

• 要实现分布式训练的最佳训练时间，请考虑以下几点。

  • 使用具有多个 GPU 的实例，而不是使用单 GPU 实例。例如，与 8 x ml.p3.2xlarge 实例相比，单个 ml.p3dn.24xlarge 实例的训练时间更短。

  • 使用支持 EFA 的实例，例如 ml.p3dn.24xlarge 和 ml.p4d.24xlarge。这些实例类型加快了联网速度并缩短了训练时间。

  • 调整数据集的 preprocessing_num_workers 参数，以使模型训练不会因预处理速度缓慢而延迟。

• 避免因对模型张量进行日志记录、检查点设置和分析而导致显式评估。要了解什么是显式评估，请考虑以下代码编译示例。

  a = b+c
  e = a+d

  编译器按如下方式解释代码，并减少变量 a 的内存占用：

  e = b+c+d

  现在考虑以下情况，即更改代码来为变量 a 添加 print 函数。

  a = b+c
  e = a+d
  print(a)

  编译器按如下方式对变量 a 进行显式评估。

  e = b+c+d
  a = b+c    # Explicit evaluation
  print(a)

  例如，在 PyTorch 中，避免使用 torch.tensor.items()，它可能会引入显式评估。在深度学习中，此类显式评估可能会产生开销，因为它们会中断模型编译图中的融合运算并导致重新计算张量。

  如果您仍需要在训练期间使用 SageMaker Training Compiler 定期评估模型，我们建议以较低的频率进行日志记录和检查点设置，以减少因显式评估而产生的开销。例如，每 10 个纪元而不是每个纪元记录一次。

• 图形编译在训练的前几个步骤中运行。因此，预计前几个步骤的速度将非常慢。不过，这是一次性编译成本，可以通过较长时间的训练来摊销，因为编译会加快未来步骤的速度。初始编译开销取决于模型的大小、输入张量的大小以及输入张量形状的分布。

• 编译 PyTorch 模型时出现的最典型错误之一是由于运算符和张量的设备类型错误导致的。要正确编译 PyTorch 模型，请确保使用 XLA 设备 (xm.xla_device())，而不是使用 CUDA 或组合使用 CUDA 设备和 XLA 设备。

• mark_step() 只是 XLA 的一个障碍。如果设置不正确，则会导致训练作业停滞。

• PyTorch/XLA 提供了额外的分布式训练 API。未能正确编程 API 会导致错误地收集梯度，致使训练收敛失败。