核心内容摘要
MogFace模型剪枝与蒸馏实战：减少耦合过度，打造轻量级子网络
随着大语言模型LLM的广泛应用如何优化推理性能已成为企业级应用部署中的核心挑战传统的无状态推理架构在处理LLM应用时面临诸多性能瓶颈每次请求被随机路由到不同的计算实例导致KV Cache无法有效复用、多轮对话上下文频繁重建、系统提示词重复处理等问题严重影响了用户体验和系统效率。
针对这些痛点Amazon SageMaker推理集群的粘性会话路由Sticky Session Routing功能通过会话绑定机制确保同一用户会话的所有请求都路由到相同的推理集群实例从而实现缓存复用和状态保持显著提升LLM推理性能和交互体验。
技术原理与优势Amazon SageMaker Sticky Session技术原理Amazon SageMaker Sticky Session通过会话标识符Session ID实现智能路由机制。
与传统的负载均衡策略不同该机制确保来自同一会话的所有请求都被定向到特定的计算实例。
这种“会话亲和性”Session Affinity设计使得GPU内存中的KV Cache、中间结果等状态信息得以持续保存和复用从而避免了重复的计算开销。
核心优势启用粘性会话后同一会话的所有请求都会路由到同一实例这样您的AI应用程序就能重复使用先前处理过的信息避免了不必要的重复计算从而减少延迟并改善用户体验。
当客户想要使用大型数据负载或需要无缝的交互体验时这项功能尤其有用。
通过利用以前的推理请求客户现在可以借助这项功能在Amazon SageMaker上构建创新的状态感知AI应用程序。
工作流程一起看看客户端在Amazon SageMaker上启动粘性会话时的事件序列Sticky Session的完整工作流程包含三个核心阶段会话创建、会话维持和会话关闭:
在第一个请求中当调用Boto3 SageMaker runtime的invoke_endpoint时在请求头中设置session-idNEW_SESSION并在请求负载中指明这是一个打开会话类型的请求。
然后Amazon SageMaker会创建一个新的会话并存储会话ID。
Amazon SageMaker会向推理服务器发起一个打开会话的请求这个API是由客户端定义的也可以使用其他名称比如start_session并返回200状态码同时附带会话ID和生存时间TTL这些信息会被发送回客户端。
后续请求中客户端在invoke_endpoint调用中传递相同的会话ID。
Amazon SageMaker路由服务根据会话ID识别对应的推理实例确保所有请求都路由到同一台服务器。
这样在LLM推理的多轮对话场景中GPU内存中的KV Cache得以保持和累积显著减少了计算开销从而显著降低推理延迟。
当对话结束时客户端可以发送关闭会话请求您可以使用invoke_endpoint传入一个包含会话ID的关闭请求。
Amazon SageMaker路由器首先会检查该会话是否存在。
如果存在路由器会向推理服务器发起关闭会话的调用服务器会返回一个成功的200响应以及会话ID这些信息随后会被发送回客户端。
在会话ID不存在的情况下路由器会返回一个400响应。
这种设计带来的性能提升主要体现在首次响应时间优化避免重复处理系统提示词和上下文信息后续对话加速利用已缓存的键值对减少大量重复计算内存使用效率合理利用GPU内存避免频繁的内存分配和释放
实践案例构建支持Sticky Session的LLM推理服务本案例将以基于SGLang推理引擎的模型部署为例展示如何在Amazon SageMaker上构建支持Sticky Session的高性能LLM推理服务。
整个实现涉及推理容器构建、模型部署和客户端调用三个核心部分。
构建支持会话管理的推理容器在推理服务器端使用Sanic Web框架构建了支持会话管理的推理服务。
核心实现包括三种请求类型的处理逻辑app.route(/invocations, methods[POST])async def generate(request): reqType request.json.get(requestType) extSessID request.json.get(extSessionID) # 处理新建Session请求 ifNEW_SESSION reqType: current_time datetime.now(dt.timezone.utc) future_time current_time timedelta(minutes int(os.environ[SES_TTL_MIN])) formatted_time future_time.strftime(%Y-%m-%dT%H:%M:%SZ) response json({}) response.headers[X-Amzn-SageMaker-Session-Id] f{extSessID}; Expires{formatted_time} return response # 处理Session关闭请求 elif CLOSE_SESSION reqType: response json({}) response.headers[X-Amzn-SageMaker-Closed-Session-Id] extSessID return response # 处理Session内推理请求 else: prompt request.json.get(inputs) ifnot prompt: return json({error: inputs is required}, status
inf_params request.json.get(parameters) # 使用部署至 SGLang 推理引擎的模型产生图里结果 result await engine.async_generate(promptprompt, sampling_paramsinf_params) return json({generation: result})左右滑动查看完整示意构建客户端调用封装为了简化Amazon SageMaker Sticky Session的使用本例封装了一个客户端工具类提供了会话管理的完整功能import boto3, json classStatefulSMEDPBuilder: def __init__(self, endpoint_name): self.endpoint_name endpoint_name ## Boto3 API self.sm_bt3_client boto
client(runtime.sagemaker) def start_session(self, extSessID): 创建新的粘性会话 payload { extSessionID: extSessID, requestType: NEW_SESSION } response self.sm_bt3_client.invoke_endpoint( EndpointNameself.endpoint_name, Bodyjson.dumps(payload), ContentTypeapplication/json, SessionIdNEW_SESSION ) return response def end_session(self, extSessID): 关闭指定会话 payload { extSessionID: extSessID, requestType: CLOSE_SESSION } response self.sm_bt3_client.invoke_endpoint( EndpointNameself.endpoint_name, Bodyjson.dumps(payload), ContentTypeapplication/json, SessionIdextSessID ) return response def invoke(self, textPayload, sampling_paramsNone, extSessIDNone): 在会话中进行推理调用 if None sampling_params: sampling_params {temperature:
9, max_new_token:128, do_sample:True} payload { inputs: textPayload, sampling_params: sampling_params, extSessionID: extSessID, requestType: SESSION } response self.sm_bt3_client.invoke_endpoint( EndpointNameself.endpoint_name, Bodyjson.dumps(payload), ContentTypeapplication/json, SessionIdextSessID ) return response def invoke_stream(self, textPayload, sampling_paramsNone, extSessIDNone): 流式推理接口 pass左右滑动查看完整示意使用示例下面展示如何使用封装好的客户端进行多轮对话import uuid # 初始化客户端client StatefulSMEDPBuilder(endpoint_namellama3-sticky-session-endpoint) # 生成唯一会话 IDsession_id fconversation-{uuid.uuid4().hex[:8]} try: # 创建新会话 session_response client.start_session(session_id) print(f会话 {session_id} 创建成功) # 第一轮对话 response1 client.invoke(你好请介绍一下自己, extSessIDsession_id) result1 json.loads(response1[Body].read().decode()) print(fAI: {result1[generation][text]}) # 第二轮对话利用上下文 response2 client.invoke(请详细说明你刚才提到的能力, extSessIDsession_id) result2 json.loads(response2[Body].read().decode()) print(fAI: {result2[generation][text]}) # 第三轮对话继续利用完整上下文 response3 client.invoke(基于我们之前的对话你觉得哪个能力最重要, extSessIDsession_id) result3 json.loads(response3[Body].read().decode()) print(fAI: {result3[generation][text]}) finally: # 会话结束后清理资源 client.end_session(session_id) print(f会话 {session_id} 已关闭)左右滑动查看完整示意通过这种方式同一会话的所有请求都会路由到相同的推理实例确保KV Cache的有效复用显著提升多轮对话的性能表现。
完整的部署和使用代码可参考项目仓库。
项目仓库https://github.com/aws-samples/sample-sagemaker-sticky-session
总结本文深入介绍了Amazon SageMaker Sticky Session在LLM推理中的技术原理和实践应用。
通过会话亲和性设计该技术有效解决了传统无状态推理架构在处理LLM应用时面临的性能瓶颈这一功能特别适合需要维持对话上下文的应用以及要求低延迟的实时交互场景。
后续博客将展示详细的性能对比测试结果。
参考链接
https://aws.amazon.com/cn/about-aws/whats-new/2024/09/sticky-session-routing-amazon-sagemaker-inference
https://aws.amazon.com/cn/blogs/machine-learning/build-ultra-low-latency-multimodal-generative-ai-applications-using-sticky-session-routing-in-amazon本篇作者胡凯亚马逊云科技解决方案架构师负责基于亚马逊云科技的云计算方案架构的咨询和设计致力于亚马逊云科技云服务在企业的应用和推广。
郑昊亚马逊云科技AI/ML解决方案架构师。
主要专注于基础模型的训练、推理及性能优化广告、排序算法等及其基于亚马逊云科技AI/ML技术栈的相关优化及方案构建拥有多年排序、定价及竞价机制等算法研发经验。
新用户注册海外区域账户可获得最高200美元服务抵扣金覆盖Amazon Bedrock生成式AI相关服务。
“免费计划”账户类型确保零花费安心试用。
星标不迷路开发更极速关注后记得星标「亚马逊云开发者」听说点完下面4个按钮就不会碰到bug了点击阅读原文查看博客获得更详细内容