KT利用Amazon SageMaker减少视觉变换器模型训练时间的历程 机器学习博客

---

KT使用Amazon SageMaker加速视觉变换器模型训练的旅程

关键要点

KT公司通过知识蒸馏技术，实现了人工智能饮食标签模型的训练速度提升29倍，得以在短时间内进行生产部署。采用Amazon SageMaker训练环境，KT与AWS合作优化了模型训练过程，包括使用数据并行库和调试器等工具。通过分析训练过程中的性能瓶颈，KT有效提升了GPU利用率，并进一步缩短了训练时间。

引入KT项目及问题定义

KT公司是韩国最大的电信企业之一，提供固定电话、移动通讯、互联网和人工智能服务。KT的AI饮食标签是基于计算机视觉模型的AI饮食管理解决方案，能够识别照片中的食物类型及营养成分。该模型基于经过大量未标记图像数据预训练的视觉变换器ViT架构，比之前的视觉模型具有更多的参数以提高精度。为了缩小生产中的模型规模，KT使用了知识蒸馏KD技术，在保持准确度的前提下减少模型参数。

在知识蒸馏中，预训练模型被称为教师模型，而轻量级输出模型被称为学生模型。学生模型的参数比教师模型少，从而减少内存需求，允许在更小、更便宜的实例上部署。尽管模型更小，学生模型仍然可以通过学习教师模型的输出维持良好的准确度。

教师模型在KD过程中保持不变，但学生模型通过教师模型的输出逻辑作为标签进行训练，以计算损失。KT最初在内部环境中使用两台GPUA100 80GB进行学生模型的训练，但这个过程耗时约40天，经历了300个epochs。为了加速训练，同时在更短时间内生成学生模型，KT与AWS进行了合作，显著降低了模型训练时间。本文将介绍KT团队如何使用Amazon SageMaker训练、SageMaker数据并行库、Amazon SageMaker调试器和Amazon SageMaker剖析器成功开发轻量级AI饮食标签模型的过程。

构建基于SageMaker的分布式训练环境

Amazon SageMaker训练是AWS提供的托管机器学习ML训练环境，提供简化训练体验的多种功能和工具，同时在分布式计算中也非常有用，如下图所示。

SageMaker用户还可以访问内置的Docker镜像，其中包含各种预装的深度学习框架及必要的Linux、NCCL和Python包，便于模型训练。数据科学家或ML工程师可以在无需配置训练基础设施或管理Docker及不同库兼容性的情况下，轻松运行模型训练。

在为期一天的工作坊中，我们成功在KT的AWS账户中搭建了基于SageMaker的分布式训练配置，利用SageMaker分布式数据并行DDP库加速了KT的训练脚本，还用两台mlp4d24xlarge实例进行了训练作业测试。

在概念验证中，KT希望通过使用针对AWS基础设施优化的SageMaker DDP库加快训练作业的速度。要从PyTorch DDP切换到SageMaker DDP，只需要声明torchsmddp包，并将后端更改为smddp，示例代码如下：

pythonimport smdistributeddataparalleltorchtorchsmddp

distinitprocessgroup(backend=smddprank=argsrankworldsize=argsworldsize)

有关SageMaker DDP库的详细信息，请参阅SageMaker数据并行库。

使用SageMaker调试器与分析器分析训练速度

优化和加速训练工作负载的第一步是理解和诊断瓶颈的原因。对于KT的训练作业，我们测量了数据加载器、前向传递和反向传递的每次迭代训练时间：

迭代次序数据加载器前向反向1000053 sec777474 sec158002 sec2000063 sec067429 sec2474539 sec3000061 sec090976 sec831253 sec4000060 sec060958 sec3093830 sec5000080 sec083237 sec841030 sec6000067 sec075715 sec2988415 sec

