44 lines
1.3 KiB
Bash
Executable File
44 lines
1.3 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
# Host-side launcher for the Megatron-LM text generation server.
|
|
# Run on g0050 from /ssd/workspace/yi/laoyao_2b_moe.
|
|
|
|
REPO_ROOT=${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}
|
|
IMAGE=${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}
|
|
SERVER_PORT=${SERVER_PORT:-5000}
|
|
MASTER_PORT=${MASTER_PORT:-6000}
|
|
CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
|
|
CONTAINER_NAME=${CONTAINER_NAME:-laoyao_megatron_server}
|
|
DETACH=${DETACH:-0}
|
|
LOG_DIR=${LOG_DIR:-${REPO_ROOT}/runs/megatron_inference}
|
|
|
|
mkdir -p "${LOG_DIR}"
|
|
|
|
DOCKER_ARGS=(
|
|
--gpus all
|
|
--network host
|
|
--ipc host
|
|
--shm-size 64g
|
|
--ulimit memlock=-1
|
|
--ulimit stack=67108864
|
|
--name "${CONTAINER_NAME}"
|
|
-e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}"
|
|
-e SERVER_PORT="${SERVER_PORT}"
|
|
-e MASTER_PORT="${MASTER_PORT}"
|
|
-v "${REPO_ROOT}:/work"
|
|
-w /work
|
|
)
|
|
|
|
docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true
|
|
|
|
if [[ "${DETACH}" == "1" ]]; then
|
|
docker run -d "${DOCKER_ARGS[@]}" "${IMAGE}" \
|
|
bash scripts/serve_laoyao_megatron_inner.sh
|
|
echo "Started ${CONTAINER_NAME}"
|
|
echo "Logs: docker logs -f ${CONTAINER_NAME}"
|
|
else
|
|
docker run --rm "${DOCKER_ARGS[@]}" "${IMAGE}" \
|
|
bash scripts/serve_laoyao_megatron_inner.sh 2>&1 | tee "${LOG_DIR}/server_$(date +%Y%m%d_%H%M%S).log"
|
|
fi
|