使用Amazon SageMaker AI将基准测试和推荐结果流式传输到MLflow

在生成式AI模型的部署优化过程中，团队往往需要评估多种GPU实例类型、服务容器、并行策略以及推测解码等高级技术。传统做法通常导致数周的手动数据整理和结果对比，效率低下且容易出错。为解决这一痛点，Amazon SageMaker AI推出了全新的MLflow集成，允许用户将基准测试和推荐结果自动流式传输到统一的SageMaker MLflow应用，从而消除数据孤岛，加速迭代周期，并确保实验的可复现性。

该集成的核心优势包括：

• 自动消除手动数据整合：多个作业的结果自动合并到同一实验名称下，用户可以在MLflow实验视图中并排比较不同运行，无需手动收集指标、日志和配置。例如，可以直观地比较qwen2-0.5b模型在ml.g4dn.12xlarge与ml.p4d.24xlarge实例上的性能差异。
• 实时监控长时间运行的作业：基准测试和推荐作业可能耗时数小时。现在，用户可以实时观察延迟和吞吐量指标在测试过程中的持续更新，如果吞吐量不符合预期，可以提前停止作业，避免资源浪费。
• 完整的审计追踪：每个实验运行都捕获了完整的上下文，包括作业参数、时间戳、检查点指标和生成的工件。审计记录可查询和重现，帮助团队追溯配置变更带来的改进或识别性能瓶颈。
• 更好的协作与治理：共享的MLflow实验成为优化工作的单一事实来源，减少重复劳动，支持团队成员之间的无缝交接。

技术实现上，用户需要先通过Amazon SageMaker Studio创建SageMaker MLflow应用，并为执行角色添加必要的权限（包括sagemaker-mlflow:*和端点调用权限）。提交基准测试或推荐作业时，在OutputConfig中指定MlflowConfig，包含MLflow应用ARN、实验名称和运行名称。基准测试作业评估现有端点，推荐作业则自动在内部配置和评估端点，并返回排名靠前的配置。例如，使用Qwen/Qwen2-0.5B-Instruct模型的团队可以通过同一个MLflow实验跟踪基准测试和推荐作业的结果，典型完成时间为45-120分钟。

关键注意事项包括：集成仅支持SageMaker MLflow应用，不支持自托管MLflow服务器；需要设置tooling.version为0.8.0或更高以支持MLflow嵌套运行；推荐作业无需指定ComputeSpec.InstanceTypes（除非有特殊限制）；执行角色必须包含所需的SageMaker、MLflow、S3和端点调用权限；S3输出桶必须与作业在同一区域。

通过这个新的MLflow集成，团队可以显著简化生成式AI模型的部署优化流程，将精力集中在模型性能提升而非数据整理上。