不想在 AI 模型微调和推理上费时间？试试 Amazon SageMaker 的 GPU 服务器！

本文是面向 GPU 服务器的 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 系列博客的一部分，[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 是亚马逊云科技提供的一个全面的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)平台。**它支持从构建、训练和部署[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)模型到监控和自动模型调优的整个[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)工作流程。** **利用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 的 GPU 服务器特性，可以帮助人工智能和[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)任务的用户像使用 [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)实例一样快速启动，登录各种 GPU 算力机资源（如 G5 A10，P4 A100，P5 H100…etc）并进行 LLM/SD 等模型的分布式训练和推理部署**，这其中包括 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Notebook 笔记本实例，[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Hyperpod 集群等方式，详见附录的文章链接。 本文中我们将介绍**如何像使用 [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)GPU 服务器一样，在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 中进行 GenAI 模型的微调训练开发及调试，及训练完成后使用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Endpoint 推理实例进行模型部署和即时推理。** ### **Amazon SageMaker Training Job 训练实例简介** [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 是一种使用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 来运行[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)模型训练作业的方式，可以方便地提供模型训练所需的数据位置、算法源代码、计算资源配置等信息。[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 会在指定的 GPU 计算实例上启动和运行训练作业，如 Llama 2 模型，或者 Stable Diffusion 模型的 fine-tune。 **在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 训练作业支持各种分布式训练的框架，客户可以选择使用业界流行的和自己熟悉的算法和框架**，比如 LoRA，DeepSpeed 等，从而使用数据并行或模型并行，加快大规模数据集上的训练过程。训练作业完成后，您可以将模型保存在 [Amazon S3](https://aws.amazon.com/cn/s3/?trk=cndc-detail) 上，之后用于模型评估、推理部署等后续步骤。 ### **Amazon SageMaker Training Job 训练实例远程连接及调试** **在生产环境中，GenAI 模型的训练任务通常会持续数天甚至数周，当遇到故障或者性能瓶颈的时候，需要登录训练实例服务器进行 trouble shooting & perf turning。** 在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 训练任务拉起的 GPU 服务器，[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 会启动一个 conda 环境的 docker 镜像容器，该容器会 fullfill 该服务器的 GPU 显卡资源以及存储空间，因此我们连接到 Training Job 主机上的容器中，即可以像登录宿主机一样进行生产环境上的训练任务的运维和调试。 要远程登录训练任务实例的容器镜像，可以使用两种方式，一种是使用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 提供的 ssh helper utility，一种是使用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 的 remote debug 的功能，两种方式实现方式类似，会在训练任务的容器进程中启动 ssm agent，并等待用户客户端的连接。 以下详细介绍两种方式下训练实例的远程连接及调试的方法。 #### **ssh helper 远程连接** 使用以下脚本在本地配置安装 ssm plugin（client）： ```js 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 和封装功能，比如 Amazon Training Job/Inference Endpoint 等的 .ssh 连接。详见附录参考资料。 训练任务的脚本中，安装 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) ssh helper utility： ```js sudo pip3 install sagemaker-ssh-helper ``` 使用 ssh helper 进行 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 的远程连接，训练任务的脚本中需要稍微做一点改造，以便能拉起 ssm agent 并等待用户客户端连接，具体如下： * 训练脚本 import ssh helper，并 setup session，加入下面示例中的两行代码，即可实现该功能，对客户自己的训练任务逻辑不会有其他的侵入性。 * 如果训练任务比较短，有可能 ssm agent 尚未拉起来，训练任务已经结束了，因此这种情况下脚本需要加入一个循环等待代码，等待 ssm agent 拉起来。代码示例如下： ```js %%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") ``` * 在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 的 Estimator SDK 中，加 ssh helper wrapper 包装器，并且以 ssm wrapper 方式启动 fit 训练任务。代码示例如下： ```js 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 连接到训练实例镜像，如下所示： ```js aws ssm start-session --target mi-0324aa48ddb90134a ``` #### **Amazon SageMaker Training Job Remote Debug 远程连接** 与 ssh helper 类似，[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 使用 remote debug 配置也可以远程连接到训练实例 GPU 服务器，区别在于不需要安装 ssh helper ，也不需要 estimator 中使用 wrapper 包装器。具体步骤如下: 1.estimator fit 启动训练任务时候，增加 remote debug 配置项：  2.在拉起训练任务后，日志显示 download image 完成，即已经拉起训练实例 instance，这个时候即可通过 ssm client 进行连接：  3.与 ssh helper 方式类似，remote debug 同样通过 ssm client 进行连接，如下所示：  #### **Amazon SageMaker Training Job IDE 远程 debug** 当发现生产环境中的训练过程中的错误或异常时，有时应用的日志不够详细，无法帮助您诊断和解决问题，可能需要在本地 IDE 代码中添加更多日志语句或断点来获取更多信息，或者有时可能需要检查模型在训练过程中的中间状态，例如权重、损失值或其他指标，以了解训练过程的进展情况，因此在本地 IDE（如 vscode）中进行训练任务脚本的远程 debug 也是业界的一个刚需。 同样的，我们可以使用 ssh helper 提供的功能连接并初始化 .ssh 隧道，以便本地 IDE 可以连接到训练实例容器进程进行调试，具体步骤如下： 1.使用 ssh helper 的 local ssh training 工具自动建立 .ssh 隧道转发 local ssh training 简单地接收 estimator 拉起的训练任务的 job 名即可建立 .ssh 隧道，如下所示： ```js sm-local-ssh-training connect sd-xl-dreambooth-finetuning-high-2024-03-18-11-01-03-475 ``` 2.ssh helper 会自动初始化 public key，并写入本地的 .ssh 密钥目录，密钥文件名默认为 sagemaker-ssh-gw.pub（公钥）/sagemaker-ssh-gw（私钥）。在本地打开转发端口（默认端口 11022），如下所示： ```js 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 ``` 3.在 vscode 配置远程连接，端口设置为 ssh helper 设定的本地转发端口 11022： ```js Host sagemaker-training HostName localhost IdentityFile ~/.ssh/sagemaker-ssh-gw Port 11022 User root ``` 4.在 vscode 中点击 .ssh 连接，选择刚才配置的 sagemaker-training connection，即可连接到训练实例容器。 与 IDE 远程调试 [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)一样，可以使用 python intepretor 远程调试 Training Job 中的脚本，默认的训练代码路径在 /opt/ml/code 下，如下图所示：  ### **Inference Endpoint 推理实例简介** **[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 推理实例是一种允许您部署训练有素的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)模型以用于生产环境的托管服务**。在部署模型时，[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 推理实例可以配置不同的 GPU 服务器的自动缩放组、使用加密、配置云监控等多种选项。推理实例会自动负载均衡传入的请求，并提供一个安全的 HTTPS 端点，使客户能够实时从 Web 应用程序、移动应用程序或物联网（IoT）设备发出预测请求。您可以根据 GenAI 场景需要的资源需求、成本考虑以及响应时间的不同，选择不同类型的 GPU 实例来托管 LLM 或者 SD 模型。 ### **Inference Endpoint 推理端点远程登录** **在[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)模型的生命周期中，模型部署和推理阶段是至关重要的一环**。无论是使用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 训练完成的模型，还是从 Hugging Face 等开源库中获取的预训练模型，最终都需要在 GPU 机器上进行部署，以便于进行模型推理和性能调优。 在部署模型之前，我们需要对模型进行全面的推理测试，以确保其在生产环境中可以正常工作。在生产环境中，模型可能会遇到各种意料之外的问题，例如内存泄漏、资源竞争等。以及在推理实例上进行各种性能测试和分析，包括 CPU/GPU 使用率、内存占用、延迟分布等。这些工程化及实施层面均需要登录和操作推理实例的 GPU 服务器，查看日志、监控系统指标，从而快速定位并解决问题。 与 Training Job 一样，[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 的 Inference Endpoint 可以通过 ssh helper utility 登录到 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) endpoint 推理实例的容器镜像中。 具体操作步骤如下： 1.与 Training Job 训练任务类似，在推理代码中，我们通过 ssh helper 的 sdk 启动 ssm 服务端： ```js 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() ``` 2.Model 创建的时候，指定 ssh helper 的 wrapper。 3.Deploy model，与 Training Job 类似，使用 SSHModelWrapper 包装常规的 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) endpoint 的 Model。代码示例如下： ```js 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 ) ``` 4.[Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Inference endpoint 部署成功后，终端节点的 instance id 同样可以通过 SSHModelWrapper 包装器 API 获得，从而可通过 ssm 客户端登录到 endpoint 推理终端节点：  ### **总结** **本文介绍了如何在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 中使用 GPU 服务器进行 GenAI 模型的微调训练和推理部署。** **通过 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 和 Inference Endpoint，可以像使用 [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)实例一样快速启动 GPU 算力资源，并远程连接到训练实例和推理实例进行调试和性能分析**。使用包括使用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) ssh helper 在内的工具，我们可以方便地进行 .ssh 隧道和远程 ssm agent 的建立等，从而在本地终端，IDE 中进行训练任务远程调试及推理端点部署后的登录运维。 参考资料 * [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Training Job 训练任务：https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-training.html#how-it-works-creating-a-training-job?trk=cndc-detail * [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Inference Endpoint 推理端点：https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-deployment.html?trk=cndc-detail * [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) ssh helper 远程登录：https://github.com/aws-samples/sagemaker-ssh-helper?trk=cndc-detail