从标准输出中可以看出，反向传递的运行时间在不同迭代之间波动很大。这种变化不寻常，会影响总的训练时间。为了找出训练速度不一致的原因，我们首先利用SageMaker调试器UI的系统监视器来识别资源瓶颈，该工具允许您调试SageMaker训练作业并查看受管训练平台的CPU、GPU、网络和I/O等资源在一定时间内的状态。

SageMaker调试器UI提供了详细且重要的数据，有助于识别和诊断训练作业中的瓶颈。特别是CPU利用率折线图和每实例的CPU/GPU利用率热力图引起了我们的注意。

在CPU利用率折线图中，我们注意到有些CPU的利用率达到了100。

在热力图中颜色越深表示利用率越高，我们注意到一些CPU核心在整个训练过程中均有较高的利用率，而GPU利用率时间上并不稳定。

从这一点来看，我们开始怀疑慢训练速度的原因之一是CPU瓶颈。为此，我们检查了训练脚本代码以查看是否有什么导致CPU瓶颈的问题。最可疑的部分是数据加载器中的numworkers值过大，因此我们将此值更改为0或1，以减少CPU利用率。然后，我们再次运行训练作业并检查结果。

以下截图显示了缓解CPU瓶颈后，CPU利用率折线图、GPU利用率和热力图。

通过简单更改numworkers，我们看到CPU利用率显著降低，GPU利用率整体提升。这一步的改变显著提高了训练速度，但我们仍想继续优化GPU利用率。因此，我们使用SageMaker Profiler进行进一步分析。

SageMaker Profiler通过提供操作利用率的可见性，帮助识别优化线索，包括跟踪训练脚本中GPU/CPU的利用率指标和内核消费。首先，要使用SageMaker Profiler，您需要在调用训练作业时通过SageMaker SDK添加ProfilerConfig，如下所示：

pythonfrom sagemaker import ProfilerConfig Profilerfrom sagemakerdebugger import (ProfilerRule ruleconfigs)

rules=[ProfilerRulesagemaker(ruleconfigsProfilerReport())]profilerconfig = ProfilerConfig(profileparams = Profiler(cpuprofilingduration=3600))

from sagemakerpytorch import PyTorchregionname = uswest2imageuri=f763104351884dkrecr{regionname}amazonawscom/pytorchtraining200gpupy310cu118ubuntu2004sagemaker

estimator = PyTorch(entrypoint=trainpysourcedir=srcrole=roleimageuri=imageuriinstancecount=4instancetype=mlp4d24xlargedistribution={smdistributed {dataparallel {enabled True}}}profilerconfig=profilerconfighyperparameters=hyperparameterssagemakersession=sagemakersession)

在SageMaker Python SDK中，您可以灵活地为SageMaker Profiler添加annotate函数，以选择需要进行分析的代码或训练脚本中的步骤。以下是您应该在训练脚本中声明的SageMaker Profiler代码示例：

pythonimport smppy

SMProf = smppySMProfilerinstance()config = smppyConfig()configprofiler = {EnableCuda 1}SMProfconfigure(config)SMProfstartprofiling()with smppyannotate(Forward)studentout = studentmodel(inp)with smppyannotate(Backward)lossbackward()SMProfstopprofiling()

添加前述代码后，如果您运行训练作业，将在训练运行一段时间后获得有关GPU内核消耗操作的信息。在KT的训练脚本中，我们运行了一个epoch，并得到了以下结果。

在检查SageMaker Profiler的GPU内核消费时间的前五个操作时，我们发现KT训练脚本中，用于矩阵积操作的时间占比较高，这是一种在GPU上执行的通用矩阵相乘GEMM操作。通过SageMaker Profiler提供的这一重要见解，我们开始探索加速这些操作并改善GPU利用率的方法。

加速训练时间

我们审查了多种减少矩阵乘法计算时间的方法，并应用了两个PyTorch函数。

使用ZeroRedundancyOptimizer分片优化器状态

