57 lines
2.0 KiB
Docker
57 lines
2.0 KiB
Docker
ARG BASE_IMAGE=nvcr.io/nvidia/nemo:26.06
|
|
FROM ${BASE_IMAGE}
|
|
|
|
ARG INSTALL_FLASH_ATTN4=0
|
|
ARG PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/
|
|
ARG PIP_EXTRA_INDEX_URL=https://pypi.org/simple
|
|
|
|
ENV PIP_INDEX_URL=${PIP_INDEX_URL}
|
|
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
|
ENV PIP_DISABLE_PIP_VERSION_CHECK=1
|
|
|
|
RUN if [ "${INSTALL_FLASH_ATTN4}" = "1" ]; then \
|
|
python3 -m pip uninstall -y flash-attn flash-attn-3 flash-attn-4 || true; \
|
|
MAX_JOBS="${MAX_JOBS:-16}" python3 -m pip install --no-cache-dir --no-build-isolation \
|
|
"flash-attn-4" "nvidia-cutlass-dsl"; \
|
|
fi
|
|
|
|
RUN set -eux; \
|
|
python3 - <<'PY'
|
|
from pathlib import Path
|
|
import shutil
|
|
|
|
site = Path("/usr/local/lib/python3.12/dist-packages")
|
|
src = site / "vllm/third_party/triton_kernels"
|
|
dst = site / "triton_kernels"
|
|
if not (src / "matmul_ogs.py").exists() or not (src / "tensor.py").exists():
|
|
raise SystemExit(f"missing vLLM triton_kernels source: {src}")
|
|
if not dst.exists():
|
|
raise SystemExit(f"missing top-level triton_kernels package: {dst}")
|
|
backup = site / "triton_kernels_nv26_backup"
|
|
if backup.exists():
|
|
shutil.rmtree(backup)
|
|
shutil.copytree(dst, backup, ignore=shutil.ignore_patterns("__pycache__"))
|
|
shutil.rmtree(dst)
|
|
shutil.copytree(src, dst, ignore=shutil.ignore_patterns("__pycache__"))
|
|
PY
|
|
|
|
RUN python3 - <<'PY'
|
|
import importlib.metadata as metadata
|
|
import importlib.util
|
|
required = ["torch", "transformer-engine", "megatron-core", "megatron-bridge"]
|
|
for package in required:
|
|
print(package, metadata.version(package))
|
|
try:
|
|
print("flash-attn-4", metadata.version("flash-attn-4"))
|
|
except metadata.PackageNotFoundError:
|
|
print("flash-attn-4 MISSING")
|
|
try:
|
|
print("nvidia-cutlass-dsl", metadata.version("nvidia-cutlass-dsl"))
|
|
except metadata.PackageNotFoundError:
|
|
print("nvidia-cutlass-dsl MISSING")
|
|
spec = importlib.util.find_spec("triton_kernels.matmul_ogs")
|
|
if spec is None:
|
|
raise SystemExit("triton_kernels.matmul_ogs MISSING")
|
|
print("triton_kernels.matmul_ogs", spec.origin)
|
|
PY
|