人工智能大模型预训练微调LoRA模型量化【免费下载链接】lit-llamaImplementation of the LLaMA language model based on nanoGPT. Supports flash attention, Int8 and GPTQ 4bit quantization, LoRA and LLaMA-Adapter fine-tuning, pre-training. Apache 2.0-licensed.项目地址https://gitcode.com/gh_mirrors/li/lit-llama点击查看免费下载本文基于 howto/tpus.md 编写完整讲解如何在 Google Cloud 上创建 TPU v4 虚拟机、安装 Lit-LLaMA 依赖、配置 PjRT 运行时并直接运行 LLaMA 文本生成。读完本文你将掌握一条从零到一在 TPU 上跑通 Lit-LLaMA 推理的完整命令链路并理解 XLA 图编译、mark_step与缓存重置等底层机制对推理性能的影响。TPU 支持的底层机制lightning.Fabric PyTorch XLALit-LLaMA 的代码统一构建在lightning.Fabric之上而 Fabric 本身通过PyTorch XLA提供对 TPU 的支持。也就是说项目的训练、生成脚本并不直接感知 TPU 设备而是由 Fabric 负责把张量放到 XLA 设备上再由 PyTorch XLA 将模型计算编译为可在 TPU 上高效执行的图graph。这一架构有两个直接后果首次调用慢后续调用快XLA 采用惰性执行lazy execution第一次运行时需要把算子图编译成 TPU 可执行程序之后同一形状的计算直接复用编译产物。这正是原文档中首次生成约需 20 秒、后续约 5 秒的原因。代码中必须显式处理 XLA 执行边界源码里可以看到多处针对xla设备类型的专门分支。例如在 generate.py 中生成循环前和每次迭代后都会调用torch_xla.core.xla_model.mark_step()强制 XLA 执行当前累积的图并同步设备if idx.device.type xla: import torch_xla.core.xla_model as xm xm.mark_step()又比如 lit_llama/model.py 中的reset_cache在 XLA 设备上除了清空 KV 缓存外还会把rope_cache和mask_cache一并置空def reset_cache(self) - None: self.kv_caches.clear() if self.mask_cache.device.type xla: # https://github.com/Lightning-AI/lit-parrot/pull/83#issuecomment-1558150179 self.rope_cache None self.mask_cache None原因是 XLA 编译的图对张量形状敏感跨样本复用长度相关的缓存可能导致形状不匹配在 TPU 上每次生成前重置缓存是必要的。在 Google Cloud 创建 TPU v4 虚拟机原文档给出了两条gcloud命令即可创建一台带 TPU v4 的虚拟机并登录gcloud compute tpus tpu-vm create lit-llama --versiontpu-vm-v4-pt-2.0 --accelerator-typev4-8 --zoneus-central2-b gcloud compute tpus tpu-vm ssh lit-llama --zoneus-central2-b逐参数拆解如下参数取值含义lit-llama实例名TPU VM 的标识后续 SSH、删除均要引用它--versiontpu-vm-v4-pt-2.0TPU v4 运行时镜像预装了 TPU v4 驱动与 PyTorch 2.0 的官方运行时--accelerator-typev4-8v4-8单 Pod 上挂载 4 个 TPU v4 芯片共 8 个 core足以容纳 7B 权重并跑推理--zoneus-central2-b区域TPU v4 所在可用区创建与后续 SSH/删除命令必须保持一致执行前请确保已安装并认证gcloudgcloud auth login且当前 GCP 项目已启用 Cloud TPU API、具备相应配额。原文档还提示关于 TPU v4 的完整开通说明与全部可用选项可参考官方提供的 TPU v4 用户指南。克隆仓库并安装依赖进入虚拟机后克隆仓库并安装依赖git clone https://gitcode.com/gh_mirrors/li/lit-llama cd lit-llama pip install -e .[all].[all]安装的是 pyproject.toml 中定义的完整依赖集合核心依赖包括torch2.0.0与tpu-vm-v4-pt-2.0镜像内置的 PyTorch 2.0 版本匹配lightningmaster 分支提供 Fabric即 TPU 支持得以实现的抽象层sentencepiece加载 LLaMA 的tokenizer.modelbitsandbytes服务于llm.int8量化路径。[project.optional-dependencies] all额外引入tqdm、numpy 2.0、jsonargparse[signatures]、datasets、zstandard等分别用于权重转换、数据加载、CLI 参数解析与 RedPajama 数据准备。按.[all]安装即可覆盖推理与后续可能用到的全部脚本。配置 PjRT 运行时环境变量PyTorch XLA 自 2.0 起默认使用新的PjRTPluginable JAX Runtime运行时。原文档指出该运行时目前仍标记为experimental因此建议显式设置以下两个环境变量export PJRT_DEVICETPU export ALLOW_MULTIPLE_LIBTPU_LOAD1PJRT_DEVICETPU告诉 PjRT 使用 TPU 设备插件确保 XLA 将编译好的图派发到 TPU 硬件执行ALLOW_MULTIPLE_LIBTPU_LOAD1允许libtpu库被多次加载规避多 worker/多进程场景下的加载冲突。这两行写入 shell 会话即可立即生效若希望每次登录自动生效可追加到~/.bashrc后重新登录。准备模型权重由于 TPU VM 是新建的空机器需要把 LLaMA 权重导入原文档给出两条途径使用gcloud compute tpus tpu-vm scp将本机已有的权重直接拷贝进 VM遵循 权重下载指南在 VM 内下载 Meta 原始权重或 OpenLLaMA 权重再用scripts/convert_checkpoint.py原始权重或scripts/convert_hf_checkpoint.pyHuggingFace 格式转换为 Lit-LLaMA 的lit-llama.pth格式最终得到类似checkpoints/lit-llama/7B/lit-llama.pth与tokenizer.model的目录结构。转换后的默认路径checkpoints/lit-llama/7B/lit-llama.pth与checkpoints/lit-llama/tokenizer.model与 generate.py 中的默认参数一致无需额外指定即可运行。若想自定义目录可参考 路径定制指南所有脚本均支持-h查看可选项并通过--checkpoint_path、--tokenizer_path显式传参。在 TPU 上运行推理权重就绪后推理开箱即用python3 generate.py --prompt Hello, my name is --num_samples 3该命令以Hello, my name is为提示词连续生成 3 段文本。原文档给出的实测现象为首次生成约需 20 秒XLA 编译图的开销之后每次生成回落到约 5 秒。generate.py的核心参数如下与 generate.py 中的函数签名一致参数默认值说明--promptHello, my name is生成所用的提示词--num_samples1生成的样本数量--max_new_tokens50每个样本最多生成的新 token 数--top_k200采样时仅从概率最高的 k 个 token 中抽取--temperature0.8采样随机性控制值越大随机性越高--checkpoint_pathcheckpoints/lit-llama/7B/lit-llama.pth模型权重路径--tokenizer_pathcheckpoints/lit-llama/tokenizer.modeltokenizer 路径结合源码可以还原 TPU 上的生成流程见 generate.py编码提示词后将输入填充为最终长度T_new的连续张量并维护input_pos记录当前位置每次迭代仅对最新位置的 token 做前向配合model.reset_cache()清理 KV 缓存使用top_k裁剪后做 softmax 与torch.multinomial采样生成下一个 token每轮循环前后调用xm.mark_step()让 XLA 推进图的执行命中eos_id则提前截断返回否则生成满max_new_tokens个 token。脚本还会逐样本打印Time for inference与tokens/sec统计方便在 TPU 上直接验证吞吐表现见 generate.py。微调支持状态截至原文档编写时TPU 上的微调标记为 Coming soon即仓库尚未提供在 TPU 上运行微调的官方教程与验证。目前 finetune 目录下的lora.py、adapter.py、full.py等脚本主要面向 GPU 环境如 README 中说明的 LoRA/Adapter 微调需要约 24 GB 显存。在 TPU 上自行尝试微调前建议先确认对应脚本在 XLA 设备上的算子兼容性。使用完毕删除实例Cloud TPU 是按使用计费的托管服务原文档特别提醒结束后务必删除实例避免产生持续费用gcloud compute tpus tpu-vm delete lit-llama --zoneus-central2-b删除命令与创建命令一样需要带上正确的--zone。注意事项与限制小结PjRT 仍处于实验阶段尽管它已成为 PyTorch XLA 2.0 的默认运行时官方仍建议显式设置PJRT_DEVICE与ALLOW_MULTIPLE_LIBTPU_LOAD两个环境变量首次编译延迟不可避免约 20 秒的首个样本延迟来自 XLA 图编译属于预期行为后续样本即恢复约 5 秒区域与镜像保持一致create、ssh、delete三条命令的--zoneus-central2-b必须统一镜像版本以tpu-vm-v4-pt-2.0PyTorch 2.0为准与仓库torch2.0.0的依赖要求吻合TPU 上不要复用 GPU 的假设例如精度选择generate.py 仅在检测到 CUDA 且支持 bf16 时才使用bf16-trueTPU 环境下会回退到32-true属于脚本的既定行为无需干预。至此你已经掌握了从创建 TPU v4 虚拟机、安装 Lit-LLaMA、配置 PjRT 环境、导入权重到最终在 TPU 上完成 LLaMA 文本生成的完整链路。更多推理细节可继续阅读 推理指南。赞分享人工智能大模型预训练微调LoRA模型量化【免费下载链接】lit-llamaImplementation of the LLaMA language model based on nanoGPT. Supports flash attention, Int8 and GPTQ 4bit quantization, LoRA and LLaMA-Adapter fine-tuning, pre-training. Apache 2.0-licensed.项目地址https://gitcode.com/gh_mirrors/li/lit-llama点击查看免费下载相关推荐20分钟搞定BLIP2-OPT-2.7B环境配置Windows/Linux/Mac全平台避坑指南20分钟搞定BLIP2 OPT 2.7B环境配置Windows/Linux/Mac全平台避坑指南 你是否曾因环境配置失败放弃AI视觉项目是否在CUDA、Py人工智能大模型预训练微调LoRA模型量化SkyPilot 上的 Cloud TPU v6eTrillium实战一键创建、Llama 3 8B 训练与 JetStream 推理服务SkyPilot 上的 Cloud TPU v6eTrillium实战一键创建、Llama 3 8B 训练与 JetStream 推理服务 本指南以仓库后端任务调度MLOps集群管理上一篇终极Office文档安全分析oletools完整指南与高效应用下一篇django-simple-history用户跟踪终极教程自动记录操作者信息创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考