KT使用Amazon SageMaker减少视觉转换模型训练时间的旅程
KT 通过使用 Amazon SageMaker 实现视觉转换模型训练时间缩减的旅程
KT Corporation 是韩国最大的电信服务提供商之一，提供固定电话、移动通信和互联网以及人工智能服务等多种服务。KT的AI Food Tag是一种基于人工智能的饮食管理解决方案，使用计算机视觉模型在照片中识别食物的类型和营养含量。KT开发的这个视觉模型依靠预先经过大量未标记图像数据训练的模型来分析各种食物的营养含量和卡路里信息。AI Food Tag可以帮助糖尿病等慢性病患者管理饮食。KT利用AWS和Amazon SageMaker训练了这个AI Food Tag模型，比以前快29倍，并使用模型蒸馏技术进行了生产部署的优化。在本文中，我们将介绍KT的模型开发过程以及使用SageMaker取得的成功。
介绍KT项目并定义问题
K与T训练的AI Food Tag模型基于Vision Transformers（ViT）架构，具有比以前的视觉模型更多的模型参数，以提高精确度。为了在生产环境中减小模型尺寸，KT使用了知识蒸馏（KD）技术，以减少模型参数的数量，而不会对精确度产生显著影响。通过知识蒸馏，预训练模型被称为教师模型，并训练一个轻量级输出模型作为学生模型，如下图所示。轻量级学生模型比教师模型具有更少的模型参数，从而降低了内存需求，可以部署在更小、更便宜的实例上。学生模型通过学习教师模型的输出来保持可接受的精确度，尽管它更小。
在知识蒸馏期间，教师模型保持不变，但使用教师模型的输出logits作为标签来计算损失来训练学生模型。在这种知识蒸馏范式中，教师和学生模型都需要在单个GPU内存上进行训练。KT最初在其内部的本地环境中使用了两个GPU（A100 80 GB）来训练学生模型，但这个过程需要约40天覆盖300个时期。为了加速训练，并在更短的时间内生成一个学生模型，KT与AWS合作。两个团队共同大大缩短了模型训练时间。本文描述了团队如何使用Amazon SageMaker Training、SageMaker Data Parallelism Library、Amazon SageMaker Debugger和Amazon SageMaker Profiler成功开发了一个轻量级的AI Food Tag模型。
使用SageMaker构建分布式训练环境
SageMaker Training是AWS上的托管机器学习（ML）训练环境，提供一套功能和工具来简化训练体验，并可以在分布式计算中发挥作用，如下图所示。
SageMaker客户还可以访问内置的Docker镜像，其中包含各种预装的深度学习框架以及用于模型训练的必要Linux、NCCL和Python软件包。希望运行模型训练的数据科学家或ML工程师可以在不必配置训练基础设施或管理Docker和不同库的兼容性的情况下进行训练。
在一天的研讨会中，我们能够在KT的AWS账户中基于SageMaker建立分布式训练配置，使用SageMaker分布式数据并行（DDP）库加速KT的训练脚本，甚至使用两个ml.p4d.24xlarge实例测试训练作业。在本节中，我们将描述KT与AWS团队合作并使用SageMaker开发模型的经验。
在概念验证中，我们希望通过使用SageMaker DDP库来加快训练作业，该库在分布式训练过程中针对AWS基础设施进行了优化。要从PyTorch DDP切换到SageMaker DDP，您只需要声明torch_smddp包并将后端更改为smddp，如下所示：
import smdistributed.dataparallel.torch.torch_smddpdist.init_process_group(backend='smddp',rank=args.rank,world_size=args.world_size)要了解有关SageMaker DDP库的更多信息，请参阅SageMaker的数据并行性库。
利用SageMaker调试器和分析仪分析训练速度慢的原因
优化和加速训练工作负载的第一步是了解和诊断瓶颈发生的位置。对于KT的训练作业，我们测量了数据加载器、前向传递和反向传递的每次迭代的训练时间：
| 1次迭代时间 – 数据加载器：0.00053秒，前向传递：7.77474秒，反向传递：1.58002秒 |
| 2次迭代时间 – 数据加载器：0.00063秒，前向传递：0.67429秒，反向传递：24.74539秒 |
| 3次迭代时间 – 数据加载器：0.00061秒，前向传递：0.90976秒，反向传递：8.31253秒 |
| 4次迭代时间 – 数据加载器：0.00060秒，前向传递：0.60958秒，反向传递：30.93830秒 |
| 5次迭代时间 – 数据加载器：0.00080秒，前向传递：0.83237秒，反向传递：8.41030秒 |
| 6次迭代时间 – 数据加载器：0.00067秒，前向传递：0.75715秒，反向传递：29.88415秒 |
通过查看每次迭代的标准输出时间，我们发现反向传递的运行时间在迭代之间有很大波动。这种变化是不寻常的，可能会影响总的训练时间。为了找到训练速度不一致的原因，我们首先尝试通过使用系统监视器（SageMaker调试器UI）来识别资源瓶颈，该工具允许您在一定的秒数内调试SageMaker训练作业，并查看托管训练平台的CPU、GPU、网络和I/O的状态。
SageMaker调试器UI提供了详细而重要的数据，可以帮助识别和诊断训练作业中的瓶颈。具体来说，CPU利用率线图和每个实例的CPU / GPU利用率热图表引起了我们的注意。
在CPU利用率线图中，我们注意到有些CPU的利用率达到了100%。
在热图中（颜色越暗表示利用率越高），我们注意到部分CPU核心在整个训练过程中利用率很高，而GPU利用率不一直很高。
从这里开始，我们开始怀疑训练速度缓慢的原因之一是CPU瓶颈。我们检查了训练脚本代码，看是否有任何东西导致CPU瓶颈。最可疑的部分是数据加载器中num_workers的大值，所以我们将这个值改为0或1以减少CPU利用率。然后我们再次运行训练作业并检查结果。
以下截图显示了缓解CPU瓶颈后的CPU利用率曲线图、GPU利用率和热图。
仅通过更改num_workers，我们看到CPU利用率显著降低，GPU利用率总体上增加。这是一个重要的改变，显著提高了训练速度。不过，我们还想看看在哪些地方可以优化GPU利用率。为此，我们使用了SageMaker Profiler。
SageMaker Profiler通过提供对操作的利用率可见性来帮助识别优化线索，包括跟踪GPU和CPU利用率指标以及训练脚本内GPU/CPU的内核消耗。它帮助用户了解哪些操作正在消耗资源。首先，要使用SageMaker Profiler，您需要向使用SageMaker SDK调用训练作业的函数中添加ProfilerConfig，如下面的代码所示：
from sagemaker import ProfilerConfig, Profilerfrom sagemaker.debugger import (ProfilerRule, rule_configs)rules=[ProfilerRule.sagemaker(rule_configs.ProfilerReport())]profiler_config = ProfilerConfig(profile_params = Profiler(cpu_profiling_duration=3600))from sagemaker.pytorch import PyTorchregion_name = 'us-west-2'image_uri=f'763104351884.dkr.ecr.{region_name}.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker'estimator = PyTorch(entry_point='train.py',source_dir='src',role=role,image_uri=image_uri,instance_count=4,instance_type='ml.p4d.24xlarge',distribution={'smdistributed': {'dataparallel': {'enabled': True}}},profiler_config=profiler_config,hyperparameters=hyperparameters,sagemaker_session=sagemaker_session,)在SageMaker Python SDK中，您可以灵活地为SageMaker Profiler添加annotate函数，以选择需要进行性能分析的代码或步骤。以下是在训练脚本中声明SageMaker Profiler的代码示例：
import smppySMProf = smppy.SMProfiler.instance()config = smppy.Config()config.profiler = {"EnableCuda": "1",}SMProf.configure(config)SMProf.start_profiling()…with smppy.annotate("Forward"):student_out = student_model(inp)with smppy.annotate("Backward"):loss.backward()…SMProf.stop_profiling()添加了上述代码后，如果使用训练脚本运行训练作业，则可以在训练运行一段时间后获取有关GPU内核消耗的操作的信息（如下图所示）。在KT的训练脚本中，我们运行了一个epoch并获得了以下结果。
当我们检查SageMaker Profiler结果中GPU内核的前五项操作耗时时，我们发现对于KT训练脚本来说，最耗时的是矩阵乘法操作，即在GPU上进行的一般矩阵乘法（GEMM）操作。有了SageMaker Profiler提供的这一重要见解，我们开始研究加速这些操作并提高GPU利用率的方法。
加速训练时间
我们回顾了各种减少矩阵乘法计算时间的方法，并应用了两个PyTorch函数。
使用ZeroRedundancyOptimizer进行优化器状态分片
如果您查看Zero Redundancy Optimizer（ZeRO），DeepSpeed/ZeRO技术通过消除模型使用的内存中的冗余，以更高的训练速度高效地训练大型模型。 PyTorch中的ZeroRedundancyOptimizer使用了优化器状态分片的技术，以降低分布式数据并行（DDP）中每个进程的内存使用量。 DDP在向后传递中使用同步梯度，以便所有优化器副本都在相同的参数和梯度值上迭代，但是每个优化器状态只保留不同DDP进程的分片，以降低内存使用量。
要使用它，您可以将现有的优化器放在optimizer_class中，并使用其余的模型参数和学习率作为参数声明一个ZeroRedundancyOptimizer。
student_optimizer = ZeroRedundancyOptimizer(student_model.parameters(),optimizer_class=torch.optim.AdamW,lr=initial_lr)自动混合精度
自动混合精度（AMP）在一些操作中使用torch.float32数据类型，而在其他操作中使用torch.bfloat16或torch.float16，以便进行快速计算和减少内存使用。特别是，因为深度学习模型在计算中对指数位比分数位更敏感，torch.bfloat16等效于torch.float32的指数位，使它们能够在减少损失的同时快速学习。torch.bfloat16仅在具有A100 NVIDIA架构（Ampere）或更高的实例上运行，例如ml.p4d.24xlarge、ml.p4de.24xlarge和ml.p5.48xlarge。
要应用AMP，在训练脚本中，您可以如上述代码中声明torch.cuda.amp.autocast，并将dtype声明为torch.bfloat16。
with torch.cuda.amp.autocast(dtype="torch.bfloat16"):teacher = teacher_model(input_data)student = student_model(input_data)loss = loss(teacher, student, target)loss.requires_grad_(True)loss.backward()student_optimizer.step()student_optimizer.zero_grad(set_to_none=True)SageMaker Profiler结果
将这两个函数应用于训练脚本并再次运行了一个周期的训练作业后，我们检查了SageMaker Profiler中GPU内核的前五项操作耗时。下图显示了我们的结果。
我们可以看到，在应用这两个Torch函数之前排在首位的GEMM操作已从前五个操作中消失，取而代之的是ReduceScatter操作，这通常在分布式训练中发生。
KT蒸馏模型的训练速度结果
我们将训练批次大小增加了128个，以考虑应用两个Torch函数带来的内存节省，最终的批次大小由1024增加到1152。最终学生模型的训练能够每天运行210个时期；KT内部训练环境和SageMaker之间的训练时间和加速比在下表中总结。
| 训练环境 | 训练GPU规格 | GPU数量 | 训练时间（小时） | 时期 | 每个时期的小时数 | 缩减比例 |
| KT内部训练环境 | A100 (80GB) | 2 | 960 | 300 | 3.20 | 29 |
| 亚马逊SageMaker | A100 (40GB) | 32 | 24 | 210 | 0.11 | 1 |
使用32个GPU代替内部的2个GPU，AWS的可扩展性使我们能够将训练作业的完成速度提高29倍。因此，在SageMaker上使用更多的GPU将会显著缩短训练时间，而整体训练成本没有任何差异。
结论
KT收敛技术中心AI2XL实验室的Vision AI Serving Technology团队领导者Park Sang-min对与AWS合作开发AI Food Tag模型发表了评论：
“最近，随着视觉领域中的基于Transformer的模型越来越多，模型参数和所需的GPU内存正在增加。我们正在使用轻量级技术解决这个问题，但一次学习需要花费很长时间，大约一个月的时间。通过与AWS的这个PoC，我们能够借助SageMaker Profiler和Debugger的帮助，识别出资源瓶颈，解决问题，并使用SageMaker的数据并行性库在四个ml.p4d.24xlarge实例上优化模型代码，以在大约一天内完成训练。”
SageMaker帮助节省了Sang-min的团队数周的模型训练和开发时间。
基于这个关于视觉模型的合作，AWS和SageMaker团队将继续与KT合作进行各种AI/ML研究项目，以通过应用SageMaker的能力改进模型开发和服务生产力。
要了解更多关于SageMaker的相关功能，请查阅以下内容：
