Add Megatron-Bridge pretrain launcher

scripts/preprocess_megatron_bridge_pretrain.sh

@@ -0,0 +1,38 @@
+#!/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}"
+SOURCE_DATA="${SOURCE_DATA:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701}"
+WORK_DIR="${WORK_DIR:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1}"
+JSONL="${JSONL:-$WORK_DIR/text.jsonl}"
+OUTPUT_PREFIX="${OUTPUT_PREFIX:-$WORK_DIR/laoyao_2b_moe_8192_text_document}"
+TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
+WORKERS="${WORKERS:-16}"
+MAX_DOCS="${MAX_DOCS:-0}"
+
+mkdir -p "$WORK_DIR"
+
+docker run --rm --ipc=host --network=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -v /mnt/beegfs:/mnt/beegfs \
+  -w "$REPO_ROOT" \
+  "$IMAGE" \
+  bash -lc "
+    set -euo pipefail
+    python3 dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py \
+      --input '$SOURCE_DATA/*.parquet' \
+      --output '$JSONL' \
+      --max-docs '$MAX_DOCS'
+    python3 /opt/Megatron-Bridge/3rdparty/Megatron-LM/tools/preprocess_data.py \
+      --input '$JSONL' \
+      --json-keys text \
+      --tokenizer-type HuggingFaceTokenizer \
+      --tokenizer-model '$TOKENIZER_MODEL' \
+      --append-eod \
+      --output-prefix '$OUTPUT_PREFIX' \
+      --workers '$WORKERS'
+    ls -lh '${OUTPUT_PREFIX}'*
+  "
+
+echo "Megatron indexed dataset prefix: $OUTPUT_PREFIX"

scripts/train_megatron_bridge_2b_moe.sh

@@ -0,0 +1,48 @@
+#!/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}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+DATA_PREFIX="${DATA_PREFIX:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1/laoyao_2b_moe_8192_text_document}"
+TRAIN_ITERS="${TRAIN_ITERS:-10}"
+SEQ_LENGTH="${SEQ_LENGTH:-8192}"
+MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-1024}"
+TP="${TP:-1}"
+PP="${PP:-1}"
+EP="${EP:-1}"
+CP="${CP:-1}"
+DRY_RUN="${DRY_RUN:-0}"
+
+if [[ "$DRY_RUN" != "1" && ! -f "${DATA_PREFIX}.idx" ]]; then
+  echo "missing Megatron indexed data prefix: $DATA_PREFIX" >&2
+  echo "run scripts/preprocess_megatron_bridge_pretrain.sh first, or set DATA_PREFIX" >&2
+  exit 1
+fi
+
+DRY_RUN_ARG=""
+if [[ "$DRY_RUN" == "1" ]]; then
+  DRY_RUN_ARG="--dry-run"
+fi
+
+docker run --rm --gpus all --ipc=host --network=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -v /mnt/beegfs:/mnt/beegfs \
+  -w "$REPO_ROOT" \
+  "$IMAGE" \
+  bash -lc "
+    set -euo pipefail
+    torchrun --nproc_per_node='$NPROC_PER_NODE' \
+      training/megatron_bridge/laoyao_2b_moe_pretrain.py \
+      --data-prefix '$DATA_PREFIX' \
+      --seq-length '$SEQ_LENGTH' \
+      --train-iters '$TRAIN_ITERS' \
+      --micro-batch-size '$MICRO_BATCH_SIZE' \
+      --global-batch-size '$GLOBAL_BATCH_SIZE' \
+      --tensor-parallel '$TP' \
+      --pipeline-parallel '$PP' \
+      --expert-parallel '$EP' \
+      --context-parallel '$CP' \
+      $DRY_RUN_ARG
+  "