面向 GPU 服务器的 SageMaker 无痛使用指南（二）—SageMaker TrainingJob

本文是面向 GPU 服务器的 SageMaker 系列博客的一部分，Amazon SageMaker 是亚马逊云提供的一个全面的机器学习平台。它支持从构建、训练和部署机器学习模型到监控和自动模型调优的整个机器学习工作流程。

利用 SageMaker 的 GPU 服务器特性，可以帮助 AIML 机器学习任务的用户像使用 EC2 实例一样快速启动，登陆各种 GPU 算力机资源（如 G5 A10，P4 A100，P5 H100…etc）并进行 LLM/SD 等模型的分布式训练和推理部署，这其中包括 SageMaker Notebook 笔记本实例，Hyperpod 集群等方式，详见附录的文章链接。

本文中我们将介绍如何像使用 EC2 GPU 服务器一样，在 SageMaker training Job 中进行 GenAI 模型的微调训练开发及调试，及训练完成后使用 SageMaker Endpoint 推理实例进行模型部署和即时推理。

SageMaker Training Job 训练实例简介

SageMaker TrainingJob 是一种使用 Amazon SageMaker 来运行机器学习模型训练作业的方式，可以方便地提供模型训练所需的数据位置、算法源代码、计算资源配置等信息。SageMaker 会在指定的 GPU 计算实例上启动和运行训练作业，如 Llama 2 模型，或者 Stable Diffusion 模型的 fine-tune。

在 SageMaker Training Job 训练作业支持各种分布式训练的框架，客户可以选择使用业界流行的和自己熟悉的算法和框架，比如 LoRA，DeepSpeed 等，从而使用数据并行或模型并行，加快大规模数据集上的训练过程。训练作业完成后,您可以将模型保存在 Amazon S3 上，之后用于模型评估、推理部署等后续步骤。

Training Job 训练实例远程连接及调试

在生产环境中，GenAI 模型的训练任务通常会持续数天甚至数周，当遇到故障或者性能瓶颈的时候，需要登陆训练实例服务器进行 trouble shooting & perf turning。

在 Training Job 训练任务拉起的 GPU 服务器，SageMaker 会启动一个 conda 环境的 docker 镜像容器，该容器会 fullfill 该服务器的 GPU 显卡资源以及存储空间，因此我们连接到 training job 主机上的容器中，即可以像登陆宿主机一样进行生产环境上的训练任务的运维和调试。

要远程登录训练任务实例的容器镜像，可以使用两种方式，一种是使用 SageMaker 提供的 ssh helper utility，一种是使用 SageMaker Training Job 的 remote debug 的功能，两种方式实现方式类似，会在训练任务的容器进程中启动 ssm agent，并等待用户客户端的连接。

以下详细介绍两种方式下训练实例的远程连接及调试的方法。

SageMaker ssh helper 远程连接

使用以下脚本在本地配置安装 ssm plugin（client）

echo "安装aws cli（v2）..."
curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
unzip awscliv2.zip
sudo ./aws/install

echo "安装linux ssm client..."
curl "https://s3.amazonaws.com/session-manager-downloads/plugin/latest/ubuntu_64bit/session-manager-plugin.deb" -o "session-manager-plugin.deb"
sudo dpkg -i session-manager-plugin.deb

ssh helper 是 Amazon 提供的一个开源 utility 工具 lib，其中包了丰富的 api 和封装功能，比如 Training Job/Inference Endpoint 等的 ssh 连接。详见附录参考资料。

训练任务的脚本中，安装 sagemaker ssh helper utility：

sudo pip3 install sagemaker-ssh-helper

使用 ssh helper 进行 training job 的远程连接，训练任务的脚本中需要稍微做一点改造，以便能拉起 SSM agent 并等待用户客户端连接，具体如下：

训练脚本 import ssh helper，并 setup session，加入下面示例中的两行代码，即可实现该功能，对客户自己的训练任务逻辑不会有其他的侵入性。

