diff --git a/docker/nemo/Dockerfile.nemo-megatron b/docker/nemo/Dockerfile.nemo-megatron index 101271d..e624cfd 100644 --- a/docker/nemo/Dockerfile.nemo-megatron +++ b/docker/nemo/Dockerfile.nemo-megatron @@ -15,8 +15,29 @@ RUN if [ "${INSTALL_FLASH_ATTN4}" = "1" ]; then \ "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)) @@ -28,4 +49,8 @@ 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