conform EXAONE MoE slurm conversion example to public roundtrip launcher

Signed-off-by: sangHa0411 <sangha110495@gmail.com>

Author: sangHa0411

examples/models/exaone/exaone_moe/README.md

@@ -82,14 +82,24 @@ The Slurm scripts warm the shared `uv` cache once before all distributed ranks e
 
 ## Slurm Checkpoint Conversion
 
-[slurm_conversion.sh](slurm_conversion.sh) runs round-trip conversion across 16 GPUs on 2 nodes with default `TP=1, PP=1, EP=16`.
+[slurm_conversion.sh](slurm_conversion.sh) uses `convert.sh roundtrip` to submit a fixed `TP=1, PP=1, EP=16` config and verify HF ↔ Megatron round-trip conversion across 16 GPUs on 2 nodes. Run the wrapper from a Slurm login node; it submits one job and waits for it by default.
+
+Set the container and account, then launch:
 
 ```bash
-mkdir -p logs
-sbatch examples/models/exaone/exaone_moe/slurm_conversion.sh
+export CONTAINER_IMAGE=/path/to/container.sqsh
+export SLURM_ACCOUNT=<your-account>
+export SLURM_PARTITION=batch
+# Optional: export CONTAINER_MOUNTS, HF_TOKEN, HF_HOME, and UV_CACHE_DIR before launching.
+bash examples/models/exaone/exaone_moe/slurm_conversion.sh
 ```
 
-Before submitting, edit the `#SBATCH --account` and `#SBATCH --partition` lines for your cluster, or override them according to your scheduler policy.
+The current checkout is mounted automatically at `/opt/Megatron-Bridge` and must be on storage visible from the compute nodes. Forward any cluster-specific `srun` options your scheduler requires, for example:
+
+```bash
+bash examples/models/exaone/exaone_moe/slurm_conversion.sh \
+    --srun-arg=--mpi=pmix
+```
 
 ## Script Configuration
 
@@ -107,4 +117,4 @@ Before submitting, edit the `#SBATCH --account` and `#SBATCH --partition` lines
 
 ## Validation
 
-Focused EXAONE conversion coverage lives in `tests/functional_tests/test_groups/models/exaone/test_exaone_conversion.py`.
+Focused EXAONE conversion coverage lives in `tests/functional_tests/test_groups/models/exaone/test_exaone_conversion.py`.

examples/models/exaone/exaone_moe/slurm_conversion.sh

@@ -13,98 +13,53 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# ==============================================================================
-# K-EXAONE-236B-A23B Conversion Round-Trip Verification (Slurm)
+# K-EXAONE-236B-A23B round-trip verification on 2 Slurm nodes (16 GPUs).
+# Run this wrapper from a Slurm login node; convert.sh submits one job and
+# waits for it by default.
 #
-# K-EXAONE is a BF16 MoE model with 128 routed experts and 8 active experts.
-# The default TP=1, PP=1, EP=16 configuration uses 16 GPUs across 2 nodes.
-#
-# Usage:
-#   1. Set CONTAINER_IMAGE and, if needed, CONTAINER_MOUNTS.
-#   2. Export HF_TOKEN, HF_HOME, and UV_CACHE_DIR on shared storage.
-#   3. Create the log directory and submit:
-#        mkdir -p logs
-#        sbatch examples/models/exaone/exaone_moe/slurm_conversion.sh
-# ==============================================================================
-
-#SBATCH --job-name=k-exaone-roundtrip
-#SBATCH --nodes=2
-#SBATCH --ntasks-per-node=8
-#SBATCH --gpus-per-node=8
-#SBATCH --time=4:00:00
-#SBATCH --account=<your-account>
-#SBATCH --partition=batch
-#SBATCH --output=logs/k_exaone_roundtrip_%j.log
-#SBATCH --exclusive
+# Required:
+#   export CONTAINER_IMAGE=/path/to/container.sqsh
+#   export SLURM_ACCOUNT=<your-account>
+# Optional:
+#   export CONTAINER_MOUNTS=/shared:/shared,/host/path:/container/path
+#   bash "$0" --srun-arg=--mpi=pmix
 
 set -euo pipefail
 
-# -- Container ---------------------------------------------------------------
-CONTAINER_IMAGE="${CONTAINER_IMAGE:-}"
-CONTAINER_MOUNTS="${CONTAINER_MOUNTS:-}"
-WORKDIR="${WORKDIR:-/opt/Megatron-Bridge}"
+: "${CONTAINER_IMAGE:?Set CONTAINER_IMAGE to the Megatron-Bridge container}"
+: "${SLURM_ACCOUNT:?Set SLURM_ACCOUNT to your Slurm account}"
 
-# -- Model / Parallelism -----------------------------------------------------
-HF_MODEL_ID="${HF_MODEL_ID:-LGAI-EXAONE/K-EXAONE-236B-A23B}"
-TP="${TP:-1}"
-PP="${PP:-1}"
-EP="${EP:-16}"
-ETP="${ETP:-1}"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(git -C "${SCRIPT_DIR}" rev-parse --show-toplevel)"
+CONVERT_SH="${CONVERT_SH:-${REPO_ROOT}/scripts/conversion/convert.sh}"
 
