使用异构集群进行训练 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

使用异构集群进行训练

使用 SageMaker Training 的异构集群功能,您可以使用多种类型的机器学习实例运行训练作业,从而更好地扩展和利用资源以实施不同的机器学习训练任务和目的。例如,如果带 GPU 实例的集群上的训练作业因 CPU 密集型任务而遇到 GPU 利用率低和 CPU 瓶颈问题,则可使用异构集群,通过添加更具成本效益的 CPU 实例组来协助卸载 CPU 密集型任务,解决此类瓶颈问题并实现更高的 GPU 利用率。

注意

此功能在 SageMaker Python SDK v2.98.0 及更高版本中可用。

注意

可通过 SageMaker PyTorchTensorFlow 框架估算器类使用此功能。支持的框架为 PyTorch v1.10 或更高版本以及 TensorFlow v2.6 或更高版本。

如何配置异构集群

此部分提供了有关如何使用由多种实例类型组成的异构集群运行训练作业的说明。

使用 SageMaker Python SDK

遵循有关如何使用 SageMaker Python SDK 为异构集群配置实例组的说明。

  1. 要为训练作业配置异构集群的实例组,请使用 sagemaker.instance_group.InstanceGroup 类。您可以指定每个实例组的自定义名称、实例类型和每个实例组的实例数。有关更多信息,请参阅 SageMaker Python SDK 文档中的 sagemaker.instance_group.InstanceGroup

    注意

    有关可用实例类型以及您可在异构集群中配置的实例组的最大数目的更多信息,请参阅 InstanceGroup API 参考。

    以下代码示例说明如何设置两个实例组,它们由两个名为 instance_group_1ml.c5.18xlarge 仅 CPU 实例和一个名为 instance_group_2ml.p3dn.24xlarge GPU 实例组成,如下图所示。

    The preceding diagram shows a conceptual example of how pre-training processes, such as data preprocessing, can be assigned to the CPU instance group and stream the preprocessed data to the GPU instance group.
    from sagemaker.instance_group import InstanceGroup instance_group_1 = InstanceGroup( "instance_group_1", "ml.c5.18xlarge", 2 ) instance_group_2 = InstanceGroup( "instance_group_2", "ml.p3dn.24xlarge", 1 )
  2. 通过使用实例组对象,设置训练输入通道,并通过 sagemaker.inputs.TrainingInput 类的 instance_group_names 参数将实例组分配给通道。instance_group_names 参数接受实例组名称的字符串列表。

    以下示例说明如何设置两个训练输入通道并分配在上一步的示例中创建的实例组。您还可以为实例组指定 s3_data 参数的 Amazon S3 存储桶路径,处理数据以供使用。

    from sagemaker.inputs import TrainingInput training_input_channel_1 = TrainingInput( s3_data_type='S3Prefix', # Available Options: S3Prefix | ManifestFile | AugmentedManifestFile s3_data='s3://your-training-data-storage/folder1', distribution='FullyReplicated', # Available Options: FullyReplicated | ShardedByS3Key input_mode='File', # Available Options: File | Pipe | FastFile instance_groups=["instance_group_1"] ) training_input_channel_2 = TrainingInput( s3_data_type='S3Prefix', s3_data='s3://your-training-data-storage/folder2', distribution='FullyReplicated', input_mode='File', instance_groups=["instance_group_2"] )

    有关 TrainingInput 的参数的更多信息,请参阅以下链接。

  3. 使用 instance_groups 参数配置 SageMaker 估算器,如以下代码示例中所示。instance_groups 参数接受 InstanceGroup 对象列表。

    PyTorch
    from sagemaker.pytorch import PyTorch estimator = PyTorch( ... entry_point='my-training-script.py', framework_version='x.y.z', # 1.10.0 or later py_version='pyxy', job_name='my-training-job-with-heterogeneous-cluster', instance_groups=[instance_group_1, instance_group_2] )
    TensorFlow
    from sagemaker.tensorflow import TensorFlow estimator = TensorFlow( ... entry_point='my-training-script.py', framework_version='x.y.z', # 2.6.0 or later py_version='pyxy', job_name='my-training-job-with-heterogeneous-cluster', instance_groups=[instance_group_1, instance_group_2] )
    注意

    instance_typeinstance_count 参数对与 SageMaker 估算器类的 instance_groups 参数是互斥的。对于同构集群训练,请使用 instance_typeinstance_count 参数对。对于异构集群训练,请使用 instance_groups

    注意

    要查找可用的框架容器、框架版本和 Python 版本的完整列表,请参阅 Amazon 深度学习容器 GitHub 存储库中的 SageMaker 框架容器

  4. 使用配置了实例组的训练输入通道配置 estimator.fit 方法,然后启动训练作业。

    estimator.fit( inputs={ 'training': training_input_channel_1, 'dummy-input-channel': training_input_channel_2 } )

使用低级 SageMaker API

如果您使用 Amazon Command Line Interface 或 Amazon SDK for Python (Boto3) 并希望使用低级 SageMaker API 来提交异构集群的训练作业请求,请参阅以下 API 参考。

使用异构集群进行分布式训练

通过 SageMaker 估算器类的 distribution 参数,您可以分配特定的实例组来运行分布式训练。例如,假定您有以下两个实例组,并且想在其中一个实例组上运行多 GPU 训练。

from sagemaker.instance_group import InstanceGroup instance_group_1 = InstanceGroup("instance_group_1", "ml.c5.18xlarge", 1) instance_group_2 = InstanceGroup("instance_group_2", "ml.p3dn.24xlarge", 2)

您可以为其中一个实例组设置分布式训练配置。例如,以下代码示例说明如何将带两个 ml.p3dn.24xlarge 实例的 training_group_2 分配给分布式训练配置。

注意

目前,只能为分发配置指定异构群集的一个实例组。

使用 MPI

PyTorch
from sagemaker.pytorch import PyTorch estimator = PyTorch( ... instance_groups=[instance_group_1, instance_group_2], distribution={ "mpi": { "enabled": True, "processes_per_host": 8 }, "instance_groups": [instance_group_2] } )
TensorFlow
from sagemaker.tensorflow import TensorFlow estimator = TensorFlow( ... instance_groups=[instance_group_1, instance_group_2], distribution={ "mpi": { "enabled": True, "processes_per_host": 8 }, "instance_groups": [instance_group_2] } )

使用 SageMaker 数据并行库

PyTorch
from sagemaker.pytorch import PyTorch estimator = PyTorch( ... instance_groups=[instance_group_1, instance_group_2], distribution={ "smdistributed": { "dataparallel": { "enabled": True } }, "instance_groups": [instance_group_2] } )
TensorFlow
from sagemaker.tensorflow import TensorFlow estimator = TensorFlow( ... instance_groups=[instance_group_1, instance_group_2], distribution={ "smdistributed": { "dataparallel": { "enabled": True } }, "instance_groups": [instance_group_2] } )
注意

使用 SageMaker 数据并行库时,请确保实例组包含库支持的实例类型

有关 SageMaker 数据并行库的更多信息,请参阅 SageMaker 数据并行训练

使用 SageMaker 模型并行库

PyTorch
from sagemaker.pytorch import PyTorch estimator = PyTorch( ... instance_groups=[instance_group_1, instance_group_2], distribution={ "smdistributed": { "modelparallel": { "enabled":True, "parameters": { ... # SageMaker model parallel parameters } } }, "instance_groups": [instance_group_2] } )
TensorFlow
from sagemaker.tensorflow import TensorFlow estimator = TensorFlow( ... instance_groups=[instance_group_1, instance_group_2], distribution={ "smdistributed": { "modelparallel": { "enabled":True, "parameters": { ... # SageMaker model parallel parameters } } }, "instance_groups": [instance_group_2] } )

有关 SageMaker 模型并行库的更多信息,请参阅 SageMaker 模型并行训练

修改训练脚本以分配实例组

使用前几个部分中的异构集群配置,您已为训练作业准备 SageMaker 训练环境和实例。要进一步将实例组分配给某些训练和数据处理任务,下一步是修改训练脚本。默认情况下,无论实例大小如何,训练作业都只为所有节点创建训练脚本副本,这可能会导致性能下降。

例如,如果您在异构集群中混合 CPU 实例和 GPU 实例,同时将深度神经网络训练脚本传递给 SageMaker 估算器的 entry_point 参数,则 entry_point 脚本将复制到每个实例。这意味着,如果没有适当的任务分配,CPU 实例还会运行整个脚本并启动专为 GPU 实例上的分布式训练设计的训练作业。因此,您必须更改要卸载并在 CPU 实例上运行的特定处理函数。您可以使用 SageMaker 环境变量来检索异构集群的信息,并相应地让特定的进程运行。

在 SageMaker 训练作业的初始化阶段查询实例组信息

当训练作业启动时,训练脚本会读取 SageMaker 训练环境信息,其中包括异构集群配置。该配置包含诸如当前实例组、每个组中的当前主机以及当前主机所在的组之类的信息。

您可以通过以下方式检索实例组信息。

(推荐)使用 SageMaker 训练工具包读取实例组信息

使用 SageMaker 训练工具包库提供的环境 Python 模块。工具包库预安装在适用于 TensorFlow 和 PyTorch 的 SageMaker 框架容器中,因此,在使用预构建容器时,无需执行额外的安装步骤。这是检索 SageMaker 环境变量的推荐方法,只需对训练脚本进行少量代码更改。

from sagemaker_training import environment env = environment.Environment()

与常规 SageMaker 训练和异构集群相关的环境变量:

  • env.is_hetero – 返回一个布尔结果,指示是否配置了异构集群。

  • env.current_host – 返回当前主机。

  • env.current_instance_type – 返回当前主机的实例的类型。

  • env.current_instance_group – 返回当前实例组的名称。

  • env.current_instance_group_hosts – 返回当前实例组中的主机列表。

  • env.instance_groups – 返回用于训练的实例组名称的列表。

  • env.instance_groups_dict – 返回训练作业的整个异构集群配置。

  • env.distribution_instance_groups – 返回分配给 SageMaker 估算器类的 distribution 参数的实例组列表。

  • env.distribution_hosts – 返回属于分配给 SageMaker 估算器类的 distribution 参数的实例组的主机列表。

例如,考虑以下包含两个实例组的异构集群示例。

from sagemaker.instance_group import InstanceGroup instance_group_1 = InstanceGroup( "instance_group_1", "ml.c5.18xlarge", 1) instance_group_2 = InstanceGroup( "instance_group_2", "ml.p3dn.24xlarge", 2)

示例异构集群的 env.instance_groups_dict 的输出应类似于以下内容。

{ "instance_group_1": { "hosts": [ "algo-2" ], "instance_group_name": "instance_group_1", "instance_type": "ml.c5.18xlarge" }, "instance_group_2": { "hosts": [ "algo-3", "algo-1" ], "instance_group_name": "instance_group_2", "instance_type": "ml.p3dn.24xlarge" } }

(可选)从资源配置 JSON 文件中读取实例组信息

如果您更喜欢以 JSON 格式检索环境变量,则可以直接使用资源配置 JSON 文件。默认情况下,SageMaker 训练实例中的 JSON 文件位于 /opt/ml/input/config/resourceconfig.json

file_path = '/opt/ml/input/config/resourceconfig.json' config = read_file_as_json(file_path) print(json.dumps(config, indent=4, sort_keys=True))

注意事项

在使用异构集群功能时,请考虑以下项目。

  • 所有实例组共享相同的 Docker 映像和训练脚本。因此,应修改训练脚本以相应地检测它所属的实例组和分叉执行。

  • SageMaker 本地模式不支持异构集群功能。

  • 异构集群训练作业的 Amazon CloudWatch 日志流不按实例组进行分组。您需要从日志中查明哪些节点属于哪个组。

  • 可通过 SageMaker PyTorchTensorFlow 框架估算器类使用异构集群功能。支持的框架为 PyTorch v1.10 或更高版本以及 TensorFlow v2.6 或更高版本。要查找可用的框架容器、框架版本和 Python 版本的完整列表,请参阅 Amazon 深度学习容器 GitHub 存储库中的 SageMaker 框架容器

  • 一个分布式训练策略只能应用于一个实例组。

示例、博客和案例研究

以下博客讨论有关使用 SageMaker 异构集群训练的案例研究。