Amazon SageMaker HyperPod 存储设计与实践（一）

SageMaker HyperPod 简介

Amazon SageMaker HyperPod （简称 HyperPod）消除了为训练基础模型（FM）构建和优化机器学习（ML）基础设施所涉及的，千篇一律的繁重工作，最多可将训练时间缩短 40％。SageMaker HyperPod 预配置了 SageMaker 的分布式训练库，使客户能够自动将训练工作负载拆分到数千个加速器上，因此可以并行处理工作负载以提高模型性能。SageMaker HyperPod 还通过定期保存检查点，来确保客户能够不间断地持续进行 FM 训练。当训练期间出现硬件故障时，SageMaker HyperPod 会自动检测故障，修复或更换有故障的实例，并从上次保存的检查点恢复训练，从而使客户无需手动管理此过程，并帮助他们在分布式设置中进行数周或数月的无中断训练。

背景与摘要

SageMaker HyperPod 能够帮助客户简化大模型训练的基础设设施搭建和管理，所以大量客户采用 SageMaker HyperPod 作为大模型训练基础设施。在 SageMaker HyperPod 集群部署中，我们往往参考 awsome-distributed-training 中的代码进行部署。但 awsome-distributed-training 在存储方面只给出了 FSx for Lustre 在 SageMaker HyperPod 中的部署参考，这远远不能满足客户的实际存储需求，这是因为：

1. 一般我们会有多位模型开发人员在 SageMaker HyperPod 训练平台上进行模型开发、测试和训练， 不同开发人员共享同一个 FSx for Lustre 的共享目录，这无法保证不同开发人员的代码环境安全隔离，如模型训练代码被其他人员误修改、删除等。
2. 因为用途不同，如训练数据往往需要提前预处理，一般我们希望将我们的训练代码与训练数据、检查点（checkpoint）等数据分别存在不同的存储上，进行不同的存储管理，如备份周期等。
3. 在 SageMaker HyperPod 在同一时刻，有不同训练任务在运行，有的训练任务只是进行模型代码 debug，往往只运行几个 Step，有的训练任务长达数周或数月，遍历所有训练数据集用于模型训练。 如果这些训练任务都同时使用 FSx for Lustre 上的数据集进行训练，势必会造成存储带宽的争抢。
4. 从经济性角度来说，对于数据集较小的训练任务，我们可以采用从 S3 上直接读取训练数据进行即可，如 LLM 模型中的模型微调任务，而不用部署 FSx for Lustre，即可满足要求。
5. 针对检查点（checkpoint）ing阶段，以往我们需要先把检查点（checkpoint）存放到 Instance Store（本地 NVMe 磁盘），再上传到 S3 上，这样增加了训练代码的复杂度，尤其是 sharded 检查点（checkpoint）来说。同时，在整个 Instance 级别故障中，若这时检查点（checkpoint）没有完成向 S3 的上传，那会丢失整个检查点（checkpoint）。为避免这种故障，我们希望直接将检查点（checkpoint）保存到 S3 上。

正是基于上面在 SageMaker HyperPod 实际使用的问题，我们需要进行存储设计，使用不同存储的组合来满足这些需求。本文的主要内容包含如何在 SageMaker HyperPod 环境中支持 EFS、MountPoint for S3，这些存储在不同场景下的设计以及针对不同模型训练的场景下的测试与实践。

SageMaker HyperPod 中的存储设计

一、SageMaker HyperPod 中的默认存储部署

参考 awsome-distributed-training 代码中，给出了 FSx for Lustre 的部署脚本 mount_fsx.sh。这个脚本的作用：

1. 检查 FSx for Lustre 挂载状态
2. 将 FSx for Lustre 信息加入 /etc/fstab
3. 执行挂载
4. 创建一个新的服务，确保 FSx for Lustre 一直处于挂载状态

