36 lines
810 B
Bash
Executable File
36 lines
810 B
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
|
|
IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
|
|
DOCKER_PULL="${DOCKER_PULL:-0}"
|
|
USE_GPUS="${USE_GPUS:-0}"
|
|
|
|
cd "$REPO_ROOT"
|
|
python3 tools/validate_model_architecture.py
|
|
|
|
if [ "$DOCKER_PULL" = "1" ]; then
|
|
docker pull "$IMAGE"
|
|
fi
|
|
|
|
GPU_ARGS=()
|
|
if [ "$USE_GPUS" = "1" ]; then
|
|
GPU_ARGS=(--gpus all)
|
|
fi
|
|
|
|
docker run --rm "${GPU_ARGS[@]}" --ipc=host --network=host \
|
|
-v /mnt/beegfs:/mnt/beegfs \
|
|
-w "$REPO_ROOT" \
|
|
"$IMAGE" \
|
|
bash -lc 'set -euo pipefail
|
|
python - <<PY
|
|
import importlib
|
|
mods = ["torch", "nemo", "megatron.core"]
|
|
for name in mods:
|
|
mod = importlib.import_module(name)
|
|
version = getattr(mod, "__version__", "")
|
|
print("IMPORT_OK", name, version)
|
|
PY
|
|
python tools/validate_model_architecture.py
|
|
'
|