ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

如何用Docker部署FlagEmbedding文本嵌入:3步从构建镜像到GPU调优

如何用Docker部署FlagEmbedding文本嵌入:3步从构建镜像到GPU调优 如何用Docker部署FlagEmbedding文本嵌入3步从构建镜像到GPU调优【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbeddingFlagEmbedding 是专注密集检索与 RAG 的开源项目提供 BGE 嵌入Embedding与重排Reranker模型的统一推理接口。这篇实操教程帮你在 Docker 里把它跑起来构建运行时镜像、验证嵌入链路、挂上 GPU 服务并调优参数照着敲命令即可得到一套可复用的推理环境。一图看懂全局这张图就是你容器里要跑的东西图中展示了 FlagEmbedding 在 RAG 系统里的位置文档先过嵌入模型向量化检索命中后交给重排模型排序最终送入 LLM 生成答案。你的容器要承载的正是嵌入 重排这一层。开工前检查先确认机器和依赖满足条件开工前先对照下表确认环境缺什么补什么能省掉后面大半的排查时间。类别要求说明GPUNVIDIA 单卡 8GB 显存起16GB 可跑微调纯 CPU 也能推理只是慢内存 / 磁盘16GB 内存、20GB 空闲磁盘磁盘主要留给 HuggingFace 模型缓存软件Docker 20.10、NVIDIA Container Toolkit装好后nvidia-smi在容器外必须能看到卡Python 依赖torch1.6、transformers4.44.2、datasets、peft由项目setup.py统一声明无需手动对齐微调额外依赖deepspeed、flash-attn只在你要做训练时才需要跑两条命令自检都能正常输出就说明环境就绪# 确认 Docker 与 GPU 驱动链路正常 docker version nvidia-smi# 若宿主机已有 PyTorch顺手确认 CUDA 可用 python3 -c import torch; print(torch.__version__, torch.cuda.is_available())10 分钟构建 FlagEmbedding 推理镜像 这一步把项目打包成 Docker 镜像核心思路是CUDA 基础镜像 PyTorch 项目本体依赖全部交给setup.py解析。先在项目根目录新建Dockerfile每行作用都写在注释里# 与 PyTorch cu117 轮子匹配的官方 CUDA 基础镜像 FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu22.04 WORKDIR /app # 系统层只装构建需要的最小工具 RUN apt-get update apt-get install -y --no-install-recommends \ git python3 python3-pip rm -rf /var/lib/apt/lists/* # 先装 PyTorch指定 cu117 源保证与基础镜像的 CUDA 一致 RUN pip3 install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu117 # 拷贝项目并安装setup.py 会带出 transformers/datasets/peft 等全部依赖 COPY . . RUN pip3 install --no-cache-dir . # 把 HuggingFace 缓存固定到容器内路径方便外部挂载复用 ENV HF_HOME/app/hf_cache CMD [bash]然后执行构建首次拉取基础镜像较慢10~20 分钟属正常docker build -t flagembedding:1.4 .5 分钟验证镜像真的能跑嵌入 这一步不是只验证 import 成功而是直接跑官方示例做一次真实推理确认整条链路可用。仓库自带了单卡推理示例运行后它会自动从 HuggingFace 下载BAAI/bge-multilingual-gemma2模型、对 200 条 query 和 passage 编码最后打印分数矩阵docker run --rm --gpus all flagembedding:1.4 \ python examples/inference/embedder/decoder_only/auto_base_single_device.py结尾处对照脚本内置的期望输出Expected Output: [[0.558 0.0212 ] [0.01651 0.526 ]]数值量级一致即说明 GPU、依赖、推理代码三者都没问题。没有 GPU 的话把示例里的devicescuda:0改成devicescpu再跑即可。挂上 GPU、固定缓存把服务稳定跑起来这一步解决两个实际问题模型缓存别每次重下训练输出别随容器销毁而丢失。做法是把宿主机目录挂载进容器的固定路径。以仓库自带的 decoder-only 微调脚本为例它用torchrun拉起FlagEmbedding.finetune.embedder.decoder_only.base模块缓存和输出目录都已挂载docker run --gpus all -it --rm \ -v $PWD/cache:/app/hf_cache \ -v $PWD/output:/app/output \ flagembedding:1.4 \ bash examples/finetune/embedder/decoder_only/base.sh两个注意点脚本默认num_gpus2单卡机器先把脚本里这一行改成1模型首次下载完成后落在$PWD/cache里之后所有容器共享启动时间从分钟级降到秒级。推理场景同理跑重排示例只需把命令换成python examples/inference/reranker/decoder_only/auto_base_single_device.py。调吞吐与显存三个真正有效的旋钮这一步针对跑得慢和OOM做调优三个参数都是仓库代码里真实存在的开关。# 推理端开启半精度、指定多卡分担编码负载 from FlagEmbedding import FlagAutoModel model FlagAutoModel.from_finetuned( BAAI/bge-multilingual-gemma2, use_fp16True, # 默认开启Ampere 以上显卡可换 use_bf16True devices[cuda:0, cuda:1], # 多卡自动分片编码 )# 训练端在 examples/finetune/embedder/decoder_only/base.sh 中调整这三处 per_device_train_batch_size2 # 显存吃紧就调小吞吐吃紧就调大 # --gradient_checkpointing # 用少量算力换显存默认已开启 # --deepspeed ../../ds_stage1.json # 多卡必开配置在 examples/ds_stage1.json显存不够时优先降 batch size 而不是换量化方案多卡时配合脚本里已有的--negatives_cross_device跨卡负采样还能顺带提升训练质量。踩坑速查症状原因解法容器内nvidia-smi报错、cuda.is_available()为 False未安装 NVIDIA Container Toolkit 或驱动不匹配安装 nvidia-container-toolkit 后systemctl restart docker宿主机先自检首次推理卡住、几分钟后才出结果正在从 HuggingFace 下载模型网络慢先在有网机器跑一次并挂载-v $PWD/cache:/app/hf_cache缓存复用微调时 CUDA OOMbatch size 超显存调小per_device_train_batch_size保留--gradient_checkpointing多卡加 deepspeed 配置import 时报 transformers 版本冲突手动装的依赖版本与项目约束打架删除手动安装的包只用pip install .让setup.py统一解析版本延伸资源Tutorials/quick_start.ipynb官方快速入门含嵌入、相似度计算、检索完整示例examples/inference/embedder 与 reranker 的官方推理示例单卡/多卡各一套FlagEmbedding/inference/推理入口源码FlagAutoModel与FlagAutoReranker的参数都在这里定义Tutorials/7_Fine-tuning/微调与难负例挖掘教程配合examples/finetune/使用到这里镜像构建、推理验证、GPU 挂载和调优四步已全部走通。下一步建议把这套容器接入你的 RAG 流程从Tutorials/6_RAG/里的 From-Scratch 示例开始验证端到端效果。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表