Zero Redundancy OptimizerZeRO是一种DeepSpeed/ZeRO技术，通过消除模型使用的内存冗余，能够有效地训练大型模型，提高训练速度。PyTorch中的ZeroRedundancyOptimizer利用分片优化器状态的技术，减少分布式数据并行DDP中每个进程的内存占用。

要使用它，您可以将现有的优化器保留在optimizerclass中，并声明一个ZeroRedundancyOptimizer，将模型的其余参数和学习率作为参数传入。

pythonstudentoptimizer = ZeroRedundancyOptimizer(studentmodelparameters()optimizerclass=torchoptimAdamWlr=initiallr)

自动混合精度

自动混合精度AMP对某些操作使用torchfloat32数据类型，而对其他操作使用torchbfloat16或torchfloat16，以实现快速计算和减少内存使用。由于深度学习模型通常对指数位比小数位更敏感，因此torchbfloat16等价于torchfloat32的指数位，使得模型能够快速学习且损失最小。torchbfloat16仅在支持A100 NVIDIA架构Ampere或更高版本的实例上运行，例如mlp4d24xlarge、mlp4de24xlarge和mlp548xlarge。

要应用AMP，您可以在训练脚本中声明torchcudaampautocast，并将dtype设置为torchbfloat16。

pythonwith torchcudaampautocast(dtype=torchbfloat16)teacher = teachermodel(inputdata)student = studentmodel(inputdata)loss = loss(teacher student target)lossrequiresgrad(True)lossbackward()studentoptimizerstep()studentoptimizerzerograd(settonone=True)

SageMaker Profiler中的结果

在将这两个函数应用到训练脚本并再次运行一个epoch的训练作业后，我们检查了SageMaker Profiler的GPU内核前五个操作的时间消耗。以下图显示了我们的结果。

可以看到，之前在列表顶部的GEMM操作在应用这两个Torch函数后消失了，取而代之的是ReduceScatter操作，而该操作通常出现在分布式训练中。

魔方网络服务加速器

KT蒸馏模型的训练速度结果

我们将训练批次大小增加了128，以考虑到应用两个Torch函数带来的内存节省，最终批次大小为1152，而非1024。最终学生模型的训练能够在一天内完成210个epochs。KT内部训练环境与SageMaker的训练时间和加速比在下表中汇总：

训练环境训练GPU规格GPU数量训练时间小时Epoch数量每个Epoch小时减少比率KT内部训练环境A10080GB296030032029Amazon SageMakerA10040GB32242100111

AWS的可扩展性使我们能够使用32个GPU而非2个，完成训练作业，并将训练时间提升至29倍。由此可见，使用SageMaker上更多GPU可以在不增加总体训练成本的情况下显著减少训练时间。

结论

KT的AI2XL实验室的Vision AI服务技术团队领导者朴尚敏对与AWS合作开发AI饮食标签模型的评论表示：

“最近，随着视觉领域中基于变换器的模型增多，模型参数和所需GPU内存也随之增加。我们正在采取轻量化技术来解决这个问题，而训练一次通常需要大约一个月。通过与AWS的此次概念验证，我们能够借助SageMaker Profiler和调试器识别资源瓶颈并解决它们，进而使用SageMaker的数据并行库在约一天内完成训练。”

SageMaker帮助朴尚敏团队节省了数周的模型训练和开发时间。

基于此次视觉模型的合作，AWS及SageMaker团队将继续与KT开展多项AI/ML研究项目，通过应用SageMaker能力提升模型开发和服务生产力。

欲了解有关SageMaker相关功能的更多信息，请访问以下链接：

训练机器学习模型SageMaker的数据并行库使用Amazon SageMaker调试器调试和改善模型性能使用Amazon SageMaker剖析器分析AWS计算资源的活动

关于作者

Youngjoon Choi，AI/ML专家SA，拥有制造业、高科技和金融等多个行业的企业IT经验，曾担任

---