mount_fsx.sh 脚本会被目录 1.architectures/5.SageMaker-HyperPod/LifecycleScripts/base-config/ 下的 lifecycle_script.py 脚本所调用，lifecycle_script.py 会将 provisioning_parameters.json 文件中的 FSx for Lustre 的挂载信息传递给 mount_fsx.sh 脚本，并且通过调用同目录下的 add_users.sh 脚本，将需要增加的每个用户的 home 目录放到 FSx for Lustre 的挂载点 /fsx 上。集群部署完成后，示例效果如下：

root@ip-10-1-84-37:/usr/bin# mount|grep lustre
10.1.35.32@tcp:/xi3zrbev on /fsx type lustre (rw,noatime,flock,lazystatfs)
root@ip-10-1-84-37:/usr/bin# su - ubuntu
ubuntu@ip-10-1-84-37:~$ pwd
/fsx/ubuntu

集群创建完后，集群会自动挂载 FSx for Lustre 文件系统，并且对于用户 ubuntuhome 目录来说，会是 /fsx 下的同名目录，同时训练数据也会保存在 /fsx 的其他子目录中。

二、SageMaker HyperPod 存储的精细化设计

在默认部署中，SageMaker HyperPod 只支持一种存储方式：FSx for Lustre，我们希望在部署中支持更多存储来满足不同的存储需求。针对不同的数据类型，我们希望有不同的存储，如下图。

1. 针对 Datasets 来说，原始数据是存在于 S3 之上，通过 FSx for Lustre 与 S3桶（或前缀）建立映射，FSx for Lustre 自动与 S3 进行同步，实现训练过程中的存储的加速获取。针对检查点（checkpoint）来说，利用 FSx for Lustre 高带宽的优势，先将检查点（checkpoint）存储于 FSx for Lustre 中，再自动同步回 S3 来实现持久性存储。
2. 但有的时候，我们的数据集不大，存储一般不会成为模型训练中的瓶颈。这时但我们可以采用直接访问 S3 的方式获取训练数据，节省宝贵的 FSx for Lustre 带宽资源用于大规模、长训练。对 S3 的访问我们可以采用下载方式到本地磁盘，再读取到 GPU 中，然而我们更推荐利用 Mountpoint for S3 对 S3 桶或前缀进行本地挂载，实现数据的读取。同时，类似于 FSx for Lustre，我们也将检查点（checkpoint）数据通过 Mountpoint for S3 存储至 S3 桶中。
3. 针对模型开发和调试人员来说，每个开发人员都是有自己的工作 home 目录，他们的工作成果应该是通过设定特定的权限来实现不同人员的访问安全隔离，比如可以设计只有自己才能访问自己的工作 home 目录，或者只有自己才有完整访问权限，其他人只有代码的读权限等。这就需要利用到 EFS Access Point 的功能，利用 EFS Access Point 实现每个开发人员都有自己的工作目录，并可以对每个 Access Point 设置单独权限。

SageMaker HyperPod 部署过程

在开始我们的配置之前，有必要熟悉 SageMaker HyperPod 的部署，这是一个复杂的，需要调用多个脚本的过程，在进行存储实施前我们简单介绍一下 SageMaker HyperPod 的部署过程。通过下图我们看一下不同脚本的调用关系，该图来自于官方 SageMaker HyperPod 文档。不同脚本在部署过程中的作用和关系详见官方文档，这里不再赘述。

从上图可以看到，如果我们需要在部署中增加对 EFS、Mountpoint for S3 的支持，我们首先要定义相应的脚本，再在 provisioning_params.json 中定义 EFS ID，S3 桶名等输入参数，最后由 lifecycle_script.py 调用自定义存储脚本即可。

SageMaker HyperPod 中存储实施

一、在 SageMaker HyperPod 中实现 EFS Access Point 支持

首先，为了实现各用户 home 目录的隔离，我们把 home 目录从 FSx for Lustre 挂载目录 /fsx 剥离出来，我们利用 EFS 中的 Access Point 来实现不同用户 home 目录的隔离，实现安全性需求，同时也方便以 EFS 文件系统为整体进行数据的备份与管理。我们知道 add_users.sh 脚本会读取同目录下的 shared_users.txt 来进行用户的创建，shared_users.txt 包含用户名，UID，home 目录位置，具体如下：

