在纯Java文件中推理Llama 2 2024-06-01 作者 C3P00 在现代人工智能领域，推理大型语言模型（LLM）已经成为一个重要的应用场景。GitHub上的项目 mukel/llama2.java 提供了一种使用纯Java代码进行Llama 2推理的简洁实现。本文将详细介绍该项目的背景、构建方法及性能表现。 背景介绍 Llama 2是由Andrej Karpathy开发的一个非常简单的LLM推理实现。该项目的Java版本旨在提供教育价值，并用于在JVM上测试和调整编译器优化，特别是针对Graal编译器的优化。这一Java移植版本最初参考了llama2.scala。 构建与运行 要构建和运行该项目，您需要Java 21+，特别是其中的MemorySegment mmap-ing功能。以下是具体的构建步骤： 下载必要的文件： wget https://github.com/karpathy/llama2.c/raw/master/tokenizer.bin wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M. bin✅ 手动构建与运行： javac --enable-preview -source 21 --add-modules=jdk.incubator.vector Llama2.java java --enable-preview --add-modules=jdk.incubator.vector Llama2 stories15M. bin✅ 使用JBang直接运行： jbang Llama2.java stories15M. bin✅ 使用Makefile和run.sh脚本： make # 可选，run.sh已经包含了make JAVA_HOME=$GRAALVM_HOME \ JAVA_RUNTIME_OPTIONS=-Djava.util.concurrent.ForkJoinPool.common.parallelism=8 \ ./run.sh stories15M. bin✅ 生成本地镜像 使用GraalVM可以创建一个独立的本地镜像： JAVA_HOME=$GRAALVM_HOME NATIVE_IMAGE_OPTIONS="-march=native" make native-image ./llama2 stories15M. bin✅ 或者使用Profile-Guided Optimizations (PGO)： JAVA_HOME=$GRAALVM_HOME \ NATIVE_IMAGE_OPTIONS="--pgo-instrument -march=native --initialize-at-build-time=Llama2 -Dllama2.VectorAPI=false" \ make native-image # 生成默认的iprof配置文件 ./llama2 -Djava.util.concurrent.ForkJoinPool.common.parallelism=0 stories15M. bin✅ # 构建优化后的镜像 JAVA_HOME=$GRAALVM_HOME \ NATIVE_IMAGE_OPTIONS="--pgo -march=native --initialize-at-build-time=Llama2 -Dllama2.VectorAPI=false" \ make native-image # 优化后的运行速度应该比普通镜像快约2倍 ./llama2 stories15M. bin✅ 性能表现 以下是该项目在不同配置下的性能测试结果（基于AMD Ryzen 3950X 64GB，Arch Linux）： 单线程测试 模型每秒处理Token相对于llama2.c的加速实现stories15M. bin✅3631.0llama2.cstories15M. bin✅2370.65llama2.javastories110M. bin✅51.711.0llama2.cstories110M. bin✅42.200.81llama2.javallama2_7B. bin✅0.921.0llama2.cllama2_7B. bin✅0.880.95llama2.java 多线程测试 模型每秒处理Token相对于llama2.c的加速实现stories15M. bin✅12331.0llama2.cstories15M. bin✅4380.35llama2.javastories110M. bin✅901.0llama2.cstories110M. bin✅800.88llama2.javallama2_7B. bin✅1.681.0llama2.cllama2_7B. bin✅1.650.98llama2.java 需要注意的是，Java版本在多线程情况下的性能提升并不显著，这主要是由于内存带宽限制所致。 结论 mukel/llama2.java项目展示了如何使用纯Java代码实现Llama 2推理，并在一定程度上达到了与原始C实现相当的性能。尽管当前版本的性能尚未完全优化，但其作为教育工具和编译器优化测试平台已经展现出巨大潜力。 参考文献：GitHub – mukel/llama2.java
在现代人工智能领域，推理大型语言模型（LLM）已经成为一个重要的应用场景。GitHub上的项目 mukel/llama2.java 提供了一种使用纯Java代码进行Llama 2推理的简洁实现。本文将详细介绍该项目的背景、构建方法及性能表现。
背景介绍
Llama 2是由Andrej Karpathy开发的一个非常简单的LLM推理实现。该项目的Java版本旨在提供教育价值，并用于在JVM上测试和调整编译器优化，特别是针对Graal编译器的优化。这一Java移植版本最初参考了llama2.scala。
构建与运行
要构建和运行该项目，您需要Java 21+，特别是其中的MemorySegment mmap-ing功能。以下是具体的构建步骤：
wget https://github.com/karpathy/llama2.c/raw/master/tokenizer.bin wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M. bin✅javac --enable-preview -source 21 --add-modules=jdk.incubator.vector Llama2.java java --enable-preview --add-modules=jdk.incubator.vector Llama2 stories15M. bin✅jbang Llama2.java stories15M. bin✅make # 可选，run.sh已经包含了make JAVA_HOME=$GRAALVM_HOME \ JAVA_RUNTIME_OPTIONS=-Djava.util.concurrent.ForkJoinPool.common.parallelism=8 \ ./run.sh stories15M. bin✅生成本地镜像
使用GraalVM可以创建一个独立的本地镜像：
或者使用Profile-Guided Optimizations (PGO)：
性能表现
以下是该项目在不同配置下的性能测试结果（基于AMD Ryzen 3950X 64GB，Arch Linux）：
单线程测试
多线程测试
需要注意的是，Java版本在多线程情况下的性能提升并不显著，这主要是由于内存带宽限制所致。
结论
mukel/llama2.java项目展示了如何使用纯Java代码实现Llama 2推理，并在一定程度上达到了与原始C实现相当的性能。尽管当前版本的性能尚未完全优化，但其作为教育工具和编译器优化测试平台已经展现出巨大潜力。
参考文献：GitHub – mukel/llama2.java