ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

FlagEmbedding容器化实战:3步从Docker镜像构建到GPU微调容器

FlagEmbedding容器化实战:3步从Docker镜像构建到GPU微调容器 FlagEmbedding容器化实战3步从Docker镜像构建到GPU微调容器【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbeddingFlagEmbedding 是一个面向密集检索与检索增强 LLM 的开源工具包内置 BGE 嵌入模型与重排序模型两大组件的推理、微调和评估代码。本文带你完成 FlagEmbedding 容器化部署全程 3 步跑通结束时会得到一个可复现的 GPU 镜像以及一个能直接执行官方微调脚本的容器。在第三台机器上重装环境时如果你的 BGE 推理先跑在个人电脑上又搬到团队 GPU 服务器再试了一次云端实例大概率经历过这样的循环torch 版本和 CUDA 对不上、transformers 被带出新版本后推理脚本报错、每次换机器都要重新下载一遍 1.3GB 的 bge-large 模型。问题根源在于运行环境散落在各台机器上每台机器的依赖组合都略有不同。容器化对它的价值恰好是两点环境被冻结torch、transformers 等依赖版本固化进镜像你机器上能跑的同事机器上原样能跑缓存可以带走把 HuggingFace 模型缓存挂到宿主机目录上换机器、重建镜像都不用重新拉模型图1FlagEmbedding 容器化部署覆盖 RAG 流程中的嵌入检索与重排两个高负载环节正是最受益于一套固定环境的模块第1步准备一台带 GPU 的 Docker 机器硬件底线如下够跑 bge-large 级模型的推理和单卡微调即可项目最低推荐CPU8 核16 核内存16 GB32 GBGPU1 块 8GB 显存1 块 16GB 显存软件侧要装两样Docker 20.10 以上以及 NVIDIA Container Toolkit提供--gpus能力。装完后验证 GPU 容器链路docker run --rm --gpus all nvidia/cuda:11.7.1-base-ubuntu22.04 nvidia-smi你应该看到终端输出显卡型号与显存占用列表。看到这张列表说明 GPU 容器通道打通如果报could not select device driver直接跳到文末的排坑清单。⚠️ 后文所有镜像统一使用 CUDA 11.7 基础镜像配 cu117 版本的 torch wheel两者必须匹配这是最容易踩的坑之一。第2步用最小 Dockerfile 构建镜像在仓库根目录setup.py所在层新建 Dockerfile核心思路只有一条把变化慢的大依赖 torch 放前面把经常改的项目代码放最后这样日常改代码重新构建时前面所有层都命中缓存。FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04 WORKDIR /opt/flagembedding # 系统依赖只装最小集 RUN apt-get update apt-get install -y --no-install-recommends \ git python3 python3-pip \ rm -rf /var/lib/apt/lists/* # torch 是最大依赖先装并锁定 cu117 wheel防止被拉成 CPU 版 RUN pip install --no-cache-dir -U pip \ pip install --no-cache-dir torch2.4.0 \ --index-url https://download.pytorch.org/whl/cu117 # 项目文件变化最频繁放最后 COPY . . RUN pip install --no-cache-dir . deepspeed # 固定模型缓存路径禁用 wandb 上报官方微调脚本同样这样做 ENV HF_HOME/root/.cache/huggingface \ HF_HUB_CACHE/root/.cache/huggingface/hub \ WANDB_MODEdisabled然后执行docker build -t flagemb:dev . docker images | grep flagemb你应该看到构建结束输出Successfully tagged flagemb:devdocker images里出现约 8GB 的镜像。首次构建要拉基础镜像和几个 GB 的 wheel耗时 10 分钟左右之后只改代码重建通常几十秒。第3步用官方示例验证模型加载先不写自己的业务代码直接跑仓库自带的推理示例——它是仓库里自带预期输出的验收脚本最适合作为容器验收标准docker run --rm --gpus all \ -v $PWD/hf-cache:/root/.cache/huggingface \ -w /opt/flagembedding \ flagemb:dev \ python examples/inference/embedder/encoder_only/base_single_device.py你应该看到首次运行先下载bge-small-en-v1.5约 130MB随后打印一个 2×2 的余弦相似度矩阵对角元素约 0.794 和 0.801明显高于交叉项且与脚本末尾自带的Expected Output一致。看到这组数字说明模型加载、GPU 前向和依赖版本链全部正确。Dockerfile 背后的三个决策点第 2 步的 Dockerfile 每一行都有理由理解这三点后你改镜像不会手抖。基础镜像与依赖分层为什么用 devel 版 CUDA 镜像[finetune]扩展里的 flash-attn 等组件含 C/CUDA 扩展编译期需要 devel 包里的头文件为什么 torch 单独先装setup.py只声明了torch1.6.0直接pip install .有概率拉进 CPU 版 torch单独装 cu117 wheel 后再装项目pip 会复用已装版本为什么 ENV 要写 HF_HOME这是后续模型缓存挂载的落点路径在镜像里写死挂载目录才能跨机器复用多阶段构建把镜像瘦身8GB 的镜像对分发和拉取都不友好。多阶段构建的做法在 builder 阶段用 devel 镜像完成全部编译安装最终镜像换用 runtime 基础镜像不含编译器和 apt 缓存通常能瘦掉 2~4GB# ---------- 阶段一builder负责安装 ---------- FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04 AS builder WORKDIR /opt/flagembedding RUN apt-get update apt-get install -y --no-install-recommends \ python3 python3-pip python3-venv \ rm -rf /var/lib/apt/lists/* RUN python3 -m venv /opt/venv \ /opt/venv/bin/pip install --no-cache-dir -U pip \ /opt/venv/bin/pip install --no-cache-dir torch2.4.0 \ --index-url https://download.pytorch.org/whl/cu117 COPY . . RUN /opt/venv/bin/pip install --no-cache-dir . deepspeed # ---------- 阶段二runtime只带运行必需的 ---------- FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu22.04 COPY --frombuilder /opt/venv /opt/venv COPY --frombuilder /opt/flagembedding /opt/flagembedding ENV PATH/opt/venv/bin:$PATH \ HF_HOME/root/.cache/huggingface \ WANDB_MODEdisabled WORKDIR /opt/flagembedding你应该看到docker images中最终镜像体积明显下降且重跑第 3 步的验证命令仍输出相同的相似度矩阵。运行容器挂载规划与常用参数前台调试还是后台任务仓库本身没有内置 HTTP 服务容器是任务型的跑一次推理、微调或评估。所以运行模式只有两种场景命令形态看输出的方式交互式调试docker run -it --rm ...直接终端微调 / 长评估docker run -d --name ...docker logs -f 容器名挂载三个目录容器内路径宿主机目录建议用途/root/.cache/huggingface$PWD/hf-cache或共享盘模型缓存跨机器复用数据目录你的训练数据所在路径微调输入工作目录你的输出路径checkpoint 与评估结果⚠️ 官方微调脚本的 checkpoint 输出到工作目录下的./test_encoder_only_base_*务必确认容器工作目录-w参数落在挂载点内否则容器删掉后产物就丢了。GPU 容器启动参数速查参数作用什么时候用--gpus all全部 GPU 进容器默认选择--gpus device0只透传指定 GPU多卡机器做任务隔离-e CUDA_VISIBLE_DEVICES1容器内再收窄可见卡与上一项组合--shm-size16g扩大/dev/shmtorchrun 多卡微调必加-v挂载目录缓存、数据、输出最容易漏的是--shm-size默认 64MB 的共享内存撑不住多卡 NCCL 通信不加它微调会在第一步就崩。把官方 encoder-only 微调脚本放进容器后台跑示例默认双卡这里用sed改成单卡脚本内的数据路径是相对工作目录的所以-w必须指向encoder_only/docker run -d --name flagemb-ft \ --gpus all --shm-size16g \ -v $PWD/hf-cache:/root/.cache/huggingface \ -w /opt/flagembedding/examples/finetune/embedder/encoder_only \ flagemb:dev \ bash -c sed -i s/^num_gpus2/num_gpus1/ base.sh \ bash base.sh 21 | tee train.log你应该看到docker logs -f flagemb-ft每秒输出一条 step 日志示例设置了logging_steps 1loss 持续打印同时挂载目录里出现train.log和 checkpoint 目录。生产环境的三个调整模型缓存挂载与预下载生产环境里第一次跑任务不该在等模型下载。在起业务容器之前先把模型预置到共享缓存目录docker run --rm \ -v /data/hf-cache:/root/.cache/huggingface \ flagemb:dev \ huggingface-cli download BAAI/bge-large-en-v1.5你应该看到宿主机缓存目录里出现模型权重分片文件之后任何容器挂载同一目录启动时直接命中缓存不再产生下载流量。多机共用一块共享盘时这个目录还能作为模型的统一分发源。批大小与显存的关系官方脚本里per_device_train_batch_size被故意设成 2配合 toy 数据生产上必须上调。方法从单卡不 OOM 的最大值开始用gradient_accumulation_steps补齐等效 batch。脚本本身已开启fp16和gradient_checkpointing前者省显存、后者用少量速度换激活内存。你应该看到训练日志无 OOMnvidia-smi里显存占用稳定落在 80%~95% 区间。日志落到磁盘docker logs跟随容器生命周期容器一删日志就没了。前文的后台命令已用tee把训练输出同时写入挂载目录这是长任务的底线配置容器可以随时删日志和 checkpoint 都留在宿主机。你应该看到宿主机挂载目录中的train.log随训练增长删掉容器后文件仍在。排坑清单最可能碰到的 4 个现象现象原因解决could not select device driver或容器内找不到 nvidia-smi未安装 NVIDIA Container Toolkit安装后执行systemctl restart docker再重跑第 1 步的验证命令多卡微调第一步即崩报错涉及/dev/shm默认共享内存仅 64MBNCCL 不够用docker run加--shm-size16gtorchrun 报进程数与可见 GPU 不匹配脚本里num_gpus与--gpus透传数量不一致两者对齐单卡就先把脚本改成num_gpus1首次构建极慢基础镜像加 torch wheel 合计数 GB换网络良好的构建机之后改代码重建只重跑最后一层几十秒完成进阶路线容器跑通之后按仓库自带的教程体系深入以下四个入口按阶段取用Tutorials/quick_start.ipynb —— 5 分钟入门适合刚接触项目、先建立基本手感再看容器细节的阶段Tutorials/7_Fine-tuning/ —— 准备在容器里微调自己数据时看含数据格式准备与参数讲解examples/finetune/README.md —— 改启动脚本时的参数速查配合examples/finetune/embedder/encoder_only/base.sh使用Tutorials/4_Evaluation/ —— 想在容器内跑 MTEB / MSMARCO 评估验证微调效果时的完整流程图2FlagEmbedding 教程地图容器化部署就绪后可沿此地图逐模块深入收尾FlagEmbedding 的容器化部署解决的不是多一层封装而是把三样原本不稳定的东西固定下来依赖版本、GPU 环境、模型缓存。镜像和挂载规划就位后装环境这件事就退化成一条 docker run 命令。接下来可以立刻做的两件事用第 4 节的多阶段 Dockerfile 重建生产镜像用docker images确认瘦身效果并重跑验证脚本确认输出不变按第 5 节的流程把BAAI/bge-large-en-v1.5预置进共享缓存然后后台启动微调示例用docker logs -f flagemb-ft观察 loss 曲线【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表