优化器状态分片 - Amazon SageMaker

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

优化器状态分片

优化器状态分片是一种非常有用的内存节省技术,它可以跨数据并行设备组对优化器状态(描述优化器状态的一组权重)进行分片。每当使用有状态优化器(例如 Adam)或优化器(存储参数FP16和参数FP32副本)时,都可以使用FP16优化器状态分片。

注意

优化器状态分片可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。

如何使用优化器状态分片

您可以通过在 modelparallel 配置设置 "shard_optimizer_state": True 来启用优化器状态分片

启用此功能后,库会根据数据并行度对模型参数集进行分区。与第 i 个分区对应的梯度,仅在第 i 个数据并行秩处缩减。在对 smp.step 修饰器函数的第一次调用结束时,被 smp.DistributedOptimizer 包装的优化器重新定义了其参数,使其仅限于与当前数据并行秩的分区相对应的参数。重新定义的参数称为虚拟参数,它们与原始参数共享底层存储。在第一次调用 optimizer.step 期间,优化器状态是根据这些重新定义的参数创建的,这些状态由于原始分区而被分片。优化器更新后,该 AllGather操作(作为optimizer.step调用的一部分)在数据 parallel 等级中运行,以实现一致的参数状态。

提示

当数据并行度大于 1 且模型的参数超过 10 亿时,优化器状态分片可能很有用。

数据并行性由 (processes_per_host * instance_count / pipeline_parallel_degree) 计算得出,smp.dp_size() 函数在后台处理大小调整。

配置 SageMaker PyTorch 估算器

mpi_options = { "enabled" : True, "processes_per_host" : 8, # 8 processes "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " } smp_options = { "enabled":True, "parameters": { "microbatches": 4, "pipeline_parallel_degree": 2, # alias for "partitions" "placement_strategy": "cluster", "tensor_parallel_degree": 2, # tp over 2 devices "ddp": True, "shard_optimizer_state": True } }

调整您的 PyTorch 训练脚本

请参阅 Tensor 并行度与流水线并行度相结合部分中的调整 PyTorch 训练脚本。无需对脚本进行其他修改。