Support g0050 SSD preprocessing path

This commit is contained in:
yi_lu
2026-07-02 21:59:02 +08:00
parent b8b5c3ffe0
commit 0e8c2a554d

View File

@@ -2,7 +2,7 @@
set -euo pipefail set -euo pipefail
REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}" 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}" 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}" 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}" TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
@@ -46,9 +46,17 @@ if [[ "$OVERWRITE" == "1" ]]; then
overwrite_args+=(--overwrite) overwrite_args+=(--overwrite)
fi 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 \ docker run --rm --ipc=host --network=host \
--ulimit memlock=-1 --ulimit stack=67108864 \ --ulimit memlock=-1 --ulimit stack=67108864 \
-v /mnt/beegfs:/mnt/beegfs \ "${DOCKER_MOUNTS[@]}" \
-w "$REPO_ROOT" \ -w "$REPO_ROOT" \
"$IMAGE" \ "$IMAGE" \
python3 dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py \ python3 dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py \