如果训练任务比较短，有可能 ssm agent 尚未拉起来，训练任务已经结束了，因此这种情况下脚本需要加入一个循环等待代码，等待 ssm agent 拉起来。代码示例如下：

%%writefile ./sd_xl_finetune_and_inference/train_and_inference.py
## setup ssm agent及建立session
import sagemaker_ssh_helper
sagemaker_ssh_helper.setup_and_start_ssh()

import time
import os
import json
import socket
start_time = time.time()

if __name__ == "__main__":
    hosts = json.loads(os.environ['SM_HOSTS'])
    current_host = os.environ['SM_CURRENT_HOST']
    ## 等待一段时间，以便ssm agent启动等待连接
    while True:
       current_time = time.time()
       if current_time - start_time >= 1200:
           break
    os.system("chmod +x ./s5cmd")
    os.system("/bin/bash train_and_inference.sh")

在 SageMaker 的 Estimator SDK 中，加 ssh helper wrapper 包装器，并且以 ssm wrapper 方式启动 fit 训练任务。代码示例如下：

ssh_wrapper = SSHEstimatorWrapper.create(estimator, connection_wait_time_seconds=600)  # <--NEW--
estimator.fit(inputs,wait=False)
instance_ids = ssh_wrapper.get_instance_ids(timeout_in_sec=900)

如上拉起带 ssh helper 的 training 训练任务后，即可以通过 wrapper 获取到远程训练任务的 instance_id，然后通过上文中安装的 SSM plugin（client）中即可通过 cli 连接到训练实例镜像，如下所示：
```
aws ssm start-session --target mi-0324aa48ddb90134a
```

SageMaker Training Job Remote Debug 远程连接

与 ssh helper 类似，SageMaker Training Job 使用 remote debug 配置也可以远程连接到训练实例 GPU 服务器，区别在于不需要安装 ssh helper ，也不需要 estimator 中使用 wrapper 包装器。具体步骤如下。

estimator fit 启动训练任务时候，增加 remote debug 配置项：
在拉起训练任务后，日志显示 download image 完成，即已经拉起训练实例 instance，这个时候即可通过 ssm client 进行连接：
与 ssh helper 方式类似，remote debug 同样通过 ssm client 进行连接，如下所示：

Training Job IDE 远程 debug

当发现生产环境中的训练过程中的错误或异常时，有时应用的日志不够详细，无法帮助您诊断和解决问题，可能需要在本地 IDE 代码中添加更多日志语句或断点来获取更多信息，或者有时可能需要检查模型在训练过程中的中间状态，例如权重、损失值或其他指标，以了解训练过程的进展情况，因此在本地 IDE（如 vscode）中进行训练任务脚本的远程 debug 也是业界的一个刚需。

同样的，我们可以使用 ssh helper 提供的功能连接并初始化 ssh 隧道，以便本地 IDE 可以连接到训练实例容器进程进行调试，具体步骤如下：

使用 ssh helper 的 local ssh training 工具自动建立 ssh 隧道转发 local ssh training 简单地接收 estimator 拉起的训练任务的 job 名即可建立 ssh 隧道，如下所示：
```
sm-local-ssh-training connect sd-xl-dreambooth-finetuning-high-2024-03-18-11-01-03-475
```

ssh helper 会自动初始化 public key，并写入本地的 .ssh 密钥目录，密钥文件名默认为 sagemaker-ssh-gw.pub（公钥）/sagemaker-ssh-gw（私钥）。在本地打开转发端口（默认端口 11022），如下所示：

tangqy@6c7e67c16c37 .ssh % ls -lat
total 40
drwxr-x---+ 63 tangqy staff 2016 4 29 10:46 ..
drwx------ 7 tangqy staff 224 4 11 09:01 .
-rw-------@ 1 tangqy staff 371 4 7 12:16 config
-rw-r—r-- 1 tangqy staff 181 3 18 20:07 sagemaker-ssh-gw.pub
-rw------- 1 tangqy staff 513 3 18 20:07 sagemaker-ssh-gw

在 vscode 配置远程连接，端口设置为 ssh helper 设定的本地转发端口 11022。

