修改您的训练脚本以分配实例组 - Amazon SageMaker

本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

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

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

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

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

在 SageMaker训练作业的初始化阶段，您可以通过以下方式查询实例组信息。

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

使用SageMaker 训练工具库提供的环境 Python 模块。工具包库已预先安装在 TensorFlow 和的SageMaker 框架容器中 PyTorch，因此在使用预构建的容器时，您无需执行额外的安装步骤。这是检索 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))