diff --git a/scripts/resume_stage1_2_pretrain_8192_8gpu_mbs14.sh b/scripts/resume_stage1_2_pretrain_8192_8gpu_mbs14.sh new file mode 100755 index 0000000..050355a --- /dev/null +++ b/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" diff --git a/training/megatron_bridge/laoyao_2b_moe_pretrain.py b/training/megatron_bridge/laoyao_2b_moe_pretrain.py index eb0edf2..329dc69 100755 --- a/training/megatron_bridge/laoyao_2b_moe_pretrain.py +++ b/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)