diff --git a/scripts/preprocess_megatron_bridge_pretrain_direct.sh b/scripts/preprocess_megatron_bridge_pretrain_direct.sh
index 42eabc7..35e7410 100755
--- a/scripts/preprocess_megatron_bridge_pretrain_direct.sh
+++ b/scripts/preprocess_megatron_bridge_pretrain_direct.sh
@@ -2,7 +2,7 @@
 set -euo pipefail
 
 REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
-IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
+IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
 SOURCE_DATA_DIRS="${SOURCE_DATA_DIRS:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_v1_20260701}"
 WORK_DIR="${WORK_DIR:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_direct_smoke_v1}"
 TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
@@ -46,9 +46,17 @@ if [[ "$OVERWRITE" == "1" ]]; then
   overwrite_args+=(--overwrite)
 fi
 
+DOCKER_MOUNTS=()
+if [[ -d /mnt/beegfs ]]; then
+  DOCKER_MOUNTS+=(-v /mnt/beegfs:/mnt/beegfs)
+fi
+if [[ -d /ssd ]]; then
+  DOCKER_MOUNTS+=(-v /ssd:/ssd)
+fi
+
 docker run --rm --ipc=host --network=host \
   --ulimit memlock=-1 --ulimit stack=67108864 \
-  -v /mnt/beegfs:/mnt/beegfs \
+  "${DOCKER_MOUNTS[@]}" \
   -w "$REPO_ROOT" \
   "$IMAGE" \
   python3 dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py \