user1,2001,/fsx/user1
user2,2002,/fsx/user2
user3,2003,/fsx/user3
user4,2004,/fsx/user4

下面我们就介绍关键的步骤：

1. 首先，我们根据 shared_users.txt 中的用户名和 UID 创建 EFS access point，并将 user 与 access point 的对应关系保存在 share_users_with_accesspoint_id.txt 文件中，关键代码如下：

new_file = open("share_users_with_accesspoint_id.txt","a" )
with open("shared_users.txt") as file:
    for line in file:
        line = line.strip()
        username, user_group_id, home_directory = line.split(sep=",")
        response = client.create_access_point(
            Tags=[
                {
                    'Key': "name",
                    'Value': username
                },
            ],
            FileSystemId=efs_id,
            RootDirectory={
                'Path': "/"+username,
                'CreationInfo': {
                    'OwnerUid': int(user_group_id),
                    'OwnerGid': int(user_group_id),
                    'Permissions': '755'
                }
            }
        )

        new_file.write("{},{},{},{}\n".format(username, user_group_id, home_directory, response.get("AccessPointId")))
        print(response.get("AccessPointId"))

2. 安装 EFS 客户端，这里脚本参考 efs-utils，不再详述。

3. 创建 EFS 挂载脚本 mount_efs.sh，用于 EFS access point 的挂载，关键代码如下：

add_to_fstab() {
  # Add FSx to /etc/fstab
  sudo echo "$EFS_FSID.efs.us-west-2.amazonaws.com:/ $MOUNT_POINT efs accesspoint=$AccessPointId,tls,_netdev,noresvport,iam 0 0" | sudo tee -a /etc/fstab  
}

mount_fs() {
  if [[ ! -d $MOUNT_POINT ]]; then
    sudo mkdir -p $MOUNT_POINT
    sudo chmod 644 $MOUNT_POINT
  fi

  sudo mount -t efs -o noresvport,iam,tls,accesspoint=$AccessPointId $EFS_FSID.efs.$AWS_DEFAULT_REGION.amazonaws.com:/ $MOUNT_POINT
  mount | grep nfs

}

4. 在 lifecycle_script.py 文件中增加代码如下，这里使用循环保证每个 access point 被挂载到正确的目录下。

if efs_id:
        with open("share_users_with_accesspoint_id.txt") as file:
            for line in file:
                line = line.strip()
                username, user_group_id, home_directory, accesspoint_id = line.split(sep=",")
                ExecuteBashScript("./mount_efs.sh").run(efs_id, username, user_group_id, home_directory, accesspoint_id)
                

5. 同时增加一个 install_efs_remount_service.sh，开启服务，用于检查 EFS 的挂载状态，确保 EFS 是一直处于挂载状态。

