42 lines
1.7 KiB
Docker
42 lines
1.7 KiB
Docker
# ==============================================================
|
|
# flash_attn 휠 빌드 전용 이미지
|
|
#
|
|
# 결과 이미지: /wheels/flash_attn-{version}-cp312-cp312-linux_x86_64.whl 포함
|
|
# Gitea 레지스트리 태그: 10.89.0.8:3000/gm/flash-attn-builder:cu128-torch211
|
|
#
|
|
# 빌드: podman build -t 10.89.0.8:3000/gm/flash-attn-builder:cu128-torch211 .
|
|
# 재빌드 필요 조건: torch 버전 변경, flash_attn 버전 변경, CUDA 버전 변경
|
|
# ==============================================================
|
|
|
|
# ── Stage 1: 컴파일 ──────────────────────────────────────────
|
|
FROM docker.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 AS builder
|
|
|
|
ENV DEBIAN_FRONTEND=noninteractive
|
|
ENV PIP_NO_CACHE_DIR=1
|
|
ENV PATH=/opt/venv/bin:$PATH
|
|
|
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
python3.12 python3.12-venv python3.12-dev python3-pip \
|
|
build-essential ninja-build git \
|
|
&& rm -rf /var/lib/apt/lists/*
|
|
|
|
RUN python3.12 -m venv /opt/venv && pip install --upgrade pip setuptools wheel
|
|
|
|
ARG TORCH_VERSION=2.11.0+cu128
|
|
ARG TORCH_INDEX=https://download.pytorch.org/whl/cu128
|
|
|
|
# torch만 설치 (flash_attn 컴파일에 torch 헤더 필요, GPU 불필요)
|
|
RUN pip install torch==${TORCH_VERSION} --extra-index-url ${TORCH_INDEX}
|
|
|
|
ARG FLASH_ATTN_VERSION=2.8.3
|
|
ARG MAX_JOBS=4
|
|
|
|
RUN mkdir /wheels && \
|
|
MAX_JOBS=${MAX_JOBS} pip wheel flash-attn==${FLASH_ATTN_VERSION} \
|
|
--no-build-isolation --no-deps -w /wheels
|
|
|
|
# ── Stage 2: 휠만 담은 최소 이미지 ──────────────────────────
|
|
FROM alpine:3
|
|
COPY --from=builder /wheels/ /wheels/
|
|
CMD ["ls", "/wheels/"]
|