38 lines
1.4 KiB
Bash
Executable File
38 lines
1.4 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
REPO_ROOT="${REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
|
|
IMAGE_TAG="${IMAGE_TAG:-laoyao/nemo-megatron:26.06-flashattn4}"
|
|
BASE_IMAGE="${BASE_IMAGE:-ti-coding-agent/nemo-bridge-flashattn4:26.06}"
|
|
INSTALL_FLASH_ATTN4="${INSTALL_FLASH_ATTN4:-0}"
|
|
|
|
export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
|
|
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
|
|
export HTTP_PROXY="${HTTP_PROXY:-$http_proxy}"
|
|
export HTTPS_PROXY="${HTTPS_PROXY:-$https_proxy}"
|
|
|
|
docker build \
|
|
--network=host \
|
|
--build-arg BASE_IMAGE="$BASE_IMAGE" \
|
|
--build-arg INSTALL_FLASH_ATTN4="$INSTALL_FLASH_ATTN4" \
|
|
--build-arg PIP_INDEX_URL="${PIP_INDEX_URL:-https://mirrors.aliyun.com/pypi/simple/}" \
|
|
--build-arg PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL:-https://pypi.org/simple}" \
|
|
-f "$REPO_ROOT/docker/nemo/Dockerfile.nemo-megatron" \
|
|
-t "$IMAGE_TAG" \
|
|
"$REPO_ROOT"
|
|
|
|
docker run --rm --gpus all --ipc=host --network=host "$IMAGE_TAG" bash -lc '
|
|
python3 - <<PY
|
|
import importlib.metadata as metadata
|
|
for package in ["torch", "transformer-engine", "flash-attn", "flash-attn-4", "nvidia-cutlass-dsl", "megatron-core", "megatron-bridge"]:
|
|
try:
|
|
print(package, metadata.version(package))
|
|
except metadata.PackageNotFoundError:
|
|
print(package, "MISSING")
|
|
import cutlass
|
|
print("cutlass import OK")
|
|
PY
|
|
'
|
|
|
|
echo "Built image: $IMAGE_TAG"
|