Host sagemaker-training
 HostName localhost
 IdentityFile ~/.ssh/sagemaker-ssh-gw
 Port 11022
 User root

在 vscode 中点击 ssh 连接，选择刚才配置的 sagemaker-training connection，即可连接到训练实例容器。
与 IDE 远程调试 EC2 一样，可以使用 python intepretor 远程调试 training job 中的脚本，默认的训练代码路径在/opt/ml/code 下，如下图所示：

Inference Endpoint 推理实例简介

SageMaker 推理实例是一种允许您部署训练有素的机器学习模型以用于生产环境的托管服务。在部署模型时，SageMaker 推理实例可以配置不同的 GPU 服务器的自动缩放组、使用加密、配置云监控等多种选项。推理实例会自动负载均衡传入的请求，并提供一个安全的 HTTPS 端点，使客户能够实时从 Web 应用程序、移动应用程序或物联网（IoT）设备发出预测请求。您可以根据 GenAI 场景需要的资源需求、成本考虑以及响应时间的不同，选择不同类型的 GPU 实例来托管 LLM 或者 SD 模型。

Inference Endpoint 推理端点远程登录

在机器学习模型的生命周期中，模型部署和推理阶段是至关重要的一环。无论是使用 SageMaker Training Job 训练完成的模型，还是从 Hugging Face 等开源库中获取的预训练模型，最终都需要在 GPU 机器上进行部署，以便于进行模型推理和性能调优。

在部署模型之前，我们需要对模型进行全面的推理测试，以确保其在生产环境中可以正常工作。在生产环境中，模型可能会遇到各种意料之外的问题，例如内存泄漏、资源竞争等。以及在推理实例上进行各种性能测试和分析，包括 CPU/GPU 使用率、内存占用、延迟分布等。这些工程化及实施层面均需要登录和操作推理实例的 GPU 服务器，查看日志、监控系统指标，从而快速定位并解决问题。

与 training job 一样，SageMaker 的 Inference Endpoint 可以通过 ssh helper utility 登陆到 SageMaker endpoint 推理实例的容器镜像中。

具体操作步骤如下：

与 training job 训练任务类似，在推理代码中，我们通过 ssh helper 的 sdk 启动 ssm 服务端。

import os
import sys

## for debug only
import os
import sagemaker_ssh_helper
sys.path.append(os.path.join(os.path.dirname(__file__), "lib"))
sagemaker_ssh_helper.setup_and_start_ssh()

Model 创建的时候，指定 ssh helper 的 wrapper。

Deploy model，与 training job 类似，使用 SSHModelWrapper 包装常规的 SageMaker endpoint 的 Model。代码示例如下：

from sagemaker_ssh_helper.wrapper import SSHModelWrapper
instance_type = "ml.g5.xlarge"
endpoint_name = sagemaker.utils.name_from_base("facefusion-byoc")

model = Model(image_uri=full_image_uri, 
              model_data=model_data, 
              role=role,dependencies=[SSHModelWrapper.dependency_dir()] )
ssh_wrapper = SSHModelWrapper.create(model, connection_wait_time_seconds=0) 

predictor = model.deploy(
    initial_instance_count=1,
    instance_type=instance_type,
    endpoint_name=endpoint_name,
    wait=True
)

SageMaker Inference endpoint 部署成功后，终端节点的 instance id 同样可以通过 sshModelWrapper 包装器 API 获得，从而可通过 ssm 客户端登陆到 endpoint 推理终端节点。

总结

本文介绍了如何在 SageMaker 中使用 GPU 服务器进行 GenAI 模型的微调训练和推理部署。

通过 SageMaker Training Job 和 Inference Endpoint，可以像使用 EC2 实例一样快速启动 GPU 算力资源，并远程连接到训练实例和推理实例进行调试和性能分析。使用包括使用 SageMaker ssh helper 在内的工具，我们可以方便地进行 ssh 隧道和远程 ssm agent 的建立等，从而在本地终端，IDE 中进行训练任务远程调试及推理端点部署后的登陆运维。