简介大模型推理部署通常被Python生态主导但在企业内网隔离、后端以Spring Boot为主且缺乏Python运维能力的场景下Java直接承载LLaMA2推理反而能显著降低系统复杂度。本文从显存规划与技术选型切入对比DJL、ONNX Runtime与TensorRT Java绑定三条路线说明如何基于DJLPyTorch引擎加载LLaMA2模型并深入解析数据并行与张量并行的多GPU实现策略。同时结合KV Cache管理、CUDA_VISIBLE_DEVICES配置、JVM内存与GC调优等实战经验给出单卡与多卡性能实测数据为Java工程师在现有基础设施中集成大模型推理能力提供完整参考。 先说结论用Java做LLaMA2的推理部署乍看像是一个伪需求。毕竟LLaMA2的官方实现、HuggingFace生态、各种优化加速库基本都是Python的天下硬要往Java上靠总给人一种脱裤子放屁的感觉。但如果你所在的企业后端是清一色的Spring Boot机房里有几块闲置的NVIDIA GPU而业务方又希望在一个内网隔离系统里直接调用大模型的能力时Java多GPU这套组合反而是性价比最高的路。这篇文章是我把一个基于Java多GPU的LLaMA2推理服务从零推到线上的完整复盘包含了算账、选型、编码、排障各阶段的细节希望对正在纠结这套方案的你有参考价值。1. 为什么要把LLaMA2塞进Java进程先聊聊这个方案的适用边界很多人一听到Java跑LLaMA2第一反应都是“何必呢”。确实如果目标只是跑通一个DemoPython一条命令就能启动服务Java连渣都赶不上。但在实际生产环境里事情往往不是“哪个语言跑模型最方便”而是“在现有系统架构下哪个方案的综合成本最低”。1.1 哪些场景下Java推理是刚需而不是炫技我接手这个项目时业务方已经有一套基于Spring Boot的文档审核系统每天要处理大量内部文档。那边的诉求是在文档上传后能自动调用大模型做摘要和关键词提取。看似简单但有两个硬约束第一服务必须部署在内网和公网完全隔离不能调用云厂商API第二团队里只有Java工程师没有人愿意为了一个摘要功能额外维护一套Python微服务。这种情况下Java直接调用LLaMA2就成了最优解。因为业务代码和推理逻辑可以打包在同一个JVM进程里走一次http调用就能拿到结果不需要跨语言、跨服务。系统监控、日志、链路追踪也能复用原有的基础设施不需要为Python单独搭建一套。对于运维来说一个Jar包丢到机器上就能跑远比“再维护一个Python环境、再装一堆pip依赖”要好接受。1.2 什么场景下不该用Java硬扛这里也要泼一盆冷水。如果你的目标是做一个高并发的LLM在线服务平台比如面向C端的聊天机器人要支撑几百路的并发推理那Java方案大概率不是最优选择。vLLM、TensorRT-LLM、Triton Inference Server这些专门针对大模型优化过的Python/C方案在吞吐量和显存利用效率上要领先Java生态一大截。我的建议是如果你的核心业务就是“卖模型推理能力”别碰Java用专业推理框架如果你的核心业务是“企业内部系统里需要一个AI能力”业务逻辑和模型推理耦合很深并发量中等那Java方案完全够用。弄清楚这个边界后面所有的技术选型都不会走偏。2. 硬件规划篇64G内存加48G显卡到底能喂饱多大参数的模型在写任何代码之前先算账。推理部署最怕的是模型下载完了代码写完了结果一跑就显存溢出然后所有工作推翻重来。显存和内存的规划是第一步也是决定方案能不能落地的一步。2.1 LLaMA2各尺寸模型的显存账本LLaMA2有7B、13B、70B三个主要尺寸。一个模型加载到显存里占用空间主要由两部分组成模型权重和KV Cache。权重的计算公式很简单参数量乘以每个参数占用的字节数。以FP16精度为例7B模型7 * 10^9 * 2字节 ≈ 14GB13B模型13 * 10^9 * 2字节 ≈ 26GB70B模型70 * 10^9 * 2字节 ≈ 140GB这只是权重部分。KV Cache是推理过程中为每个请求缓存的历史Key-Value向量它的占用和请求的并发数、上下文长度有直接关系。即便是一个中等并发场景KV Cache吃掉几个GB显存也是常态。所以实际显存需求要比权重占用高出一截。模型尺寸精度权重占用预估总显存含KV Cache适合的GPU配置7BFP16~14GB~18-20GB单张24GB13BFP16~26GB~34-40GB单张48GB13BINT8~13GB~20-24GB单张24GB70BFP16~140GB~160GB至少4张48GB/80GB70BINT8~70GB~90GB2张48GB或4张24GB结合标题里的硬件配置来看64G内存 48G显卡或者多张合计48G显存跑7B是绰绰有余的跑13B FP16也基本够用但显存余量不会太大。如果想跑70B就必须走量化或者多卡张量并行。我在实测中比较推荐的组合是48G单卡跑13B FP16或者双24G卡用数据并行跑两个7B实例吞吐量会更理想。2.2 多GPU分区、拓扑与CUDA_VISIBLE_DEVICES的坑多GPU环境有个细节很多人容易忽略——物理拓扑。两张卡如果走PCIe总线通信带宽一般在16GB/s到32GB/s之间如果能走NVLink带宽能达到数百GB/s。如果你要做张量并行同一份模型切分到多张卡跨卡通信频繁PCIe会成为严重瓶颈导致加速比极低。但如果你做的是数据并行每张卡跑独立模型副本卡间通信很少PCIe完全够用。还有一个在实际中非常容易踩的坑CUDA_VISIBLE_DEVICES配置。不设置这个环境变量时CUDA驱动默认按物理编号给进程分配GPU一旦你屏蔽了某张卡后续所有卡号都会顺移。比如你设置CUDA_VISIBLE_DEVICES2,3程序里看到的设备编号就变成了0和1而不是2和3。这在多卡部署时特别容易让人糊涂。我现在的习惯是每个GPU进程单独配一个环境变量再在程序启动日志里打印当前可见的GPU数量和名称一旦发现和预期不符第一时间就能定位。至于GPU分区如果你用的是A100/A800这类支持MIG的卡可以做显存分区把一张卡切成多个实例。但MIG会限制显存和算力配额对LLaMA2这种动辄十几GB的模型来说分区收益不大。如果你只是想让多个进程分摊不同的卡用CUDA_VISIBLE_DEVICES做隔离就够了不需要动硬件层的分区。3. Java侧技术选型DJL、ONNX Runtime还是TensorRT的Java绑定硬件账算完之后就进入正题Java侧到底怎么加载和运行LLaMA2模型。这个环节的方案选择基本决定了整个项目的开发效率和上线后的稳定性值得花点篇幅好好聊聊。3.1 三条可行的技术路线对比我在调研阶段梳理了三条路线Deep Java LibraryDJL、ONNX Runtime Java API、以及自己封装TensorRT的JNI接口。三条路线思路完全不同优劣差异明显。DJL是AWS开源的一个Java深度学习框架定位是Java生态的“标准推理入口”。它像是一个适配层底层引擎可以切换成PyTorch、TensorFlow、ONNX Runtime甚至MXNet。对开发者来说只需要面向DJL的API写代码模型文件路径、输入输出格式、预处理逻辑统一由DJL管理。ONNX Runtime的Java API则另辟蹊径。它的核心思路是把模型先转换成ONNX格式然后通过ONNX Runtime的Java绑定直接加载推理。优势是ONNX格式非常通用很多大模型都能转换而且ONNX Runtime的CPU和GPU优化做得很好。还有一条更激进的路直接用TensorRT C API手写Java层JNI绑定。TensorRT是NVIDIA官方的推理加速神器延迟和吞吐表现都是顶级的但JNI封装的工作量极大而且你要熟悉CUDA内存管理和生命周期没有几周时间下不来。方案开发成本底层推理引擎性能表现社区活跃度DJL低PyTorch/ONNX/TensorRT中高高ONNX Runtime Java API中ONNX Runtime高中TensorRT JNI手写高TensorRT最高低3.2 我为什么选择了DJL加PyTorch引擎最终这个项目选了DJL加PyTorch引擎核心理由是两点。第一DJL对模型格式的兼容性好。LLaMA2开源权重虽然在HuggingFace上是PyTorch格式但通过DJL加载时不需要做复杂的格式转换。你用HuggingFace下载的PyTorch权重在DJL里直接指定模型路径就能跑。这种“零转换”能力对时间紧、任务重的项目相当关键。ONNX方案则需要先做torch.onnx.exportLLaMA2这种带KV Cache的动态图模型导出ONNX并不方便稍不注意就报一堆算子不支持的错误。第二DJL的并发和生命周期管理做得省心。Java应用里跑大模型推理最怕的是显存泄漏和CUDA context冲突。DJL内部对PyTorch的Native层做了内存管理封装模型加载、推理结束后的显存释放都有配套处理。相比之下如果自己用JNI管理稍微漏掉一个释放函数服务跑一晚上就可能把显存耗尽这个风险在线上是不能接受的。DJL的代码风格也比较清爽。加载模型只需要构建一个Criteria对象指定模型地址、输入输出类型和推理引擎即可CriteriaTextPrompt, GeneratedText criteria Criteria.builder() .optApplication(Application.NLP.TEXT_GENERATION) .setTypes(TextPrompt.class, GeneratedText.class) .optModelUrls(djl://ai.djl.pytorch/llama2-7b-chat) .optEngine(PyTorch) .optProgress(new ProgressBar()) .build(); ZooModelTextPrompt, GeneratedText model criteria.loadModel(); PredictorTextPrompt, GeneratedText predictor model.newPredictor(); GeneratedText result predictor.predict(new TextPrompt(《三国演义》的主要人物有)); System.out.println(result.getText());3.3 TensorRT的安装部署与Java的间接接轨虽然项目主用DJL加PyTorch引擎跑通了但我还是花时间把TensorRT环境搭了起来做了对比测试毕竟在GPU推理领域TensorRT的加速效果确实显著。TensorRT的安装是个相对繁琐的过程核心在于版本对齐TensorRT版本必须和你的CUDA版本、cuDNN版本严格匹配否则加载引擎时会直接报错。比如CUDA 11.8对应TensorRT 8.5和8.6CUDA 12.0对应TensorRT 8.6和9.0这个对应关系在NVIDIA官方文档里有明确矩阵我建议先查文档再安装不要凭感觉。在Java侧虽然TensorRT官方没有提供完整的Java API但可以通过DJL的TensorRT引擎间接调用。也就是说你先把PyTorch模型转成TensorRT的engine文件再用DJL加载这样既享受了TensorRT的加速能力又绕过了手写JNI的痛苦。4. 多GPU推理的核心实现从数据并行到张量并行模型选型定下来之后最核心的一个问题就浮出水面了怎么把多张GPU用起来。我在这个项目里同时尝试了数据并行和张量并行两条路最终跑了性能对比选出了适合业务场景的方案。4.1 数据并行最简单也最稳的多卡方案如果你的显存能够容纳单个模型副本数据并行是成本最低、收益最稳的方案。核心思想很直接每张GPU卡都加载一份完整的模型副本请求进来时按负载均衡策略分发到不同的卡上。在Java侧实现数据并行不需要在代码里做复杂的多卡通信。常见的做法是将模型加载和推理服务分成独立部署单元利用CUDA_VISIBLE_DEVICES让每个JVM进程绑定一张卡外层再用Nginx或者自研的router做请求分发。比如你有两台推理节点每台有一张24G显卡就启动两个Java进程# 节点1绑定GPU 0 export CUDA_VISIBLE_DEVICES0 java -Xmx16G -jar llama2-deploy.jar --server.port18080 # 节点2绑定GPU 1 export CUDA_VISIBLE_DEVICES1 java -Xmx16G -jar llama2-deploy.jar --server.port18081外层路由可以做一个简单的轮询或根据请求量动态分发。我实测下来这种方案的吞吐量基本是随GPU数量线性增长的两个实例就是接近两倍。而且因为各个进程互不干扰一个进程OOM不会影响另一个故障隔离性很好。如果你的需求是把多卡能力集成在一个进程里也有一种实现在代码中动态加载多个模型实例并把每个实例固定到指定的GPU上。DJL支持通过Criteria的optDevice选项指定设备编号。不过我个人不推荐一个进程内跑多个大模型因为JVM的GC和CUDA的显存分配混在一起排查问题时会比较复杂。4.2 张量并行显存不够时的突破方案当模型太大单卡放不下的时候数据并行就不灵了。你必须做模型切分把不同层或同一层内的不同部分放到多张卡上协同计算这就叫张量并行。张量并行的落地难度要比数据并行高一个量级。在Java生态里最现实的路径还是依赖底层的PyTorch引擎。DJL在加载模型时可以通过模型参数指定tensorParallel配置底层实际上是通过PyTorch的torch.distributed来协调跨卡通信的。比如你要把13B模型切到两张卡上可以这样配置MapString, String options new HashMap(); options.put(tensorParallel, 2); options.put(tensorParallelCount, 2); CriteriaTextPrompt, GeneratedText criteria Criteria.builder() .setTypes(TextPrompt.class, GeneratedText.class) .optModelUrls(/data/models/llama2-13b-chat) .optEngine(PyTorch) .optOptions(options) .build();这段逻辑的本质是PyTorch在加载模型时把各个层切分到两张卡上初始化的耗时比单卡加载稍长但推理时每一层的计算都在两张卡并行完成。从我的实测看张量并行能够明显降低单请求延迟但也引入了跨卡通信开销。如果两张卡走的是PCIe而不是NVLink通信延迟会让性能大打折扣。在PCIe环境下13B模型的张量并行我测下来加速比只有1.3倍左右远低于数据并行的2倍而NVLink环境下能达到1.7倍。所以我的建议比较明确只要显存够优先用数据并行只有单卡装不下模型、又不想用量化损失精度的情况下才考虑张量并行。4.3 KV Cache的并发管理这是多卡场景下很容易翻车的点还有一个必须单独提的点是KV Cache的并发管理。LLM推理和传统DNN推理不一样它不仅依赖当前输入还要维护之前所有token的KV Cache。请求并发高时KV Cache的显存占用会快速上涨。如果你的多卡方案用的是张量并行KV Cache还会分散在多张卡上任何一张卡的显存爆掉整次推理都会失败。我在项目中通过两个手段控制这个风险。第一是限定单请求的最大生成长度从源头限制KV Cache的增长上限第二是给推理服务加并发控制信号量限制同时在推理的请求数。宁可让请求排队也不要让显存被突发流量打穿。5. 实测数据与排坑实录OOM、环境变量与JVM内存那些事这部分我打算写点真正来自一线的东西。理论说得再好不如把实际踩过的坑和测过的数据摆出来。下面三个子环节每一个都是我在项目推进中真实遇到并花时间解决的。5.1 一次典型OOM的完整排查链路项目联调阶段服务刚启动是好的跑了几百个请求之后突然开始频繁抛java.lang.OutOfMemoryError: Insufficient memory。当时第一反应是JVM堆内存不够于是把-Xmx从16G调到了24G结果重启后问题依旧。这时候我才意识到事情没那么简单。我把错误堆栈重新看了一遍发现栈顶信息指向的是Deep Learning框架的Native层调用。这时候判断问题不是堆内存而是Native内存或者显存。于是我去查JVM的Direct Memory配置发现默认的MaxDirectMemorySize是-Xmx的值理论上24G应该够用于是进一步怀疑是显存问题。最终我用nvidia-smi反复观测发现每处理一个请求显存占用都会涨几百MB而且请求结束后不会回落。确诊了是模型的KV Cache或者推理过程中的临时张量没有被显式释放。DJL对这类问题有自己的缓冲池机制但需要你在代码里主动调用predictor.close()或者复用Predictor实例而不是每次请求都new一个。我当时的代码恰好踩了这个坑——每次都新建Predictor用完没有关闭导致Native层的显存上下文一直在累积。修正代码结构之后显存占用曲线变得非常平稳。这个坑的教训是排查Java侧的OOM不能只盯JVM堆。要让应用启动时先打印实际可见的GPU数量和显存总量并在请求处理完打一条日志记录当前显存占用才能快速区分是“真堆溢出”还是“显存泄漏”。5.2 Java环境变量配置与CUDA版本的匹配问题如果把项目交给其他同事跑十个里面有八个会卡在环境配置上。这个项目涉及的底层组件多JDK、CUDA、cuDNN、PyTorch Native库、TensorRT之间都存在版本联动任何一个对不上启动时就会报各种native library加载失败。先说Java环境变量。基础要求是JAVA_HOME配到JDK根目录PATH里加上bin目录。这个看似简单但最容易出现的问题是机器上有多个JDK系统默认走了旧版本。LLaMA2推理底层PyTorch Native库对JDK版本有要求我这边用的是JDK 17。如果默认JDK是8加载PyTorch库时经常报UnsupportedClassVersionError而这个问题通常只在运行阶段暴露白等编译半天才发现。CUDA和PyTorch的版本匹配也值得留意。DJL的PyTorch引擎在djl.properties里有默认绑定的CUDA版本比如DJL 0.26系列默认对应CUDA 11.8。如果你本机装的是CUDA 12.0以上直接用默认配置跑会提示找不到libcudart.so。解决方法是添加CUDA依赖参数指定DJL去加载对应CUDA版本的引擎dependency groupIdai.djl.pytorch/groupId artifactIdpytorch-engine/artifactId version0.26.0/version /dependency dependency groupIdai.djl.pytorch/groupId artifactIdpytorch-native-cu118/artifactId version2.0.1/version classifierlinux-x86_64/classifier /dependency经验之谈在写代码之前先把CUDA版本、PyTorch版本、DJL版本三个检查好记到README里。这一项能省下大量的环境排查时间。5.3 性能实测单卡与多卡的真实差距最后放一组我在项目里实测的数据。测试模型是LLaMA2-13B-Chat输入长度256 token输出长度128 token单请求压测和并发压测各跑了一轮。单卡48G单请求平均耗时约520ms双卡数据并行两个7B实例单请求平均耗时约280ms但吞吐量几乎翻倍双卡张量并行跑13B单请求平均耗时约390ms比单卡48G快25%左右但没有达到2倍。场景单卡 48G (13B)双卡数据并行 (2x7B)双卡张量并行 (13B)单请求延迟520ms280ms390ms并发8路吞吐约8 req/s约15 req/s约11 req/s显存占用峰值38GB18GB x222GB x2这个结果验证了前面的判断数据并行在吞吐冲刺上优势明显张量并行在单请求延迟上表现更好。如果你面对的是大量独立请求建议数据并行如果面对的是单个复杂任务的长上下文处理张量并行会更划算。另外还要提醒一点如果你是用Java跑推理别忽视JVM的GC对延迟的影响。我在压测时发现服务跑久了之后偶尔会有个别请求延迟飙到两秒以上排查发现是GC暂停。后来给JVM配置了G1垃圾回收器并设置了目标暂停时间延迟抖动明显减少。对于在线推理服务建议在启动参数中加上java -Xms16G -Xmx16G -XX:UseG1GC -XX:MaxGCPauseMillis200 -jar llama2-deploy.jar一套做下来Java多GPU跑LLaMA2这个方案在我这边算是稳定落地了。从最初被质疑“Java怎么可能跑大模型”到最终承接每日数千次的线上调用整个过程中我最大的体会是技术选型没有绝对的优劣关键在于匹配场景。Java生态在大模型推理上确实不如Python灵活但只要你愿意花时间做显存规划、选对框架、控制好并发它一样能成为一个可靠的大模型服务底座。如果你正在这个方向上摸索希望这篇复盘能帮你少踩几个坑。本文还有配套的精品资源点击获取 Java+多GPU部署LLaMA2推理服务实战复盘与性能调优
简介大模型推理部署通常被Python生态主导但在企业内网隔离、后端以Spring Boot为主且缺乏Python运维能力的场景下Java直接承载LLaMA2推理反而能显著降低系统复杂度。本文从显存规划与技术选型切入对比DJL、ONNX Runtime与TensorRT Java绑定三条路线说明如何基于DJLPyTorch引擎加载LLaMA2模型并深入解析数据并行与张量并行的多GPU实现策略。同时结合KV Cache管理、CUDA_VISIBLE_DEVICES配置、JVM内存与GC调优等实战经验给出单卡与多卡性能实测数据为Java工程师在现有基础设施中集成大模型推理能力提供完整参考。 先说结论用Java做LLaMA2的推理部署乍看像是一个伪需求。毕竟LLaMA2的官方实现、HuggingFace生态、各种优化加速库基本都是Python的天下硬要往Java上靠总给人一种脱裤子放屁的感觉。但如果你所在的企业后端是清一色的Spring Boot机房里有几块闲置的NVIDIA GPU而业务方又希望在一个内网隔离系统里直接调用大模型的能力时Java多GPU这套组合反而是性价比最高的路。这篇文章是我把一个基于Java多GPU的LLaMA2推理服务从零推到线上的完整复盘包含了算账、选型、编码、排障各阶段的细节希望对正在纠结这套方案的你有参考价值。1. 为什么要把LLaMA2塞进Java进程先聊聊这个方案的适用边界很多人一听到Java跑LLaMA2第一反应都是“何必呢”。确实如果目标只是跑通一个DemoPython一条命令就能启动服务Java连渣都赶不上。但在实际生产环境里事情往往不是“哪个语言跑模型最方便”而是“在现有系统架构下哪个方案的综合成本最低”。1.1 哪些场景下Java推理是刚需而不是炫技我接手这个项目时业务方已经有一套基于Spring Boot的文档审核系统每天要处理大量内部文档。那边的诉求是在文档上传后能自动调用大模型做摘要和关键词提取。看似简单但有两个硬约束第一服务必须部署在内网和公网完全隔离不能调用云厂商API第二团队里只有Java工程师没有人愿意为了一个摘要功能额外维护一套Python微服务。这种情况下Java直接调用LLaMA2就成了最优解。因为业务代码和推理逻辑可以打包在同一个JVM进程里走一次http调用就能拿到结果不需要跨语言、跨服务。系统监控、日志、链路追踪也能复用原有的基础设施不需要为Python单独搭建一套。对于运维来说一个Jar包丢到机器上就能跑远比“再维护一个Python环境、再装一堆pip依赖”要好接受。1.2 什么场景下不该用Java硬扛这里也要泼一盆冷水。如果你的目标是做一个高并发的LLM在线服务平台比如面向C端的聊天机器人要支撑几百路的并发推理那Java方案大概率不是最优选择。vLLM、TensorRT-LLM、Triton Inference Server这些专门针对大模型优化过的Python/C方案在吞吐量和显存利用效率上要领先Java生态一大截。我的建议是如果你的核心业务就是“卖模型推理能力”别碰Java用专业推理框架如果你的核心业务是“企业内部系统里需要一个AI能力”业务逻辑和模型推理耦合很深并发量中等那Java方案完全够用。弄清楚这个边界后面所有的技术选型都不会走偏。2. 硬件规划篇64G内存加48G显卡到底能喂饱多大参数的模型在写任何代码之前先算账。推理部署最怕的是模型下载完了代码写完了结果一跑就显存溢出然后所有工作推翻重来。显存和内存的规划是第一步也是决定方案能不能落地的一步。2.1 LLaMA2各尺寸模型的显存账本LLaMA2有7B、13B、70B三个主要尺寸。一个模型加载到显存里占用空间主要由两部分组成模型权重和KV Cache。权重的计算公式很简单参数量乘以每个参数占用的字节数。以FP16精度为例7B模型7 * 10^9 * 2字节 ≈ 14GB13B模型13 * 10^9 * 2字节 ≈ 26GB70B模型70 * 10^9 * 2字节 ≈ 140GB这只是权重部分。KV Cache是推理过程中为每个请求缓存的历史Key-Value向量它的占用和请求的并发数、上下文长度有直接关系。即便是一个中等并发场景KV Cache吃掉几个GB显存也是常态。所以实际显存需求要比权重占用高出一截。模型尺寸精度权重占用预估总显存含KV Cache适合的GPU配置7BFP16~14GB~18-20GB单张24GB13BFP16~26GB~34-40GB单张48GB13BINT8~13GB~20-24GB单张24GB70BFP16~140GB~160GB至少4张48GB/80GB70BINT8~70GB~90GB2张48GB或4张24GB结合标题里的硬件配置来看64G内存 48G显卡或者多张合计48G显存跑7B是绰绰有余的跑13B FP16也基本够用但显存余量不会太大。如果想跑70B就必须走量化或者多卡张量并行。我在实测中比较推荐的组合是48G单卡跑13B FP16或者双24G卡用数据并行跑两个7B实例吞吐量会更理想。2.2 多GPU分区、拓扑与CUDA_VISIBLE_DEVICES的坑多GPU环境有个细节很多人容易忽略——物理拓扑。两张卡如果走PCIe总线通信带宽一般在16GB/s到32GB/s之间如果能走NVLink带宽能达到数百GB/s。如果你要做张量并行同一份模型切分到多张卡跨卡通信频繁PCIe会成为严重瓶颈导致加速比极低。但如果你做的是数据并行每张卡跑独立模型副本卡间通信很少PCIe完全够用。还有一个在实际中非常容易踩的坑CUDA_VISIBLE_DEVICES配置。不设置这个环境变量时CUDA驱动默认按物理编号给进程分配GPU一旦你屏蔽了某张卡后续所有卡号都会顺移。比如你设置CUDA_VISIBLE_DEVICES2,3程序里看到的设备编号就变成了0和1而不是2和3。这在多卡部署时特别容易让人糊涂。我现在的习惯是每个GPU进程单独配一个环境变量再在程序启动日志里打印当前可见的GPU数量和名称一旦发现和预期不符第一时间就能定位。至于GPU分区如果你用的是A100/A800这类支持MIG的卡可以做显存分区把一张卡切成多个实例。但MIG会限制显存和算力配额对LLaMA2这种动辄十几GB的模型来说分区收益不大。如果你只是想让多个进程分摊不同的卡用CUDA_VISIBLE_DEVICES做隔离就够了不需要动硬件层的分区。3. Java侧技术选型DJL、ONNX Runtime还是TensorRT的Java绑定硬件账算完之后就进入正题Java侧到底怎么加载和运行LLaMA2模型。这个环节的方案选择基本决定了整个项目的开发效率和上线后的稳定性值得花点篇幅好好聊聊。3.1 三条可行的技术路线对比我在调研阶段梳理了三条路线Deep Java LibraryDJL、ONNX Runtime Java API、以及自己封装TensorRT的JNI接口。三条路线思路完全不同优劣差异明显。DJL是AWS开源的一个Java深度学习框架定位是Java生态的“标准推理入口”。它像是一个适配层底层引擎可以切换成PyTorch、TensorFlow、ONNX Runtime甚至MXNet。对开发者来说只需要面向DJL的API写代码模型文件路径、输入输出格式、预处理逻辑统一由DJL管理。ONNX Runtime的Java API则另辟蹊径。它的核心思路是把模型先转换成ONNX格式然后通过ONNX Runtime的Java绑定直接加载推理。优势是ONNX格式非常通用很多大模型都能转换而且ONNX Runtime的CPU和GPU优化做得很好。还有一条更激进的路直接用TensorRT C API手写Java层JNI绑定。TensorRT是NVIDIA官方的推理加速神器延迟和吞吐表现都是顶级的但JNI封装的工作量极大而且你要熟悉CUDA内存管理和生命周期没有几周时间下不来。方案开发成本底层推理引擎性能表现社区活跃度DJL低PyTorch/ONNX/TensorRT中高高ONNX Runtime Java API中ONNX Runtime高中TensorRT JNI手写高TensorRT最高低3.2 我为什么选择了DJL加PyTorch引擎最终这个项目选了DJL加PyTorch引擎核心理由是两点。第一DJL对模型格式的兼容性好。LLaMA2开源权重虽然在HuggingFace上是PyTorch格式但通过DJL加载时不需要做复杂的格式转换。你用HuggingFace下载的PyTorch权重在DJL里直接指定模型路径就能跑。这种“零转换”能力对时间紧、任务重的项目相当关键。ONNX方案则需要先做torch.onnx.exportLLaMA2这种带KV Cache的动态图模型导出ONNX并不方便稍不注意就报一堆算子不支持的错误。第二DJL的并发和生命周期管理做得省心。Java应用里跑大模型推理最怕的是显存泄漏和CUDA context冲突。DJL内部对PyTorch的Native层做了内存管理封装模型加载、推理结束后的显存释放都有配套处理。相比之下如果自己用JNI管理稍微漏掉一个释放函数服务跑一晚上就可能把显存耗尽这个风险在线上是不能接受的。DJL的代码风格也比较清爽。加载模型只需要构建一个Criteria对象指定模型地址、输入输出类型和推理引擎即可CriteriaTextPrompt, GeneratedText criteria Criteria.builder() .optApplication(Application.NLP.TEXT_GENERATION) .setTypes(TextPrompt.class, GeneratedText.class) .optModelUrls(djl://ai.djl.pytorch/llama2-7b-chat) .optEngine(PyTorch) .optProgress(new ProgressBar()) .build(); ZooModelTextPrompt, GeneratedText model criteria.loadModel(); PredictorTextPrompt, GeneratedText predictor model.newPredictor(); GeneratedText result predictor.predict(new TextPrompt(《三国演义》的主要人物有)); System.out.println(result.getText());3.3 TensorRT的安装部署与Java的间接接轨虽然项目主用DJL加PyTorch引擎跑通了但我还是花时间把TensorRT环境搭了起来做了对比测试毕竟在GPU推理领域TensorRT的加速效果确实显著。TensorRT的安装是个相对繁琐的过程核心在于版本对齐TensorRT版本必须和你的CUDA版本、cuDNN版本严格匹配否则加载引擎时会直接报错。比如CUDA 11.8对应TensorRT 8.5和8.6CUDA 12.0对应TensorRT 8.6和9.0这个对应关系在NVIDIA官方文档里有明确矩阵我建议先查文档再安装不要凭感觉。在Java侧虽然TensorRT官方没有提供完整的Java API但可以通过DJL的TensorRT引擎间接调用。也就是说你先把PyTorch模型转成TensorRT的engine文件再用DJL加载这样既享受了TensorRT的加速能力又绕过了手写JNI的痛苦。4. 多GPU推理的核心实现从数据并行到张量并行模型选型定下来之后最核心的一个问题就浮出水面了怎么把多张GPU用起来。我在这个项目里同时尝试了数据并行和张量并行两条路最终跑了性能对比选出了适合业务场景的方案。4.1 数据并行最简单也最稳的多卡方案如果你的显存能够容纳单个模型副本数据并行是成本最低、收益最稳的方案。核心思想很直接每张GPU卡都加载一份完整的模型副本请求进来时按负载均衡策略分发到不同的卡上。在Java侧实现数据并行不需要在代码里做复杂的多卡通信。常见的做法是将模型加载和推理服务分成独立部署单元利用CUDA_VISIBLE_DEVICES让每个JVM进程绑定一张卡外层再用Nginx或者自研的router做请求分发。比如你有两台推理节点每台有一张24G显卡就启动两个Java进程# 节点1绑定GPU 0 export CUDA_VISIBLE_DEVICES0 java -Xmx16G -jar llama2-deploy.jar --server.port18080 # 节点2绑定GPU 1 export CUDA_VISIBLE_DEVICES1 java -Xmx16G -jar llama2-deploy.jar --server.port18081外层路由可以做一个简单的轮询或根据请求量动态分发。我实测下来这种方案的吞吐量基本是随GPU数量线性增长的两个实例就是接近两倍。而且因为各个进程互不干扰一个进程OOM不会影响另一个故障隔离性很好。如果你的需求是把多卡能力集成在一个进程里也有一种实现在代码中动态加载多个模型实例并把每个实例固定到指定的GPU上。DJL支持通过Criteria的optDevice选项指定设备编号。不过我个人不推荐一个进程内跑多个大模型因为JVM的GC和CUDA的显存分配混在一起排查问题时会比较复杂。4.2 张量并行显存不够时的突破方案当模型太大单卡放不下的时候数据并行就不灵了。你必须做模型切分把不同层或同一层内的不同部分放到多张卡上协同计算这就叫张量并行。张量并行的落地难度要比数据并行高一个量级。在Java生态里最现实的路径还是依赖底层的PyTorch引擎。DJL在加载模型时可以通过模型参数指定tensorParallel配置底层实际上是通过PyTorch的torch.distributed来协调跨卡通信的。比如你要把13B模型切到两张卡上可以这样配置MapString, String options new HashMap(); options.put(tensorParallel, 2); options.put(tensorParallelCount, 2); CriteriaTextPrompt, GeneratedText criteria Criteria.builder() .setTypes(TextPrompt.class, GeneratedText.class) .optModelUrls(/data/models/llama2-13b-chat) .optEngine(PyTorch) .optOptions(options) .build();这段逻辑的本质是PyTorch在加载模型时把各个层切分到两张卡上初始化的耗时比单卡加载稍长但推理时每一层的计算都在两张卡并行完成。从我的实测看张量并行能够明显降低单请求延迟但也引入了跨卡通信开销。如果两张卡走的是PCIe而不是NVLink通信延迟会让性能大打折扣。在PCIe环境下13B模型的张量并行我测下来加速比只有1.3倍左右远低于数据并行的2倍而NVLink环境下能达到1.7倍。所以我的建议比较明确只要显存够优先用数据并行只有单卡装不下模型、又不想用量化损失精度的情况下才考虑张量并行。4.3 KV Cache的并发管理这是多卡场景下很容易翻车的点还有一个必须单独提的点是KV Cache的并发管理。LLM推理和传统DNN推理不一样它不仅依赖当前输入还要维护之前所有token的KV Cache。请求并发高时KV Cache的显存占用会快速上涨。如果你的多卡方案用的是张量并行KV Cache还会分散在多张卡上任何一张卡的显存爆掉整次推理都会失败。我在项目中通过两个手段控制这个风险。第一是限定单请求的最大生成长度从源头限制KV Cache的增长上限第二是给推理服务加并发控制信号量限制同时在推理的请求数。宁可让请求排队也不要让显存被突发流量打穿。5. 实测数据与排坑实录OOM、环境变量与JVM内存那些事这部分我打算写点真正来自一线的东西。理论说得再好不如把实际踩过的坑和测过的数据摆出来。下面三个子环节每一个都是我在项目推进中真实遇到并花时间解决的。5.1 一次典型OOM的完整排查链路项目联调阶段服务刚启动是好的跑了几百个请求之后突然开始频繁抛java.lang.OutOfMemoryError: Insufficient memory。当时第一反应是JVM堆内存不够于是把-Xmx从16G调到了24G结果重启后问题依旧。这时候我才意识到事情没那么简单。我把错误堆栈重新看了一遍发现栈顶信息指向的是Deep Learning框架的Native层调用。这时候判断问题不是堆内存而是Native内存或者显存。于是我去查JVM的Direct Memory配置发现默认的MaxDirectMemorySize是-Xmx的值理论上24G应该够用于是进一步怀疑是显存问题。最终我用nvidia-smi反复观测发现每处理一个请求显存占用都会涨几百MB而且请求结束后不会回落。确诊了是模型的KV Cache或者推理过程中的临时张量没有被显式释放。DJL对这类问题有自己的缓冲池机制但需要你在代码里主动调用predictor.close()或者复用Predictor实例而不是每次请求都new一个。我当时的代码恰好踩了这个坑——每次都新建Predictor用完没有关闭导致Native层的显存上下文一直在累积。修正代码结构之后显存占用曲线变得非常平稳。这个坑的教训是排查Java侧的OOM不能只盯JVM堆。要让应用启动时先打印实际可见的GPU数量和显存总量并在请求处理完打一条日志记录当前显存占用才能快速区分是“真堆溢出”还是“显存泄漏”。5.2 Java环境变量配置与CUDA版本的匹配问题如果把项目交给其他同事跑十个里面有八个会卡在环境配置上。这个项目涉及的底层组件多JDK、CUDA、cuDNN、PyTorch Native库、TensorRT之间都存在版本联动任何一个对不上启动时就会报各种native library加载失败。先说Java环境变量。基础要求是JAVA_HOME配到JDK根目录PATH里加上bin目录。这个看似简单但最容易出现的问题是机器上有多个JDK系统默认走了旧版本。LLaMA2推理底层PyTorch Native库对JDK版本有要求我这边用的是JDK 17。如果默认JDK是8加载PyTorch库时经常报UnsupportedClassVersionError而这个问题通常只在运行阶段暴露白等编译半天才发现。CUDA和PyTorch的版本匹配也值得留意。DJL的PyTorch引擎在djl.properties里有默认绑定的CUDA版本比如DJL 0.26系列默认对应CUDA 11.8。如果你本机装的是CUDA 12.0以上直接用默认配置跑会提示找不到libcudart.so。解决方法是添加CUDA依赖参数指定DJL去加载对应CUDA版本的引擎dependency groupIdai.djl.pytorch/groupId artifactIdpytorch-engine/artifactId version0.26.0/version /dependency dependency groupIdai.djl.pytorch/groupId artifactIdpytorch-native-cu118/artifactId version2.0.1/version classifierlinux-x86_64/classifier /dependency经验之谈在写代码之前先把CUDA版本、PyTorch版本、DJL版本三个检查好记到README里。这一项能省下大量的环境排查时间。5.3 性能实测单卡与多卡的真实差距最后放一组我在项目里实测的数据。测试模型是LLaMA2-13B-Chat输入长度256 token输出长度128 token单请求压测和并发压测各跑了一轮。单卡48G单请求平均耗时约520ms双卡数据并行两个7B实例单请求平均耗时约280ms但吞吐量几乎翻倍双卡张量并行跑13B单请求平均耗时约390ms比单卡48G快25%左右但没有达到2倍。场景单卡 48G (13B)双卡数据并行 (2x7B)双卡张量并行 (13B)单请求延迟520ms280ms390ms并发8路吞吐约8 req/s约15 req/s约11 req/s显存占用峰值38GB18GB x222GB x2这个结果验证了前面的判断数据并行在吞吐冲刺上优势明显张量并行在单请求延迟上表现更好。如果你面对的是大量独立请求建议数据并行如果面对的是单个复杂任务的长上下文处理张量并行会更划算。另外还要提醒一点如果你是用Java跑推理别忽视JVM的GC对延迟的影响。我在压测时发现服务跑久了之后偶尔会有个别请求延迟飙到两秒以上排查发现是GC暂停。后来给JVM配置了G1垃圾回收器并设置了目标暂停时间延迟抖动明显减少。对于在线推理服务建议在启动参数中加上java -Xms16G -Xmx16G -XX:UseG1GC -XX:MaxGCPauseMillis200 -jar llama2-deploy.jar一套做下来Java多GPU跑LLaMA2这个方案在我这边算是稳定落地了。从最初被质疑“Java怎么可能跑大模型”到最终承接每日数千次的线上调用整个过程中我最大的体会是技术选型没有绝对的优劣关键在于匹配场景。Java生态在大模型推理上确实不如Python灵活但只要你愿意花时间做显存规划、选对框架、控制好并发它一样能成为一个可靠的大模型服务底座。如果你正在这个方向上摸索希望这篇复盘能帮你少踩几个坑。本文还有配套的精品资源点击获取