Support stage1-2 MoE resume on g0050

scripts/resume_stage1_2_pretrain_8192_8gpu_mbs14.sh

@@ -0,0 +1,119 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Host path layout on g0050 after /ssd was replaced by /data.
+HOST_WORKSPACE="${HOST_WORKSPACE:-/data/workspace/yi}"
+CONTAINER_WORKSPACE="${CONTAINER_WORKSPACE:-/ssd/workspace/yi}"
+
+REPO_ROOT="${REPO_ROOT:-${CONTAINER_WORKSPACE}/laoyao_2b_moe}"
+HOST_REPO_ROOT="${HOST_REPO_ROOT:-${HOST_WORKSPACE}/laoyao_2b_moe}"
+IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
+RUN_NAME="${RUN_NAME:-stage1_2_8192_8gpu_dp8_mbs14_full_recompute_general55_code15_math15_science15}"
+RUN_DIR="${RUN_DIR:-${REPO_ROOT}/runs/${RUN_NAME}}"
+CKPT_DIR="${CKPT_DIR:-${RUN_DIR}/checkpoints}"
+TENSORBOARD_DIR="${TENSORBOARD_DIR:-${RUN_DIR}/tensorboard}"
+LOAD_DIR="${LOAD_DIR:-${CKPT_DIR}}"
+DATA_MANIFEST="${DATA_MANIFEST:-${CONTAINER_WORKSPACE}/laoyao_2b_moe_pretraining_dataset/megatron_bridge/stage2_8192_by_category_general55_code15_math15_science15_100b_v1/stage1_2_weighted_manifest.json}"
+VALIDATION_PREFIX="${VALIDATION_PREFIX:-${REPO_ROOT}/dataset/val/megatron_8192_glm52/heldout_2p8k_text_document}"
+TRAIN_ITERS="${TRAIN_ITERS:-281245}"
+SEQ_LENGTH="${SEQ_LENGTH:-8192}"
+MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-14}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-112}"
+SAVE_INTERVAL="${SAVE_INTERVAL:-2500}"
+EVAL_INTERVAL="${EVAL_INTERVAL:-15000}"
+EVAL_ITERS="${EVAL_ITERS:-10}"
+DATASET_WORKERS="${DATASET_WORKERS:-4}"
+LR="${LR:-3.0e-4}"
+MIN_LR="${MIN_LR:-5.0e-6}"
+WARMUP_FRACTION="${WARMUP_FRACTION:-0.01}"
+LOG_FILE="${LOG_FILE:-/tmp/laoyao_stage1_2_8192_8gpu_mbs14.log}"
+CONTAINER_NAME="${CONTAINER_NAME:-laoyao_stage1_2_pretrain}"
+
+HOST_DATA_MANIFEST="${DATA_MANIFEST/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+HOST_LOAD_DIR="${LOAD_DIR/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+HOST_VALIDATION_PREFIX="${VALIDATION_PREFIX/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+HOST_CKPT_DIR="${CKPT_DIR/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+HOST_TENSORBOARD_DIR="${TENSORBOARD_DIR/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+
+if [[ ! -d "${HOST_REPO_ROOT}" ]]; then
+  echo "ERROR: host repo dir not found: ${HOST_REPO_ROOT}" >&2
+  exit 1
+fi
+if [[ ! -f "${HOST_DATA_MANIFEST}" ]]; then
+  echo "ERROR: host data manifest not found: ${HOST_DATA_MANIFEST}" >&2
+  exit 1
+fi
+if [[ ! -d "${HOST_LOAD_DIR}" ]]; then
+  echo "ERROR: host load checkpoint dir not found: ${HOST_LOAD_DIR}" >&2
+  exit 1
+fi
+if [[ ! -f "${HOST_VALIDATION_PREFIX}.idx" ]]; then
+  echo "WARNING: host validation prefix not found: ${HOST_VALIDATION_PREFIX}" >&2
+fi
+
+if docker ps --format '{{.Names}}' | grep -Fxq "${CONTAINER_NAME}"; then
+  echo "ERROR: container already running: ${CONTAINER_NAME}" >&2
+  exit 1
+fi
+
+mkdir -p "${HOST_CKPT_DIR}" "${HOST_TENSORBOARD_DIR}"
+
+echo "Starting Laoyao 2B MoE stage1-2 pretraining"
+echo "host_workspace: ${HOST_WORKSPACE}"
+echo "container_workspace: ${CONTAINER_WORKSPACE}"
+echo "repo: ${REPO_ROOT}"
+echo "data_manifest: ${DATA_MANIFEST}"
+echo "save_checkpoint_dir: ${CKPT_DIR}"
+echo "load_checkpoint_dir: ${LOAD_DIR}"
+echo "load_checkpoint_iteration: $(cat "${HOST_LOAD_DIR}/latest_checkpointed_iteration.txt" 2>/dev/null || echo unknown)"
+echo "train_iters: ${TRAIN_ITERS}"
+echo "log_file: ${LOG_FILE}"
+
+cd "${HOST_REPO_ROOT}"
+
+nohup docker run --rm --name "${CONTAINER_NAME}" \
+  --gpus all \
+  --ipc=host \
+  --network=host \
+  --ulimit memlock=-1 \
+  --ulimit stack=67108864 \
+  -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+  -v "${HOST_WORKSPACE}:${CONTAINER_WORKSPACE}" \
+  -w "${REPO_ROOT}" \
+  "${IMAGE}" \
+  bash -lc "torchrun --nproc_per_node=8 \
+      training/megatron_bridge/laoyao_2b_moe_pretrain.py \
+      --data-manifest '${DATA_MANIFEST}' \
+      --validation-prefix '${VALIDATION_PREFIX}' \
+      --seq-length '${SEQ_LENGTH}' \
+      --train-iters '${TRAIN_ITERS}' \
+      --micro-batch-size '${MICRO_BATCH_SIZE}' \
+      --global-batch-size '${GLOBAL_BATCH_SIZE}' \
+      --tensor-parallel 1 \
+      --pipeline-parallel 1 \
+      --expert-parallel 1 \
+      --context-parallel 1 \
+      --split 999,1,0 \
+      --dataset-workers '${DATASET_WORKERS}' \
+      --lr '${LR}' \
+      --min-lr '${MIN_LR}' \
+      --warmup-fraction '${WARMUP_FRACTION}' \
+      --save-dir '${CKPT_DIR}' \
+      --load-dir '${LOAD_DIR}' \
+      --tensorboard-dir '${TENSORBOARD_DIR}' \
+      --save-interval '${SAVE_INTERVAL}' \
+      --keep-last-checkpoints 10 \
+      --log-interval 10 \
+      --eval-interval '${EVAL_INTERVAL}' \
+      --eval-iters '${EVAL_ITERS}' \
+      --use-distributed-optimizer \
+      --overlap-grad-reduce \
+      --overlap-param-gather \
+      --recompute-granularity full \
+      --recompute-method uniform \
+      --recompute-num-layers 1" \
+  > "${LOG_FILE}" 2>&1 &
+
+echo "Launched. Check with:"
+echo "  tail -f ${LOG_FILE}"
+echo "  docker ps"

training/megatron_bridge/laoyao_2b_moe_pretrain.py

@@ -44,7 +44,7 @@ def _token_weights_from_manifest(prefix_list: list[str], manifest: dict) -> list
     token_by_prefix = {}
     for result in manifest.get("results") or []:
         output_prefix = result.get("output_prefix")
-        tokens = result.get("tokens_estimated_from_int32_bin")
+        tokens = result.get("tokens_estimated_from_int32_bin") or result.get("tokens")
         if output_prefix and tokens:
             token_by_prefix[str(output_prefix) + "_text_document"] = float(tokens)
 
@@ -55,6 +55,42 @@ def _token_weights_from_manifest(prefix_list: list[str], manifest: dict) -> list
     return _normalize_weights(weights)
 
 
+def _target_mix_weights_from_manifest(prefix_list: list[str], manifest: dict) -> list[float] | None:
+    target_mix = manifest.get("target_mix") or {}
+    if not target_mix:
+        return None
+
+    prefix_to_category = {}
+    prefix_to_tokens = {}
+    category_totals = {}
+    for result in manifest.get("results") or []:
+        output_prefix = result.get("output_prefix")
+        category = result.get("mixture_category") or result.get("category")
+        tokens = result.get("tokens")
+        if not output_prefix or not category or not tokens:
+            continue
+
+        prefix = str(output_prefix) + "_text_document"
+        prefix_to_category[prefix] = str(category)
+        prefix_to_tokens[prefix] = float(tokens)
+        category_totals[str(category)] = category_totals.get(str(category), 0.0) + float(tokens)
+
+    if not prefix_to_category:
+        return None
+
+    raw_weights = []
+    for prefix in prefix_list:
+        category = prefix_to_category.get(prefix)
+        tokens = prefix_to_tokens.get(prefix)
+        category_total = category_totals.get(category or "", 0.0)
+        target_ratio = float(target_mix.get(category or "", 0.0))
+        if not category or not tokens or category_total <= 0 or target_ratio <= 0:
+            return None
+        raw_weights.append(target_ratio * (tokens / category_total))
+
+    return _normalize_weights(raw_weights)
+
+
 def _target_mix_weights_from_category_stats(
     prefix_list: list[str],
     category_stats_path: Path,
@@ -91,7 +127,9 @@ def load_data_blend(args: argparse.Namespace) -> tuple[list[str], list[float] |
             raise ValueError(f"{manifest_path} has no ok_prefixes")
         prefix_list = [str(prefix) for prefix in prefixes]
         category_stats_path = manifest_path.with_name("prefix_category_stats.json")
-        weights = _target_mix_weights_from_category_stats(prefix_list, category_stats_path)
+        weights = _target_mix_weights_from_manifest(prefix_list, manifest)
+        if weights is None:
+            weights = _target_mix_weights_from_category_stats(prefix_list, category_stats_path)
         if weights is None:
             weights = _token_weights_from_manifest(prefix_list, manifest)
         return (prefix_list, weights)