install_remount_service() {
  
  if [[ ! -d /opt/ml/scripts ]]; then
    mkdir -p /opt/ml/scripts
    chmod 644 /opt/ml/scripts
    echo "Created dir /opt/ml/scripts"
  fi
  CHECK_MOUNT_FILE=/opt/ml/scripts/check_mount_efs.sh
  cat > $CHECK_MOUNT_FILE << EOF
#!/bin/bash
if ! grep -qs "127.0.0.1:/" /proc/mounts; then
  /usr/bin/mount -a
else
  systemctl stop check_efs_mount.timer
fi
EOF
  chmod +x $CHECK_MOUNT_FILE
  cat > /etc/systemd/system/check_efs_mount.service << EOF
[Unit]
Description=Check and remount efs filesystems if necessary
[Service]
ExecStart=$CHECK_MOUNT_FILE
EOF
  cat > /etc/systemd/system/check_efs_mount.timer << EOF
[Unit]
Description=Run check_efs_mount.service every minute
[Timer]
OnBootSec=1min
OnUnitActiveSec=1min
[Install]
WantedBy=timers.target
EOF

6. 使用以上脚本创建 SageMaker HyperPod 集群后，在 /home 目录下会有各个用户的 home 目录，并且不同用户对自已有完全访问权限，对其他用户 home 目录只有只读权限（当然，也可以设置更严格权限禁止其他用户访问自己目录）。

• a. 查看 Mount 状态，可以看到 4 个用户所对应的 Access Point 已经挂载到对应的 home 目录上。

root@ip-10-1-41-18:/usr/bin# mount|grep 127
127.0.0.1:/ on /home/user1 type nfs4 (rw,relatime,vers=4.1,rsize=1048576,wsize=1048576,namlen=255,hard,noresvport,proto=tcp,port=20438,timeo=600,retrans=2,sec=sys,clientaddr=127.0.0.1,local_lock=none,addr=127.0.0.1,_netdev)
127.0.0.1:/ on /home/user2 type nfs4 (rw,relatime,vers=4.1,rsize=1048576,wsize=1048576,namlen=255,hard,noresvport,proto=tcp,port=20187,timeo=600,retrans=2,sec=sys,clientaddr=127.0.0.1,local_lock=none,addr=127.0.0.1,_netdev)
127.0.0.1:/ on /home/user3 type nfs4 (rw,relatime,vers=4.1,rsize=1048576,wsize=1048576,namlen=255,hard,noresvport,proto=tcp,port=20191,timeo=600,retrans=2,sec=sys,clientaddr=127.0.0.1,local_lock=none,addr=127.0.0.1,_netdev)
127.0.0.1:/ on /home/user4 type nfs4 (rw,relatime,vers=4.1,rsize=1048576,wsize=1048576,namlen=255,hard,noresvport,proto=tcp,port=20248,timeo=600,retrans=2,sec=sys,clientaddr=127.0.0.1,local_lock=none,addr=127.0.0.1,_netdev)

• b. 通过切换用户，可以看到不同用户自动进入所在的 home 目录。

root@ip-10-1-41-18:/usr/bin# su - user1
user1@ip-10-1-41-18:~$ pwd
/home/user1
user1@ip-10-1-41-18:~$ exit
logout
root@ip-10-1-41-18:/usr/bin# su - user2
user2@ip-10-1-41-18:~$ pwd
/home/user2
user2@ip-10-1-41-18:~$

• c. 通过下面演示可以看到，每个用户都可以在自己目录进行创建、修改、删除文件，但对于其他用户的目录，只有查看权限（当然，我们可以设置权限拒绝其他用户查看自己的 home 目录的内容）。

root@ip-10-1-41-18:/usr/bin# su - user1
user1@ip-10-1-41-18:~$ pwd
/home/user1
user1@ip-10-1-41-18:~$ touch user1-created-file1
user1@ip-10-1-41-18:~$ cd ..
user1@ip-10-1-41-18:/home$ ls
ubuntu user1 user2 user3 user4
user1@ip-10-1-41-18:/home$ cd user2
user1@ip-10-1-41-18:/home/user2$ touch user1-created-file2
touch: cannot touch 'user1-created-file2': Permission denied
user1@ip-10-1-41-18:/home/user2$ exit
logout
root@ip-10-1-41-18:/usr/bin# su - user2
user2@ip-10-1-41-18:~$ cd ../user1
user2@ip-10-1-41-18:/home/user1$ ls
user1-created-file1 user1file
user2@ip-10-1-41-18:/home/user1$ rm -rf user1-created-file1
rm: cannot remove 'user1-created-file1': Permission denied
user2@ip-10-1-41-18:/home/user1$

总结

FSx for Lustre 作为默认挂载的共享存储无法实现模型开发中的安全隔离需求。通过在 SageMaker HyperPod 中实现 EFS Access Point 支持，每个用户都有自己独立的 home 目录，实现了用户工作环境的隔离和安全性。用户对自己的 home 目录拥有完全访问权限，但对其他用户的 home 目录只有读权限或无权限，从而保证了集群用户的数据安全和代码环境。此外 EFS 文件系统还方便进行整体数据备份和管理。我们将在 Amazon SageMaker HyperPod 存储设计与实践（二）中实现 SageMaker HyperPod 集群创建中对 Mountpoint for S3 的支持，以及在不同场景下的存储使用与实践。

本篇作者