-# -- Environment -------------------------------------------------------------
-# Keep HF_HOME and UV_CACHE_DIR on storage shared by all nodes.
-# export HF_TOKEN="..."
-# export HF_HOME="/path/to/shared/HF_HOME"
-# export UV_CACHE_DIR="/path/to/shared/uv_cache"
 export TORCH_NCCL_AVOID_RECORD_STREAMS=1
 export NCCL_NVLS_ENABLE=0
-export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
-export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
-
-if [[ -z "$CONTAINER_IMAGE" ]]; then
-    echo "ERROR: Set CONTAINER_IMAGE to the Enroot .sqsh image path."
-    exit 1
-fi
-
-WORLD_SIZE="${SLURM_NTASKS:?SLURM_NTASKS is not set}"
-MODEL_PARALLEL_SIZE=$((TP * PP * EP))
-if [[ "$MODEL_PARALLEL_SIZE" -ne "$WORLD_SIZE" ]]; then
-    echo "ERROR: TP*PP*EP=$MODEL_PARALLEL_SIZE must equal allocated tasks=$WORLD_SIZE."
-    exit 1
-fi
-if ((128 % EP != 0)); then
-    echo "ERROR: EP=$EP must divide K-EXAONE's 128 routed experts."
-    exit 1
-fi
-
-SRUN_ARGS=(--mpi=pmix --container-image="$CONTAINER_IMAGE" --no-container-mount-home)
-if [[ -n "$CONTAINER_MOUNTS" ]]; then
-    SRUN_ARGS+=(--container-mounts="$CONTAINER_MOUNTS")
-fi
-
-echo "======================================"
-echo "K-EXAONE Round-Trip Conversion"
-echo "Job: $SLURM_JOB_ID | Nodes: $SLURM_JOB_NUM_NODES | Tasks: $WORLD_SIZE"
-echo "Model: $HF_MODEL_ID"
-echo "TP=$TP PP=$PP EP=$EP ETP=$ETP"
-echo "======================================"
 
-# Warm the shared uv cache once before all ranks enter distributed setup.
-srun --nodes=1 --ntasks=1 "${SRUN_ARGS[@]}" \
-    bash -c 'cd "$1" && uv sync' bash "$WORKDIR"
+MOUNT_ARGS=(--mount "${REPO_ROOT}:/opt/Megatron-Bridge")
+IFS=',' read -r -a EXTRA_MOUNTS <<< "${CONTAINER_MOUNTS:-}"
+for mount in "${EXTRA_MOUNTS[@]}"; do
+    if [[ -n "${mount}" ]]; then
+        MOUNT_ARGS+=(--mount "${mount}")
+    fi
+done
 
-srun "${SRUN_ARGS[@]}" \
-    bash -c '
-        cd "$1"
-        uv run --no-sync python examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
-            --hf-model-id "$2" \
-            --tp "$3" \
-            --pp "$4" \
-            --ep "$5" \
-            --etp "$6" \
-            --trust-remote-code
-    ' bash "$WORKDIR" "$HF_MODEL_ID" "$TP" "$PP" "$EP" "$ETP"
+ENV_ARGS=()
+for name in HF_TOKEN HF_HOME UV_CACHE_DIR TORCH_NCCL_AVOID_RECORD_STREAMS NCCL_NVLS_ENABLE; do
+    if [[ -n "${!name:-}" ]]; then
+        ENV_ARGS+=(--env "${name}")
+    fi
+done
 
-echo "======================================"
-echo "Round-trip conversion completed"
-echo "======================================"
+"${CONVERT_SH}" roundtrip \
+    --executor slurm --device gpu \
+    --nodes 2 --gpus-per-node 8 \
+    --account "${SLURM_ACCOUNT}" --partition "${SLURM_PARTITION:-batch}" --time 04:00:00 \
+    --container-image "${CONTAINER_IMAGE}" \
+    "${MOUNT_ARGS[@]}" \
+    "${ENV_ARGS[@]}" \
+    --experiment-name k-exaone-moe-roundtrip \
+    --hf-model LGAI-EXAONE/K-EXAONE-236B-A23B \
+    --tp 1 --pp 1 --ep 16 --etp 1 \
+    --trust-remote-code \
+    "$@"

tests/unit_tests/conversion/launcher/test_slurm_examples.py

@@ -84,6 +84,16 @@
             "trust_remote_code": True,
         },
     ),
+    (
+        "examples/models/exaone/exaone_moe/slurm_conversion.sh",
+        {
+            "nodes": "2",
+            "time": "04:00:00",
+            "hf_model": "LGAI-EXAONE/K-EXAONE-236B-A23B",
+            "config": ("1", "1", "16", "1"),
+            "trust_remote_code": True,
+        },
+    ),
 ]