diff --git a/.gitignore b/.gitignore index ae67856..377aa36 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,10 @@ core.* dataset/pretrain/data/* !dataset/pretrain/data/.gitkeep +# Generated Megatron indexed validation data and GPTDataset caches. +# Rebuild from dataset/val/data/heldout_2p8k_packed_text.jsonl when needed. +dataset/val/megatron_8192_glm52/ + # Python/cache/editor __pycache__/ *.py[cod] diff --git a/README.md b/README.md index 8bf9c67..c4cbb64 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,9 @@ - `dataset/`: 预训练数据和验证集的构建、manifest、数据落盘位置。 - `model/`: 2B MoE 架构定义，优先用 NeMo/Megatron 配置表达。 - `training/`: 训练 recipe、评估 recipe、并行和优化超参。 -- `scripts/`: 在 g0033 上同步数据、拉起训练的 shell 入口。 +- `scripts/`: 数据下载、tokenization、训练恢复、Megatron 推理服务和 HF 导出 smoke 的 shell 入口。 +- `tools/`: checkpoint 检查、Megatron DCP 探测、Megatron-to-HF 导出和 HF custom model 推理工具。 +- `docs/`: 训练/导出过程中沉淀的问题记录。 ## 当前数据计划 @@ -34,7 +36,35 @@ /mnt/beegfs/yi/laoyao_2b_moe/tokenizer/glm5.2/tokenizer.json ``` -该 tokenizer 从 ModelScope zai-org/GLM-5.2 下载，当前解析出的 vocab size 为 154,820，tokenizer.json sha256 为 19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d。后续训练数据应按这个 tokenizer 重新统计/切分 token budget。 +该 tokenizer 从 ModelScope zai-org/GLM-5.2 下载，当前训练配置使用的 vocab size 为 154,856，tokenizer.json sha256 为 19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d。后续训练数据应按这个 tokenizer 重新统计/切分 token budget。 + +## g0050 训练状态入口 + +g0050 上的实际训练 repo 路径： + +```bash +/ssd/workspace/yi/laoyao_2b_moe +``` + +当前主训练入口为： + +```bash +bash scripts/resume_pretrain_8192_8gpu_mbs14.sh +``` + +该入口默认从 `runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints` 恢复，使用： + +- `seq_length=8192` +- `micro_batch_size=14` +- `global_batch_size=112` +- `tensor/pipeline/expert/context parallel = 1/1/1/1` +- `distributed optimizer + overlap grad reduce + overlap param gather` +- `full recompute, uniform, recompute_num_layers=1` +- `save_interval=2500` +- `eval_interval=15000` +- heldout validation prefix: `dataset/val/megatron_8192_glm52/heldout_2p8k_text_document` + +Megatron indexed validation 数据是本地派生产物，被 `.gitignore` 忽略；需要在目标机器上从 `dataset/val/data/heldout_2p8k_packed_text.jsonl` 重新生成或从训练机器同步。 ## g0033 路径约定 diff --git a/dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py b/dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py index 136a5a4..43e3ee0 100755 --- a/dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py +++ b/dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py @@ -148,6 +148,8 @@ def convert_one_parquet(task: tuple[str, argparse.Namespace]) -> dict[str, objec "status": "skip", "source": str(parquet_path), "output_prefix": str(output_prefix), + "bin_file": str(bin_file), + "idx_file": str(idx_file), "reason": "existing bin/idx", } @@ -330,7 +332,11 @@ def main() -> None: "total_split_source_docs": sum(int(r.get("split_source_docs", 0)) for r in results), "total_split_extra_docs": sum(int(r.get("split_extra_docs", 0)) for r in results), "total_tokens": sum(int(r.get("tokens", 0)) for r in results), - "ok_prefixes": [str(r["output_prefix"]) for r in results if r.get("status") in {"ok", "skip"}], + "ok_prefixes": [ + str(Path(str(r["bin_file"])).with_suffix("")) + for r in results + if r.get("status") in {"ok", "skip"} + ], } manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8") print(f"manifest={manifest_path}", flush=True) diff --git a/dataset/val/README.md b/dataset/val/README.md index b9dd29c..82799b2 100644 --- a/dataset/val/README.md +++ b/dataset/val/README.md @@ -19,3 +19,13 @@ - `build_stats.json`: 构建统计和来源分布。 这些样本来源包括 GSM8K、COIG-CQIA、UltraChat、HumanEval、MBPP、BBH logical deduction、ARC-Challenge、MMLU formal logic/logical fallacies、CEVAL 等。 + +## Megatron indexed validation + +g0050 训练时使用本地派生的 Megatron indexed validation prefix： + +```text +dataset/val/megatron_8192_glm52/heldout_2p8k_text_document +``` + +该目录下的 `.bin/.idx` 和 `cache/GPTDataset_indices/*.npy` 是运行派生产物，不进入 git。训练脚本通过 `--validation-prefix` 或默认值读取这个 prefix；如果目标机器没有该目录，需要先用 GLM-5.2 tokenizer 从 `heldout_2p8k_packed_text.jsonl` 重新生成 Megatron indexed dataset。 diff --git a/docs/megatron_to_hf_export_notes.md b/docs/megatron_to_hf_export_notes.md new file mode 100644 index 0000000..2e5e932 --- /dev/null +++ b/docs/megatron_to_hf_export_notes.md @@ -0,0 +1,111 @@ +# Megatron 到 HuggingFace 导出与推理修复记录 + +本文记录 Laoyao 2B MoE 从 Megatron distributed checkpoint 导出到 HuggingFace 自定义模型时遇到的问题、修复和当前结论。 + +## 当前结论 + +Megatron 原生 checkpoint 可以正常通过 Megatron text generation server 推理。导出到 HuggingFace 时，权重本身可以转换为 `safetensors`，但 HF 自定义模型必须显式对齐 Megatron 的模型细节，并且默认禁用 KV cache。 + +当前最重要的修复是： + +- `config.use_cache = false`，因为 HF 自定义模型没有实现 KV cache。 +- `rms_norm_eps = 1e-5`，与 Megatron runtime 参数一致。 +- RoPE 使用 Megatron 默认的 non-interleaved/NeoX half-split 形式。 +- MoE router 模拟 Megatron 的 `moe_expert_capacity_factor=1.25` token dropping 行为。 +- HF dynamic module 使用相对导入，保证 `AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)` 能从导出目录独立加载。 + +## 已确认的问题 + +### KV cache + +当前 HF 模型没有实现 `past_key_values`、`cache_position` 和 per-layer K/V cache 管理。Transformers `generate()` 默认倾向于使用 cache；如果不禁用 cache，后续 step 可能只输入新 token，而模型会把该 token 当作完整上下文从 position 0 计算，导致上下文丢失和异常重复。 + +现象： + +- `model.generate()` 默认 cache 路径输出 `Paris, and the 1. The...`、中文重复 `在在在...`。 +- 设置 `use_cache=False` 后，输出恢复为连续句子，例如 `Paris. The capital of Belgium...`。 + +修复： + +- `configuration_laoyao_moe.py` 默认 `use_cache=False`。 +- `convert_laoyao_dcp_to_hf.py` 写出的 `config.json` 包含 `"use_cache": false`。 +- `generate_laoyao_hf.py` 调用 `model.generate(..., use_cache=False)`。 + +### RMSNorm epsilon + +Megatron server 日志显示： + +```text +layernorm_epsilon ............................... 1e-05 +``` + +HF 初始实现使用 `1e-6`，已修正为 `1e-5`。这个差异会造成逐层数值漂移，尤其可能影响 MoE router 的 top-k 边界。 + +### RoPE layout + +Megatron runtime 参数： + +```text +rotary_interleaved .............................. False +rotary_base ..................................... 10000 +rotary_percent .................................. 1.0 +``` + +Megatron 源码在 `rotary_interleaved=False` 时使用： + +- `emb = torch.cat((freqs, freqs), dim=-1)` +- `_rotate_half(x) = torch.cat((-x2, x1), dim=-1)`，其中 `x1, x2 = torch.chunk(x, 2, dim=-1)` + +因此 HF 侧应使用 NeoX half-split RoPE，而不是 even/odd interleaved RoPE。 + +### MoE expert capacity + +Megatron runtime 参数包含： + +```text +moe_expert_capacity_factor ...................... 1.25 +moe_token_drop_policy ........................... probs +moe_pad_expert_input_to_capacity ................ False +``` + +Megatron `TopKRouter` 在 `moe_expert_capacity_factor` 非空时会调用 `apply_router_token_dropping`，按 expert 维度保留 routing prob 较高的 assignment，超出 capacity 的 token-expert assignment 会被 mask 掉。 + +HF 初始实现是 dropless top-k routing，已补齐 sparse routing map + capacity mask 逻辑。注意：这不是本次 `1/在` 重复的主因，但它是和 Megatron 对齐必须处理的行为。 + +### HF dynamic module 导入 + +HF export 目录中的 `modeling_laoyao_moe.py` 必须使用： + +```python +from .configuration_laoyao_moe import LaoyaoMoeConfig +``` + +不能使用裸导入 `from configuration_laoyao_moe import ...`，否则脱离工具目录后 `AutoModelForCausalLM.from_pretrained` 会失败。 + +## 脚本入口 + +- `tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py` + - 从 Megatron DCP checkpoint 读取权重。 + - 写出 HF `model.safetensors`、`config.json`、tokenizer 文件和 custom code。 + +- `tools/hf_laoyao_moe/generate_laoyao_hf.py` + - 从 HF export 目录加载模型。 + - 默认 `use_cache=False` 跑生成。 + - 支持 `--print-token-details` 和 `--top-k-debug` 做逐 token 调试。 + +- `scripts/serve_laoyao_megatron.sh` + - 启动 Megatron 原生 text generation server。 + +- `scripts/laoyao_megatron_inference_setup.sh` + - 容器启动时热补丁 Megatron server： + - 修 `inference_max_sequence_length` 参数兼容。 + - 修 token list concat。 + - 支持 generated-only logprobs，用于调试。 + +## 当前限制 + +- HF export 能正确 full-context 生成，但没有高效 KV cache。 +- HF 模型没有实现 `_reorder_cache()`，不支持 beam search cache 重排。 +- MoE 和 attention 目前是简单 PyTorch 实现，不是高性能推理实现。 +- Megatron 原生 server 仍是更可信的推理基线。 + diff --git a/scripts/README.md b/scripts/README.md index 44518ab..e20018a 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -7,6 +7,13 @@ - `train_megatron_bridge_2b_moe.sh`: 当前主训练入口，使用 NeMo 26.06 镜像中的 Megatron-Bridge。 - `train_nemo_megatron_2b_moe.sh`: NeMo/Megatron 训练入口占位，包含 image、mount、路径检查。 - `g0050_download_and_setup_from_modelscope.sh`: 在 g0050 上一键准备训练环境并从 ModelScope 下载未 tokenize parquet 数据。 +- `g0050_wait_and_tokenize_glm52_8192.sh`: 等待 g0050 ModelScope parquet 下载完成后，按 GLM-5.2 tokenizer 和 `max_seq_len=8192` 转成 Megatron indexed dataset。 +- `resume_pretrain_8192_8gpu_mbs14.sh`: g0050 当前主恢复训练入口，默认 8 卡 DP、`mbs=14`、`gbs=112`、full recompute、distributed optimizer、2500 iter 保存、15000 iter validation。 +- `serve_laoyao_megatron.sh` / `serve_laoyao_megatron_inner.sh`: 启动 Megatron 原生 text generation server，用于 checkpoint 推理 smoke。 +- `laoyao_megatron_inference_setup.sh`: 容器内热补丁 Megatron inference server 的参数兼容和返回格式问题。 +- `check_laoyao_megatron_inference_ready.sh`: 检查 checkpoint、端口、GPU 和相关容器状态。 +- `query_laoyao_megatron_server.sh`: 对 Megatron inference server 发送简单 prompt。 +- `stop_laoyao_megatron_server.sh`: 停止 Megatron inference server 容器。 ## g0050 下载与部署 @@ -32,3 +39,23 @@ GIT_REPO_URL=https://yi_lu:@git.deeepseek.net/yi_lu/laoyao_2b_moe.git \ MODELSCOPE_API_TOKEN=ms-... \ bash scripts/g0050_download_and_setup_from_modelscope.sh ``` + +## g0050 当前训练恢复 + +在 g0050 上： + +```bash +cd /ssd/workspace/yi/laoyao_2b_moe +bash scripts/resume_pretrain_8192_8gpu_mbs14.sh +``` + +默认配置与当前长跑实验一致： + +- data manifest: `/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_glm52_direct_v1/manifest.json` +- validation prefix: `dataset/val/megatron_8192_glm52/heldout_2p8k_text_document` +- train iters: `184209` +- `seq_length=8192` +- `micro_batch_size=14`, `global_batch_size=112` +- `save_interval=2500`, `eval_interval=15000`, `eval_iters=10` +- `--use-distributed-optimizer --overlap-grad-reduce --overlap-param-gather` +- `--recompute-granularity full --recompute-method uniform --recompute-num-layers 1` diff --git a/scripts/check_laoyao_megatron_inference_ready.sh b/scripts/check_laoyao_megatron_inference_ready.sh new file mode 100755 index 0000000..97578ed --- /dev/null +++ b/scripts/check_laoyao_megatron_inference_ready.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT=${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe} +CHECKPOINT=${CHECKPOINT:-${REPO_ROOT}/runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints} +SERVER_PORT=${SERVER_PORT:-5000} + +cd "${REPO_ROOT}" + +echo "== checkpoint ==" +test -f "${CHECKPOINT}/latest_checkpointed_iteration.txt" +cat "${CHECKPOINT}/latest_checkpointed_iteration.txt" +ls -ld "${CHECKPOINT}"/iter_* | tail -5 + +echo +echo "== port ${SERVER_PORT} ==" +if command -v ss >/dev/null 2>&1; then + ss -ltnp | grep ":${SERVER_PORT} " || true +else + netstat -ltnp 2>/dev/null | grep ":${SERVER_PORT} " || true +fi + +echo +echo "== gpu ==" +nvidia-smi --query-gpu=index,name,memory.used,memory.total,utilization.gpu,power.draw --format=csv,noheader,nounits + +echo +echo "== existing containers ==" +docker ps --format 'table {{.Names}}\t{{.Image}}\t{{.Status}}' | grep -E 'laoyao|sglang|vllm|NAME' || true diff --git a/scripts/g0050_wait_and_tokenize_glm52_8192.sh b/scripts/g0050_wait_and_tokenize_glm52_8192.sh new file mode 100755 index 0000000..2fb96b5 --- /dev/null +++ b/scripts/g0050_wait_and_tokenize_glm52_8192.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT="${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}" +DATA_ROOT="${DATA_ROOT:-/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset}" +LOG_DIR="${LOG_DIR:-$DATA_ROOT/logs}" +SLEEP_SECONDS="${SLEEP_SECONDS:-300}" +EXPECTED_MIN_PARQUETS="${EXPECTED_MIN_PARQUETS:-1600}" + +mkdir -p "$LOG_DIR" + +echo "[wait] started at $(date -Is)" +echo "[wait] repo=$REPO_ROOT" +echo "[wait] data_root=$DATA_ROOT" + +while true; do + if ! docker ps --format '{{.ID}} {{.Image}} {{.Command}} {{.Names}}' \ + | grep -q 'modelscope.*download.*laoyao_2b_moe_pretrain_parquet_20260702'; then + if ! pgrep -af 'modelscope download --dataset eigentom/laoyao_2b_moe_pretrain_parquet_20260702' >/dev/null; then + break + fi + fi + echo "[wait] modelscope download still running at $(date -Is)" + du -sh "$DATA_ROOT" || true + find "$DATA_ROOT" -name '*.parquet' | wc -l || true + sleep "$SLEEP_SECONDS" +done + +echo "[wait] download appears complete at $(date -Is)" +du -sh "$DATA_ROOT" || true +parquet_count="$(find "$DATA_ROOT" -name '*.parquet' | wc -l | tr -d ' ')" +echo "[wait] parquet_count=$parquet_count expected_min=$EXPECTED_MIN_PARQUETS" +if [[ "$parquet_count" -lt "$EXPECTED_MIN_PARQUETS" ]]; then + echo "[wait] refusing to tokenize: parquet_count is below expected minimum" >&2 + exit 3 +fi + +cd "$REPO_ROOT" + +export REPO_ROOT +export IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}" +export SOURCE_DATA_DIRS="${SOURCE_DATA_DIRS:-$DATA_ROOT/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:$DATA_ROOT/train/logic_topup_proof_pile_17b_v1_20260701}" +export WORK_DIR="${WORK_DIR:-$DATA_ROOT/megatron_bridge/pretrain_8192_glm52_direct_v1}" +export TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}" +export OUTPUT_PREFIX_PREFIX="${OUTPUT_PREFIX_PREFIX:-laoyao_2b_moe_glm52_8192}" +export TEXT_KEY="${TEXT_KEY:-text}" +export PARALLEL_FILES="${PARALLEL_FILES:-4}" +export WORKERS_PER_FILE="${WORKERS_PER_FILE:-8}" +export BATCH_SIZE="${BATCH_SIZE:-8192}" +export CHUNKSIZE="${CHUNKSIZE:-128}" +export MAX_FILES="${MAX_FILES:-0}" +export MAX_DOCS="${MAX_DOCS:-0}" +export MAX_SEQ_LEN="${MAX_SEQ_LEN:-8192}" +export MIN_FREE_GB="${MIN_FREE_GB:-800}" +export OVERWRITE="${OVERWRITE:-0}" + +echo "[tokenize] started at $(date -Is)" +echo "[tokenize] work_dir=$WORK_DIR" +echo "[tokenize] max_seq_len=$MAX_SEQ_LEN parallel_files=$PARALLEL_FILES workers_per_file=$WORKERS_PER_FILE" + +bash scripts/preprocess_megatron_bridge_pretrain_direct.sh + +echo "[tokenize] finished at $(date -Is)" +du -sh "$WORK_DIR" || true +find "$WORK_DIR" -maxdepth 1 -type f | wc -l || true +ls -lh "$WORK_DIR/manifest.json" || true diff --git a/scripts/laoyao_megatron_inference_setup.sh b/scripts/laoyao_megatron_inference_setup.sh new file mode 100755 index 0000000..5ee21c9 --- /dev/null +++ b/scripts/laoyao_megatron_inference_setup.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +set -euo pipefail + +MEGATRON_PATH="${MEGATRON_PATH:-/opt/Megatron-Bridge/3rdparty/Megatron-LM}" +SERVER_PY="${MEGATRON_PATH}/tools/run_text_generation_server.py" +ENGINE_PY="${MEGATRON_PATH}/megatron/core/inference/text_generation_server/run_mcore_engine.py" + +python3 - <<'PY' +import importlib.util +missing = [name for name in ["flask", "flask_restful"] if importlib.util.find_spec(name) is None] +if missing: + raise SystemExit(f"missing python packages: {missing}; install with: pip install flask-restful") +print("flask dependencies ok") +PY + +python3 - "$SERVER_PY" "$ENGINE_PY" <<'PY' +from pathlib import Path +import sys + +server = Path(sys.argv[1]) +engine = Path(sys.argv[2]) + +text = server.read_text() +old = "inference_context = StaticInferenceContext(args.inference_max_requests, args.inference_max_sequence_length)" +new = "inference_context = StaticInferenceContext(args.inference_max_requests, getattr(args, 'inference_max_sequence_length', getattr(args, 'inference_max_seq_length', args.seq_length)))" +if old in text: + text = text.replace(old, new) + server.write_text(text) + print(f"patched {server}: inference_max_sequence_length fallback") +elif new in text: + print(f"already patched {server}: inference_max_sequence_length fallback") +else: + print(f"warning: expected inference_context line not found in {server}") + +text = engine.read_text() +old = '"tokens": [x.prompt_tokens + x.generated_tokens.tolist() for x in result],' +new = '"tokens": [x.prompt_tokens.tolist() + x.generated_tokens if hasattr(x.prompt_tokens, "tolist") else x.prompt_tokens + x.generated_tokens for x in result],' +if old in text: + text = text.replace(old, new) + engine.write_text(text) + print(f"patched {engine}: token list concatenation") +elif "x.prompt_tokens.tolist() + x.generated_tokens" in text: + print(f"already patched {engine}: token list concatenation") +else: + print(f"warning: expected token concat line not found in {engine}") + +text = engine.read_text() +old = " skip_prompt_log_probs=False," +new = " skip_prompt_log_probs=True," +if old in text: + text = text.replace(old, new) + engine.write_text(text) + print(f"patched {engine}: skip prompt logprobs") +elif new in text: + print(f"already patched {engine}: skip prompt logprobs") +else: + print(f"warning: expected skip_prompt_log_probs line not found in {engine}") + +text = engine.read_text() +old = " response_logprobs = [x.prompt_log_probs + x.generated_log_probs for x in result]" +new = " response_logprobs = [x.generated_log_probs for x in result]" +if old in text: + text = text.replace(old, new) + engine.write_text(text) + print(f"patched {engine}: generated-only logprobs response") +elif new in text: + print(f"already patched {engine}: generated-only logprobs response") +else: + print(f"warning: expected logprobs response line not found in {engine}") +PY + +echo "Megatron inference patch setup complete." diff --git a/scripts/query_laoyao_megatron_server.sh b/scripts/query_laoyao_megatron_server.sh new file mode 100755 index 0000000..8bf1d47 --- /dev/null +++ b/scripts/query_laoyao_megatron_server.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +set -euo pipefail + +SERVER_URL=${SERVER_URL:-http://127.0.0.1:5000/api} + +curl -sS -X PUT "${SERVER_URL}" \ + -H "Content-Type: application/json" \ + -d '{ + "prompts": [ + "The capital of France is", + "中国的首都是", + "1 + 1 =" + ], + "tokens_to_generate": 64, + "temperature": 0.7, + "top_k": 0, + "top_p": 0.9 + }' diff --git a/scripts/resume_pretrain_8192_8gpu_mbs14.sh b/scripts/resume_pretrain_8192_8gpu_mbs14.sh new file mode 100755 index 0000000..335255d --- /dev/null +++ b/scripts/resume_pretrain_8192_8gpu_mbs14.sh @@ -0,0 +1,96 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT="${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}" +IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}" +RUN_NAME="${RUN_NAME:-pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000}" +RUN_DIR="${RUN_DIR:-${REPO_ROOT}/runs/${RUN_NAME}}" +CKPT_DIR="${CKPT_DIR:-${RUN_DIR}/checkpoints}" +TENSORBOARD_DIR="${TENSORBOARD_DIR:-${RUN_DIR}/tensorboard}" +DATA_MANIFEST="${DATA_MANIFEST:-/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_glm52_direct_v1/manifest.json}" +VALIDATION_PREFIX="${VALIDATION_PREFIX:-${REPO_ROOT}/dataset/val/megatron_8192_glm52/heldout_2p8k_text_document}" +TRAIN_ITERS="${TRAIN_ITERS:-184209}" +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}" +LOG_FILE="${LOG_FILE:-/tmp/laoyao_pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000_resume.log}" +CONTAINER_NAME="${CONTAINER_NAME:-laoyao_pretrain_resume}" + +if [[ ! -d "${REPO_ROOT}" ]]; then + echo "ERROR: repo dir not found: ${REPO_ROOT}" >&2 + exit 1 +fi +if [[ ! -f "${DATA_MANIFEST}" ]]; then + echo "ERROR: data manifest not found: ${DATA_MANIFEST}" >&2 + exit 1 +fi +if [[ ! -d "${CKPT_DIR}" ]]; then + echo "ERROR: checkpoint dir not found: ${CKPT_DIR}" >&2 + exit 1 +fi +if [[ ! -f "${VALIDATION_PREFIX}.idx" ]]; then + echo "WARNING: validation prefix not found: ${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 "${CKPT_DIR}" "${TENSORBOARD_DIR}" + +echo "Resuming Laoyao 2B MoE training" +echo "repo: ${REPO_ROOT}" +echo "checkpoint_dir: ${CKPT_DIR}" +echo "latest_checkpointed_iteration: $(cat "${CKPT_DIR}/latest_checkpointed_iteration.txt" 2>/dev/null || echo unknown)" +echo "log_file: ${LOG_FILE}" + +cd "${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 /ssd:/ssd \ + -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}' \ + --save-dir '${CKPT_DIR}' \ + --load-dir '${CKPT_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/scripts/serve_laoyao_megatron.sh b/scripts/serve_laoyao_megatron.sh new file mode 100755 index 0000000..0f239b4 --- /dev/null +++ b/scripts/serve_laoyao_megatron.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Host-side launcher for the Megatron-LM text generation server. +# Run on g0050 from /ssd/workspace/yi/laoyao_2b_moe. + +REPO_ROOT=${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe} +IMAGE=${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4} +SERVER_PORT=${SERVER_PORT:-5000} +MASTER_PORT=${MASTER_PORT:-6000} +CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0} +CONTAINER_NAME=${CONTAINER_NAME:-laoyao_megatron_server} +DETACH=${DETACH:-0} +LOG_DIR=${LOG_DIR:-${REPO_ROOT}/runs/megatron_inference} + +mkdir -p "${LOG_DIR}" + +DOCKER_ARGS=( + --gpus all + --network host + --ipc host + --shm-size 64g + --ulimit memlock=-1 + --ulimit stack=67108864 + --name "${CONTAINER_NAME}" + -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" + -e SERVER_PORT="${SERVER_PORT}" + -e MASTER_PORT="${MASTER_PORT}" + -v "${REPO_ROOT}:/work" + -w /work +) + +docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true + +if [[ "${DETACH}" == "1" ]]; then + docker run -d "${DOCKER_ARGS[@]}" "${IMAGE}" \ + bash scripts/serve_laoyao_megatron_inner.sh + echo "Started ${CONTAINER_NAME}" + echo "Logs: docker logs -f ${CONTAINER_NAME}" +else + docker run --rm "${DOCKER_ARGS[@]}" "${IMAGE}" \ + bash scripts/serve_laoyao_megatron_inner.sh 2>&1 | tee "${LOG_DIR}/server_$(date +%Y%m%d_%H%M%S).log" +fi diff --git a/scripts/serve_laoyao_megatron_inner.sh b/scripts/serve_laoyao_megatron_inner.sh new file mode 100755 index 0000000..41016b3 --- /dev/null +++ b/scripts/serve_laoyao_megatron_inner.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Container-side launcher. It applies the Megatron inference server patches +# documented in pretrain_kaiyuan2b before starting the server, because docker +# --rm containers do not retain in-place edits across runs. + +bash scripts/laoyao_megatron_inference_setup.sh + +MEGATRON_PATH=${MEGATRON_PATH:-/opt/Megatron-Bridge/3rdparty/Megatron-LM} +CHECKPOINT=${CHECKPOINT:-/work/runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints} +TOKENIZER_PATH=${TOKENIZER_PATH:-/work/tokenizer/glm5.2} +SERVER_PORT=${SERVER_PORT:-5000} +MASTER_PORT=${MASTER_PORT:-6000} +INFERENCE_MAX_SEQ_LENGTH=${INFERENCE_MAX_SEQ_LENGTH:-8192} + +DISTRIBUTED_ARGS=( + --nproc_per_node 1 + --nnodes 1 + --node_rank 0 + --master_addr localhost + --master_port "${MASTER_PORT}" +) + +export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} + +torchrun "${DISTRIBUTED_ARGS[@]}" "${MEGATRON_PATH}/tools/run_text_generation_server.py" \ + --load "${CHECKPOINT}" \ + --tensor-model-parallel-size 1 \ + --pipeline-model-parallel-size 1 \ + --num-layers 12 \ + --hidden-size 1536 \ + --ffn-hidden-size 4608 \ + --num-attention-heads 24 \ + --num-query-groups 4 \ + --group-query-attention \ + --seq-length 8192 \ + --max-position-embeddings 8192 \ + --position-embedding-type rope \ + --rotary-base 10000 \ + --swiglu \ + --disable-bias-linear \ + --normalization RMSNorm \ + --make-vocab-size-divisible-by 1 \ + --vocab-size 154856 \ + --tokenizer-type HuggingFaceTokenizer \ + --tokenizer-model "${TOKENIZER_PATH}" \ + --bf16 \ + --micro-batch-size 1 \ + --num-experts 12 \ + --moe-ffn-hidden-size 6144 \ + --moe-layer-freq '[0,0,1,0,1,0,1,0,1,0,1,0]' \ + --moe-router-topk 4 \ + --moe-router-load-balancing-type aux_loss \ + --moe-aux-loss-coeff 0.02 \ + --moe-z-loss-coeff 0.001 \ + --moe-token-dispatcher-type alltoall \ + --moe-expert-capacity-factor 1.25 \ + --moe-grouped-gemm \ + --no-load-rng \ + --no-load-optim \ + --exit-on-missing-checkpoint \ + --inference-max-requests 1 \ + --inference-max-seq-length "${INFERENCE_MAX_SEQ_LENGTH}" \ + --port "${SERVER_PORT}" diff --git a/scripts/stop_laoyao_megatron_server.sh b/scripts/stop_laoyao_megatron_server.sh new file mode 100755 index 0000000..88b2b80 --- /dev/null +++ b/scripts/stop_laoyao_megatron_server.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash +set -euo pipefail + +CONTAINER_NAME=${CONTAINER_NAME:-laoyao_megatron_server} +docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true +echo "Stopped ${CONTAINER_NAME}" diff --git a/tokenizer/README.md b/tokenizer/README.md index 57c2eeb..1eca9c5 100644 --- a/tokenizer/README.md +++ b/tokenizer/README.md @@ -17,7 +17,7 @@ glm5.2/generation_config.json ```text tokenizer.json sha256: 19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d -vocab_size: 154820 +vocab_size used by training config: 154856 config.model_type: glm_moe_dsa source: https://modelscope.ai/models/zai-org/GLM-5.2/files ``` diff --git a/tools/README.md b/tools/README.md new file mode 100644 index 0000000..6de6baa --- /dev/null +++ b/tools/README.md @@ -0,0 +1,35 @@ +# Tools + +本目录保存 Laoyao 2B MoE 训练后的 checkpoint 检查、Megatron DCP 探测、HF 导出和推理调试工具。 + +## Megatron checkpoint 检查 + +- `laoyao_dcp_probe.py`: 用 `torch.distributed.checkpoint` 读取 Megatron distributed checkpoint metadata 和单个 tensor，确认 DCP checkpoint 可以被 CPU/gloo 进程读取。 +- `inspect_laoyao_ckpt_metadata.py`: 打印 checkpoint 的 `metadata.json`、`common.pt` 和 `run_config.yaml` 关键信息。 +- `dump_laoyao_model_keys.py`: dump 模型 key，用于对齐 Megatron 权重名和 HF custom model 权重名。 +- `laoyao_cpu_infer_probe.py` / `laoyao_cpu_generate_from_distcp.py`: 早期 CPU 推理探测脚本，主要用于定位 checkpoint 读取和模型定义问题。 + +## HuggingFace 导出 + +`hf_laoyao_moe/` 是 Megatron distributed checkpoint 到 HuggingFace custom model 的导出适配层。 + +核心入口： + +```bash +python3 tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py \ + --checkpoint-dir runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints/iter_0107500 \ + --tokenizer-dir tokenizer/glm5.2 \ + --output-dir runs/hf_exports/iter_0107500 +``` + +生成测试： + +```bash +python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --device cuda \ + --max-new-tokens 32 \ + --prompt "The capital of France is" +``` + +注意：HF custom model 当前没有实现 KV cache，因此导出配置和生成脚本默认 `use_cache=false`。更多细节见 `docs/megatron_to_hf_export_notes.md`。 diff --git a/tools/dump_laoyao_model_keys.py b/tools/dump_laoyao_model_keys.py new file mode 100644 index 0000000..cee5735 --- /dev/null +++ b/tools/dump_laoyao_model_keys.py @@ -0,0 +1,33 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse + +from torch.distributed.checkpoint.metadata import Metadata +from torch.distributed.checkpoint.filesystem import FileSystemReader + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint_dir") + parser.add_argument("--limit", type=int, default=240) + args = parser.parse_args() + + reader = FileSystemReader(args.checkpoint_dir) + metadata: Metadata = reader.read_metadata() + count = 0 + skip_prefixes = ("optimizer.", "opt_param_scheduler", "rng_state", "iteration") + for key in sorted(metadata.state_dict_metadata): + if key.startswith(skip_prefixes): + continue + meta = metadata.state_dict_metadata[key] + size = getattr(meta, "size", None) + print(f"{key}\t{tuple(size) if size is not None else '?'}") + count += 1 + if args.limit and count >= args.limit: + break + print(f"printed={count}") + + +if __name__ == "__main__": + main() diff --git a/tools/hf_laoyao_moe/README.md b/tools/hf_laoyao_moe/README.md new file mode 100644 index 0000000..f133a22 --- /dev/null +++ b/tools/hf_laoyao_moe/README.md @@ -0,0 +1,69 @@ +# HF Laoyao MoE 工具说明 + +本目录包含把 Laoyao 2B MoE Megatron checkpoint 导出为 HuggingFace custom model 的脚本。 + +## 导出 + +```bash +python3 tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py \ + --checkpoint-dir runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints/iter_0107500 \ + --tokenizer-dir tokenizer/glm5.2 \ + --output-dir runs/hf_exports/iter_0107500 +``` + +导出目录包含： + +- `model.safetensors` +- `config.json` +- `configuration_laoyao_moe.py` +- `modeling_laoyao_moe.py` +- tokenizer 文件 + +## 生成测试 + +```bash +python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --device cuda \ + --fix-mistral-regex \ + --max-new-tokens 32 \ + --prompt "The capital of France is" +``` + +当前 HF custom model 没有实现 KV cache，因此生成脚本强制 `use_cache=False`。不要删除这个设置；否则 Transformers 默认 cache path 会导致上下文丢失和异常重复。 + +## 调试 + +打印逐 token 结果： + +```bash +python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --device cuda \ + --fix-mistral-regex \ + --max-new-tokens 8 \ + --print-token-details \ + --prompt "The capital of France is" +``` + +打印每步 top-k： + +```bash +python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --device cuda \ + --fix-mistral-regex \ + --max-new-tokens 5 \ + --top-k-debug 8 \ + --prompt "The capital of France is" +``` + +## 已对齐的 Megatron 细节 + +- `rms_norm_eps=1e-5` +- RoPE `rotary_interleaved=False`，即 NeoX half-split layout +- GQA QKV group-major layout +- tied embedding/output head +- MoE top-k router with `moe_expert_capacity_factor=1.25` +- `use_cache=false` + diff --git a/tools/hf_laoyao_moe/configuration_laoyao_moe.py b/tools/hf_laoyao_moe/configuration_laoyao_moe.py new file mode 100755 index 0000000..85358d1 --- /dev/null +++ b/tools/hf_laoyao_moe/configuration_laoyao_moe.py @@ -0,0 +1,41 @@ +from transformers import PretrainedConfig + + +class LaoyaoMoeConfig(PretrainedConfig): + model_type = "laoyao_moe" + + def __init__( + self, + vocab_size=154856, + hidden_size=1536, + num_hidden_layers=12, + num_attention_heads=24, + num_key_value_heads=4, + intermediate_size=4608, + moe_intermediate_size=6144, + num_experts=12, + num_experts_per_tok=4, + moe_expert_capacity_factor=1.25, + moe_layers=(2, 4, 6, 8, 10), + max_position_embeddings=8192, + rope_theta=10000.0, + rms_norm_eps=1e-5, + tie_word_embeddings=True, + use_cache=False, + **kwargs, + ): + super().__init__(tie_word_embeddings=tie_word_embeddings, use_cache=use_cache, **kwargs) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_experts = num_experts + self.num_experts_per_tok = num_experts_per_tok + self.moe_expert_capacity_factor = moe_expert_capacity_factor + self.moe_layers = list(moe_layers) + self.max_position_embeddings = max_position_embeddings + self.rope_theta = rope_theta + self.rms_norm_eps = rms_norm_eps diff --git a/tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py b/tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py new file mode 100755 index 0000000..08139e9 --- /dev/null +++ b/tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +import shutil +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from safetensors.torch import save_file + + +CONFIG = { + "architectures": ["LaoyaoMoeForCausalLM"], + "auto_map": { + "AutoConfig": "configuration_laoyao_moe.LaoyaoMoeConfig", + "AutoModelForCausalLM": "modeling_laoyao_moe.LaoyaoMoeForCausalLM", + }, + "model_type": "laoyao_moe", + "vocab_size": 154856, + "hidden_size": 1536, + "num_hidden_layers": 12, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "intermediate_size": 4608, + "moe_intermediate_size": 6144, + "num_experts": 12, + "num_experts_per_tok": 4, + "moe_expert_capacity_factor": 1.25, + "moe_layers": [2, 4, 6, 8, 10], + "max_position_embeddings": 8192, + "rope_theta": 10000.0, + "rms_norm_eps": 1e-5, + "tie_word_embeddings": True, + "use_cache": False, + "torch_dtype": "bfloat16", +} + + +def init_dist() -> None: + if not dist.is_initialized(): + dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29691", rank=0, world_size=1) + + +def load_dcp_tensors(checkpoint_dir: Path) -> dict[str, torch.Tensor]: + init_dist() + metadata = dcp.FileSystemReader(str(checkpoint_dir)).read_metadata() + state = {} + for key, meta in metadata.state_dict_metadata.items(): + if key.startswith(("optimizer.", "opt_param_scheduler", "rng_state", "iteration")): + continue + if not hasattr(meta, "size"): + continue + state[key] = torch.empty(tuple(meta.size), dtype=meta.properties.dtype) + print(f"[load] tensors={len(state)} checkpoint={checkpoint_dir}", flush=True) + dcp.load(state, checkpoint_id=str(checkpoint_dir)) + return state + + +def convert_state(src: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + dst = { + "model.embed_tokens.weight": src["embedding.word_embeddings.weight"].contiguous(), + "model.norm.weight": src["decoder.final_layernorm.weight"].contiguous(), + } + for layer in range(CONFIG["num_hidden_layers"]): + sp = f"decoder.layers.{layer}" + dp = f"model.layers.{layer}" + dst[f"{dp}.input_layernorm.weight"] = src[f"{sp}.self_attention.linear_qkv.layer_norm_weight"].contiguous() + dst[f"{dp}.self_attn.qkv_proj.weight"] = src[f"{sp}.self_attention.linear_qkv.weight"].contiguous() + dst[f"{dp}.self_attn.o_proj.weight"] = src[f"{sp}.self_attention.linear_proj.weight"].contiguous() + if layer in CONFIG["moe_layers"]: + dst[f"{dp}.pre_mlp_layernorm.weight"] = src[f"{sp}.pre_mlp_layernorm.weight"].contiguous() + dst[f"{dp}.mlp.router.weight"] = src[f"{sp}.mlp.router.weight"].contiguous() + dst[f"{dp}.mlp.experts_fc1"] = src[f"{sp}.mlp.experts.experts.linear_fc1.weight"].contiguous() + dst[f"{dp}.mlp.experts_fc2"] = src[f"{sp}.mlp.experts.experts.linear_fc2.weight"].contiguous() + else: + dst[f"{dp}.pre_mlp_layernorm.weight"] = src[f"{sp}.mlp.linear_fc1.layer_norm_weight"].contiguous() + dst[f"{dp}.mlp.fc1.weight"] = src[f"{sp}.mlp.linear_fc1.weight"].contiguous() + dst[f"{dp}.mlp.fc2.weight"] = src[f"{sp}.mlp.linear_fc2.weight"].contiguous() + return dst + + +def copy_tokenizer(tokenizer_dir: Path, output_dir: Path) -> None: + for path in tokenizer_dir.iterdir(): + if path.is_file(): + shutil.copy2(path, output_dir / path.name) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint-dir", required=True) + parser.add_argument("--tokenizer-dir", default="tokenizer/glm5.2") + parser.add_argument("--output-dir", required=True) + args = parser.parse_args() + + checkpoint_dir = Path(args.checkpoint_dir) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + src = load_dcp_tensors(checkpoint_dir) + dst = convert_state(src) + print(f"[save] hf_tensors={len(dst)} output={output_dir}", flush=True) + save_file(dst, str(output_dir / "model.safetensors"), metadata={"format": "pt"}) + + copy_tokenizer(Path(args.tokenizer_dir), output_dir) + (output_dir / "config.json").write_text(json.dumps(CONFIG, indent=2, ensure_ascii=False) + "\n") + shutil.copy2(Path(__file__).with_name("configuration_laoyao_moe.py"), output_dir / "configuration_laoyao_moe.py") + shutil.copy2(Path(__file__).with_name("modeling_laoyao_moe.py"), output_dir / "modeling_laoyao_moe.py") + print("[done]", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/hf_laoyao_moe/generate_laoyao_hf.py b/tools/hf_laoyao_moe/generate_laoyao_hf.py new file mode 100755 index 0000000..761d98a --- /dev/null +++ b/tools/hf_laoyao_moe/generate_laoyao_hf.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse + +import torch +import torch.nn.functional as F +from transformers import AutoModelForCausalLM, AutoTokenizer + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model-dir", required=True) + parser.add_argument("--prompt", action="append", default=None) + parser.add_argument("--max-new-tokens", type=int, default=32) + parser.add_argument("--device", default="cpu") + parser.add_argument("--fix-mistral-regex", action="store_true") + parser.add_argument("--print-token-details", action="store_true") + parser.add_argument("--top-k-debug", type=int, default=0) + args = parser.parse_args() + + tokenizer = AutoTokenizer.from_pretrained( + args.model_dir, + trust_remote_code=True, + fix_mistral_regex=args.fix_mistral_regex, + ) + model = AutoModelForCausalLM.from_pretrained( + args.model_dir, + trust_remote_code=True, + torch_dtype=torch.bfloat16, + ) + model.to(args.device) + model.eval() + prompts = args.prompt or ["The capital of France is", "中国的首都是"] + for prompt in prompts: + inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(args.device) + with torch.inference_mode(): + output_ids = model.generate( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + use_cache=False, + pad_token_id=tokenizer.eos_token_id, + ) + print("PROMPT:", prompt) + print(tokenizer.decode(output_ids[0], skip_special_tokens=False)) + if args.print_token_details: + print("GENERATE_OUTPUT_TOKENS:", output_ids[0].tolist()) + ids = inputs["input_ids"].clone() + generated = [] + logprobs = [] + pieces = [] + with torch.inference_mode(): + for _ in range(args.max_new_tokens): + logits = model(input_ids=ids).logits[:, -1, :] + logp = F.log_softmax(logits.float(), dim=-1) + next_id = torch.argmax(logits, dim=-1) + generated.append(int(next_id.item())) + logprobs.append(float(logp[0, next_id].item())) + pieces.append(tokenizer.decode([int(next_id.item())], skip_special_tokens=False)) + ids = torch.cat([ids, next_id[:, None]], dim=1) + print("INPUT_TOKENS:", inputs["input_ids"][0].tolist()) + print("GENERATED_TOKENS:", generated) + print("GENERATED_PIECES:", pieces) + print("GENERATED_LOGPROBS:", logprobs) + if args.top_k_debug > 0: + ids = inputs["input_ids"].clone() + with torch.inference_mode(): + for step in range(args.max_new_tokens): + logits = model(input_ids=ids).logits[:, -1, :] + logp = F.log_softmax(logits.float(), dim=-1) + vals, toks = torch.topk(logp[0], k=args.top_k_debug) + print(f"TOPK_STEP {step} PREFIX_IDS {ids[0].tolist()}") + for token_id, value in zip(toks.tolist(), vals.tolist()): + print(token_id, repr(tokenizer.decode([token_id], skip_special_tokens=False)), value) + ids = torch.cat([ids, toks[:1][None, :]], dim=1) + + +if __name__ == "__main__": + main() diff --git a/tools/hf_laoyao_moe/hf_laoyao_moe_configuration.py b/tools/hf_laoyao_moe/hf_laoyao_moe_configuration.py new file mode 100644 index 0000000..1f6b0be --- /dev/null +++ b/tools/hf_laoyao_moe/hf_laoyao_moe_configuration.py @@ -0,0 +1,38 @@ +from transformers import PretrainedConfig + + +class LaoyaoMoeConfig(PretrainedConfig): + model_type = "laoyao_moe" + + def __init__( + self, + vocab_size=154856, + hidden_size=1536, + num_hidden_layers=12, + num_attention_heads=24, + num_key_value_heads=4, + intermediate_size=4608, + moe_intermediate_size=6144, + num_experts=12, + num_experts_per_tok=4, + moe_layers=(2, 4, 6, 8, 10), + max_position_embeddings=8192, + rope_theta=10000.0, + rms_norm_eps=1e-5, + tie_word_embeddings=True, + **kwargs, + ): + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_experts = num_experts + self.num_experts_per_tok = num_experts_per_tok + self.moe_layers = list(moe_layers) + self.max_position_embeddings = max_position_embeddings + self.rope_theta = rope_theta + self.rms_norm_eps = rms_norm_eps diff --git a/tools/hf_laoyao_moe/modeling_laoyao_moe.py b/tools/hf_laoyao_moe/modeling_laoyao_moe.py new file mode 100755 index 0000000..1d92d11 --- /dev/null +++ b/tools/hf_laoyao_moe/modeling_laoyao_moe.py @@ -0,0 +1,202 @@ +import math +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import PreTrainedModel +from transformers.generation import GenerationMixin +from transformers.modeling_outputs import CausalLMOutputWithPast + +from .configuration_laoyao_moe import LaoyaoMoeConfig + + +class LaoyaoRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + return (y * self.weight.float()).to(dtype=x.dtype) + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = torch.chunk(x, 2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def apply_rope(q: torch.Tensor, k: torch.Tensor, rope_theta: float) -> tuple[torch.Tensor, torch.Tensor]: + seq_len = q.shape[1] + head_dim = q.shape[-1] + inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, device=q.device, dtype=torch.float32) / head_dim)) + pos = torch.arange(seq_len, device=q.device, dtype=torch.float32) + freqs = torch.einsum("t,d->td", pos, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos()[None, :, None, :] + sin = emb.sin()[None, :, None, :] + qf = q.float() + kf = k.float() + return (qf * cos + rotate_half(qf) * sin).to(q.dtype), (kf * cos + rotate_half(kf) * sin).to(k.dtype) + + +class LaoyaoAttention(nn.Module): + def __init__(self, config: LaoyaoMoeConfig): + super().__init__() + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.hidden_size // config.num_attention_heads + self.hidden_size = config.hidden_size + self.rope_theta = config.rope_theta + qkv_out = (self.num_heads + 2 * self.num_kv_heads) * self.head_dim + self.qkv_proj = nn.Linear(config.hidden_size, qkv_out, bias=False) + self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + qkv = self.qkv_proj(x) + bsz, seq_len, _ = x.shape + heads_per_group = self.num_heads // self.num_kv_heads + qkv = qkv.view(bsz, seq_len, self.num_kv_heads, heads_per_group + 2, self.head_dim) + q = qkv[:, :, :, :heads_per_group, :].reshape(bsz, seq_len, self.num_heads, self.head_dim) + k = qkv[:, :, :, heads_per_group, :] + v = qkv[:, :, :, heads_per_group + 1, :] + q, k = apply_rope(q, k, self.rope_theta) + repeat = self.num_heads // self.num_kv_heads + k = k.repeat_interleave(repeat, dim=2) + v = v.repeat_interleave(repeat, dim=2) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim) + mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool), diagonal=1) + attn = attn.masked_fill(mask[None, None, :, :], torch.finfo(attn.dtype).min) + out = torch.matmul(torch.softmax(attn, dim=-1), v) + out = out.transpose(1, 2).contiguous().view(bsz, seq_len, self.hidden_size).to(x.dtype) + return self.o_proj(out) + + +class LaoyaoDenseMLP(nn.Module): + def __init__(self, config: LaoyaoMoeConfig): + super().__init__() + self.fc1 = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + a, b = self.fc1(x).chunk(2, dim=-1) + return self.fc2(F.silu(a) * b) + + +class LaoyaoMoeMLP(nn.Module): + def __init__(self, config: LaoyaoMoeConfig): + super().__init__() + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.capacity_factor = config.moe_expert_capacity_factor + self.router = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.experts_fc1 = nn.Parameter(torch.empty(config.num_experts, 2 * config.moe_intermediate_size, config.hidden_size)) + self.experts_fc2 = nn.Parameter(torch.empty(config.num_experts, config.hidden_size, config.moe_intermediate_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + flat_x = x.reshape(-1, x.shape[-1]) + router_logits = F.linear(flat_x, self.router.weight) + topv, topi = torch.topk(router_logits, self.top_k, dim=-1, sorted=torch.is_grad_enabled()) + probs = torch.softmax(topv, dim=-1) + routing_probs = torch.zeros_like(router_logits).scatter(1, topi, probs.type_as(router_logits)) + routing_map = torch.zeros_like(router_logits, dtype=torch.int32).scatter(1, topi, 1).bool() + if self.capacity_factor is not None: + num_tokens = routing_probs.shape[0] + capacity = math.ceil((num_tokens * self.top_k / self.num_experts) * self.capacity_factor) + if capacity <= num_tokens: + _, keep = torch.topk(routing_probs, k=capacity, dim=0, sorted=False) + capacity_mask = torch.zeros_like(routing_probs, dtype=torch.int32).scatter(0, keep, 1).bool() + routing_map = routing_map & capacity_mask + routing_probs = routing_probs * routing_map + flat_out = torch.zeros_like(flat_x) + for expert_id in range(self.num_experts): + mask = routing_map[:, expert_id] + if not mask.any(): + continue + token_idx = mask.nonzero(as_tuple=True)[0] + expert_in = flat_x[token_idx] + hidden = F.linear(expert_in, self.experts_fc1[expert_id]) + a, b = hidden.chunk(2, dim=-1) + hidden = F.silu(a) * b + expert_out = F.linear(hidden, self.experts_fc2[expert_id]) + flat_out[token_idx] += expert_out * routing_probs[token_idx, expert_id].unsqueeze(-1) + return flat_out.reshape_as(x).to(x.dtype) + + +class LaoyaoDecoderLayer(nn.Module): + def __init__(self, config: LaoyaoMoeConfig, layer_idx: int): + super().__init__() + self.input_layernorm = LaoyaoRMSNorm(config.hidden_size, config.rms_norm_eps) + self.self_attn = LaoyaoAttention(config) + self.is_moe = layer_idx in set(config.moe_layers) + self.pre_mlp_layernorm = LaoyaoRMSNorm(config.hidden_size, config.rms_norm_eps) + self.mlp = LaoyaoMoeMLP(config) if self.is_moe else LaoyaoDenseMLP(config) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.self_attn(self.input_layernorm(x)) + x = x + self.mlp(self.pre_mlp_layernorm(x)) + return x + + +class LaoyaoMoePreTrainedModel(PreTrainedModel): + config_class = LaoyaoMoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = False + + +class LaoyaoMoeModel(LaoyaoMoePreTrainedModel): + def __init__(self, config: LaoyaoMoeConfig): + super().__init__(config) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([LaoyaoDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)]) + self.norm = LaoyaoRMSNorm(config.hidden_size, config.rms_norm_eps) + + def forward(self, input_ids: torch.LongTensor) -> torch.Tensor: + x = self.embed_tokens(input_ids) + for layer in self.layers: + x = layer(x) + return self.norm(x) + + +class LaoyaoMoeForCausalLM(LaoyaoMoePreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + all_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config: LaoyaoMoeConfig): + super().__init__(config) + self.config.use_cache = False + self.model = LaoyaoMoeModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.tie_weights() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, value): + self.lm_head = value + + def forward( + self, + input_ids: torch.LongTensor, + labels: Optional[torch.LongTensor] = None, + **kwargs, + ) -> CausalLMOutputWithPast: + hidden = self.model(input_ids) + logits = self.lm_head(hidden).float() + loss = None + if labels is not None: + shift_logits = logits[:, :-1, :].contiguous() + shift_labels = labels[:, 1:].contiguous() + loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + return CausalLMOutputWithPast(loss=loss, logits=logits) diff --git a/tools/inspect_laoyao_ckpt_metadata.py b/tools/inspect_laoyao_ckpt_metadata.py new file mode 100644 index 0000000..f55db0a --- /dev/null +++ b/tools/inspect_laoyao_ckpt_metadata.py @@ -0,0 +1,26 @@ +#!/usr/bin/env python3 +from pathlib import Path + +import torch + + +ckpt = Path( + "runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/" + "checkpoints/iter_0107500" +) + +print("metadata.json:", (ckpt / "metadata.json").read_text()) +common = torch.load(ckpt / "common.pt", map_location="cpu") +print("common iteration:", common.get("iteration")) +print("common keys:", list(common.keys())) +content_metadata = common.get("content_metadata") +print("content_metadata type:", type(content_metadata)) +print("content_metadata:", content_metadata) + +print("run_config relevant lines:") +for line in (ckpt / "run_config.yaml").read_text().splitlines(): + if any( + needle in line.lower() + for needle in ["hf", "hugging", "auto", "architecture", "tokenizer", "model_provider", "bridge"] + ): + print(line[:240]) diff --git a/tools/laoyao_cpu_generate_from_distcp.py b/tools/laoyao_cpu_generate_from_distcp.py new file mode 100644 index 0000000..f2cf5fb --- /dev/null +++ b/tools/laoyao_cpu_generate_from_distcp.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import math +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +import torch.nn.functional as F +from tokenizers import Tokenizer + + +CKPT = Path( + "runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/" + "checkpoints/iter_0107500" +) +TOKENIZER = Path("tokenizer/glm5.2/tokenizer.json") + +HIDDEN = 1536 +HEADS = 24 +KV_HEADS = 4 +HEAD_DIM = 64 +DENSE_FFN = 4608 +MOE_FFN = 6144 +MOE_LAYERS = {2, 4, 6, 8, 10} +TOP_K = 4 +N_LAYERS = 12 +ROPE_BASE = 10000.0 + + +def init_dist() -> None: + if not dist.is_initialized(): + dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29641", rank=0, world_size=1) + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + y = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + return (y * weight.float()).to(torch.bfloat16) + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + out = torch.stack((-x2, x1), dim=-1) + return out.flatten(-2) + + +def apply_rope(q: torch.Tensor, k: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / (ROPE_BASE ** (torch.arange(0, HEAD_DIM, 2, dtype=torch.float32) / HEAD_DIM)) + pos = torch.arange(seq_len, dtype=torch.float32) + freqs = torch.einsum("t,d->td", pos, inv_freq) + emb = torch.repeat_interleave(freqs, 2, dim=-1) + cos = emb.cos()[None, :, None, :] + sin = emb.sin()[None, :, None, :] + qf = q.float() + kf = k.float() + return (qf * cos + rotate_half(qf) * sin).to(torch.bfloat16), ( + kf * cos + rotate_half(kf) * sin + ).to(torch.bfloat16) + + +def swiglu(x: torch.Tensor) -> torch.Tensor: + a, b = x.chunk(2, dim=-1) + return F.silu(a.float()) * b.float() + + +def load_model_tensors(ckpt: Path) -> dict[str, torch.Tensor]: + init_dist() + metadata = dcp.FileSystemReader(str(ckpt)).read_metadata() + state: dict[str, torch.Tensor] = {} + for key, meta in metadata.state_dict_metadata.items(): + if key.startswith("optimizer.") or not hasattr(meta, "size"): + continue + state[key] = torch.empty(tuple(meta.size), dtype=meta.properties.dtype) + print(f"loading {len(state)} model tensors from {ckpt}", flush=True) + dcp.load(state, checkpoint_id=str(ckpt)) + print("loaded model tensors", flush=True) + return state + + +def dense_mlp(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor: + prefix = f"decoder.layers.{layer}.mlp" + h = rms_norm(x, state[f"{prefix}.linear_fc1.layer_norm_weight"]) + fc1 = F.linear(h.float(), state[f"{prefix}.linear_fc1.weight"].float()) + hidden = swiglu(fc1) + return F.linear(hidden, state[f"{prefix}.linear_fc2.weight"].float()).to(torch.bfloat16) + + +def moe_mlp(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor: + prefix = f"decoder.layers.{layer}.mlp" + h = rms_norm(x, state[f"decoder.layers.{layer}.pre_mlp_layernorm.weight"]) + router = F.linear(h.float(), state[f"{prefix}.router.weight"].float()) + topv, topi = torch.topk(router, TOP_K, dim=-1) + probs = torch.softmax(topv, dim=-1) + + fc1_w = state[f"{prefix}.experts.experts.linear_fc1.weight"].float() + fc2_w = state[f"{prefix}.experts.experts.linear_fc2.weight"].float() + out = torch.zeros_like(h.float()) + flat_h = h.float().reshape(-1, HIDDEN) + flat_topi = topi.reshape(-1, TOP_K) + flat_probs = probs.reshape(-1, TOP_K) + flat_out = out.reshape(-1, HIDDEN) + for expert_id in range(fc1_w.shape[0]): + mask = flat_topi == expert_id + if not mask.any(): + continue + token_idx, choice_idx = mask.nonzero(as_tuple=True) + expert_in = flat_h[token_idx] + expert_hidden = swiglu(F.linear(expert_in, fc1_w[expert_id])) + expert_out = F.linear(expert_hidden, fc2_w[expert_id]) + flat_out[token_idx] += expert_out * flat_probs[token_idx, choice_idx].unsqueeze(-1) + return out.to(torch.bfloat16) + + +def layer_forward(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor: + prefix = f"decoder.layers.{layer}" + attn_in = rms_norm(x, state[f"{prefix}.self_attention.linear_qkv.layer_norm_weight"]) + qkv = F.linear(attn_in.float(), state[f"{prefix}.self_attention.linear_qkv.weight"].float()) + q, k, v = torch.split(qkv, [HEADS * HEAD_DIM, KV_HEADS * HEAD_DIM, KV_HEADS * HEAD_DIM], dim=-1) + bsz, seq_len, _ = q.shape + q = q.view(bsz, seq_len, HEADS, HEAD_DIM) + k = k.view(bsz, seq_len, KV_HEADS, HEAD_DIM) + v = v.view(bsz, seq_len, KV_HEADS, HEAD_DIM) + q, k = apply_rope(q, k, seq_len) + repeat = HEADS // KV_HEADS + k = k.repeat_interleave(repeat, dim=2) + v = v.repeat_interleave(repeat, dim=2) + q = q.transpose(1, 2).float() + k = k.transpose(1, 2).float() + v = v.transpose(1, 2).float() + scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(HEAD_DIM) + mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1) + scores = scores.masked_fill(mask[None, None, :, :], float("-inf")) + ctx = torch.matmul(torch.softmax(scores, dim=-1), v) + ctx = ctx.transpose(1, 2).contiguous().view(bsz, seq_len, HIDDEN) + attn_out = F.linear(ctx, state[f"{prefix}.self_attention.linear_proj.weight"].float()).to(torch.bfloat16) + x = x + attn_out + x = x + (moe_mlp(x, state, layer) if layer in MOE_LAYERS else dense_mlp(x, state, layer)) + return x + + +@torch.inference_mode() +def logits_for_ids(ids: list[int], state: dict[str, torch.Tensor]) -> torch.Tensor: + input_ids = torch.tensor(ids, dtype=torch.long).unsqueeze(0) + x = F.embedding(input_ids, state["embedding.word_embeddings.weight"]) + for layer in range(N_LAYERS): + x = layer_forward(x, state, layer) + x = rms_norm(x, state["decoder.final_layernorm.weight"]) + return torch.matmul(x[:, -1, :].float(), state["embedding.word_embeddings.weight"].float().T) + + +def generate(prompt: str, state: dict[str, torch.Tensor], tokenizer: Tokenizer, max_new_tokens: int) -> str: + ids = tokenizer.encode(prompt).ids + print(f"prompt={prompt!r} input_tokens={len(ids)}", flush=True) + for step in range(max_new_tokens): + logits = logits_for_ids(ids, state) + next_id = int(torch.argmax(logits, dim=-1)[0]) + ids.append(next_id) + piece = tokenizer.decode([next_id], skip_special_tokens=False) + print(f"step={step} id={next_id} piece={piece!r}", flush=True) + return tokenizer.decode(ids, skip_special_tokens=False) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--max-new-tokens", type=int, default=8) + parser.add_argument("--prompt", action="append", default=None) + args = parser.parse_args() + + tokenizer = Tokenizer.from_file(str(TOKENIZER)) + prompts = args.prompt or [ + "The capital of France is", + "中国的首都是", + ] + state = load_model_tensors(CKPT) + for prompt in prompts: + text = generate(prompt, state, tokenizer, args.max_new_tokens) + print("=== GENERATED ===") + print(text) + + +if __name__ == "__main__": + main() diff --git a/tools/laoyao_cpu_infer_probe.py b/tools/laoyao_cpu_infer_probe.py new file mode 100644 index 0000000..bdb26bf --- /dev/null +++ b/tools/laoyao_cpu_infer_probe.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import os +from pathlib import Path + +import torch +import torch.distributed as dist + + +def force_megatron_cpu_mode() -> None: + # Megatron-Core allocates several modules with device=torch.cuda.current_device(). + # Returning "cpu" lets those constructors allocate CPU tensors for this probe. + torch.cuda.current_device = lambda: "cpu" # type: ignore[assignment] + torch.cuda.set_device = lambda *_args, **_kwargs: None # type: ignore[assignment] + + +def init_single_process() -> None: + force_megatron_cpu_mode() + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29621") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + if not dist.is_initialized(): + dist.init_process_group("gloo", rank=0, world_size=1) + + from megatron.core import parallel_state + + if not parallel_state.is_initialized(): + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + context_parallel_size=1, + expert_model_parallel_size=1, + ) + + +def build_model(seq_length: int): + import torch.nn.functional as F + from megatron.bridge.models.gpt_provider import GPTModelProvider + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec + from megatron.core.transformer.torch_norm import WrappedTorchNorm + from megatron.core.utils import init_method_normal, scaled_init_method_normal + + def cpu_rmsnorm_block_spec(config, vp_stage=None): + spec = get_gpt_decoder_block_spec( + config, + use_transformer_engine=False, + normalization="RMSNorm", + vp_stage=vp_stage, + ) + spec.layer_norm = WrappedTorchNorm + return spec + + from megatron.core.process_groups_config import ProcessGroupCollection + + provider = GPTModelProvider( + num_layers=12, + hidden_size=1536, + num_attention_heads=24, + num_query_groups=4, + kv_channels=64, + ffn_hidden_size=4608, + seq_length=seq_length, + vocab_size=154856, + should_pad_vocab=True, + share_embeddings_and_output_weights=True, + position_embedding_type="rope", + normalization="RMSNorm", + gated_linear_unit=True, + activation_func=F.silu, + add_bias_linear=False, + num_moe_experts=12, + moe_layer_freq=[1 if idx in {2, 4, 6, 8, 10} else 0 for idx in range(12)], + moe_ffn_hidden_size=6144, + moe_router_topk=4, + moe_router_load_balancing_type="aux_loss", + moe_aux_loss_coeff=0.02, + moe_z_loss_coeff=0.001, + moe_token_dispatcher_type="allgather", + moe_expert_capacity_factor=1.25, + moe_router_enable_expert_bias=False, + moe_router_bias_update_rate=0.02, + moe_grouped_gemm=True, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + context_parallel_size=1, + sequence_parallel=False, + transformer_impl="local", + attention_backend="unfused", + transformer_layer_spec=cpu_rmsnorm_block_spec, + init_method=init_method_normal(0.02), + output_layer_init_method=scaled_init_method_normal(0.02, 12), + init_method_std=0.02, + ) + provider.perform_initialization = False + provider._pg_collection = ProcessGroupCollection.use_mpu_process_groups() + return provider.provide(pre_process=True, post_process=True) + + +def inspect_model(args: argparse.Namespace) -> None: + init_single_process() + model = build_model(args.seq_length) + print("model_built", type(model)) + print("param_count", sum(p.numel() for p in model.parameters())) + if args.mode == "sharded-keys": + state = model.sharded_state_dict(prefix="") + print("sharded_keys_sample") + else: + state = model.state_dict() + print("state_keys_sample") + for key in list(state.keys())[:120]: + print(key) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--seq-length", type=int, default=8192) + parser.add_argument("--mode", choices=["state-keys", "sharded-keys"], default="state-keys") + args = parser.parse_args() + inspect_model(args) + + +if __name__ == "__main__": + main() diff --git a/tools/laoyao_dcp_probe.py b/tools/laoyao_dcp_probe.py new file mode 100644 index 0000000..00a0745 --- /dev/null +++ b/tools/laoyao_dcp_probe.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp + + +CKPT = Path( + "runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/" + "checkpoints/iter_0107500" +) + + +def init_dist() -> None: + if not dist.is_initialized(): + dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29631", rank=0, world_size=1) + + +def main() -> None: + init_dist() + reader = dcp.FileSystemReader(str(CKPT)) + metadata = reader.read_metadata() + keys = list(metadata.state_dict_metadata.keys()) + print("num_keys", len(keys)) + for key in keys: + item = metadata.state_dict_metadata[key] + if hasattr(item, "size"): + print("TENSOR", key, tuple(item.size), item.properties.dtype) + + target = "embedding.word_embeddings.weight" + meta = metadata.state_dict_metadata[target] + shape = tuple(meta.size) + tensor = torch.empty(shape, dtype=torch.bfloat16) + state = {target: tensor} + dcp.load(state, checkpoint_id=str(CKPT)) + print("loaded", target, state[target].shape, state[target].dtype, state[target].float().mean().item()) + + +if __name__ == "__main__": + main() diff --git a/training/megatron_bridge/README.md b/training/megatron_bridge/README.md index 61a1fe9..7255171 100644 --- a/training/megatron_bridge/README.md +++ b/training/megatron_bridge/README.md @@ -8,10 +8,12 @@ - 训练上下文先用 `seq_length=8192`，不要一开始上 16K。 - tokenizer 使用 repo 内已验证的 GLM-5.2 tokenizer：`tokenizer/glm5.2`。 - 从零预训练使用 Megatron indexed dataset，不能直接把 parquet 喂给 Bridge pretrain。 +- 训练数据 blend 优先读取 `prefix_category_stats.json`，按目标 category mix 计算 shard 权重；如果没有 category stats，则退回按 manifest token 数加权。 +- validation 使用单独 heldout 2.8k Megatron indexed prefix，不再从训练 split 里切 valid。 ## 文件 -- `laoyao_2b_moe_pretrain.py`：自定义 Megatron-Bridge recipe/launcher。支持 `--dry-run`，用于在不启动训练 loop 的情况下检查配置。 +- `laoyao_2b_moe_pretrain.py`：自定义 Megatron-Bridge recipe/launcher。包含 GLM-5.2 vocab、MoE grouped GEMM、heldout validation、checkpoint 保留、recompute、distributed optimizer/overlap 参数。 - `../../scripts/preprocess_megatron_bridge_pretrain.sh`：从 parquet 导出 JSONL，并调用 Megatron-LM `preprocess_data.py` 生成 `.bin/.idx`。 - `../../scripts/train_megatron_bridge_2b_moe.sh`：Docker + torchrun 启动入口。 @@ -32,6 +34,17 @@ bash scripts/preprocess_megatron_bridge_pretrain.sh 注意：`preprocess_data.py` 是按文档 tokenization，不在预处理阶段固定切成 8192 行；训练时由 `GPTDatasetConfig.seq_length=8192` 生成固定长度训练 sample。 +当前 g0050 主训练使用直接 parquet -> Megatron indexed dataset 的路径： + +```bash +SOURCE_DATA_DIRS=/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_v1_20260701 \ +WORK_DIR=/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_glm52_direct_v1 \ +MAX_SEQ_LEN=8192 \ +bash scripts/preprocess_megatron_bridge_pretrain_direct.sh +``` + +该转换脚本会在 `manifest.json` 中写入可直接传给 Megatron 的 `_text_document` prefix；不要手动拼接错误的 prefix。 + ## Dry Run H200 被占用时可以先跑单进程 dry-run： diff --git a/training/megatron_bridge/laoyao_2b_moe_pretrain.py b/training/megatron_bridge/laoyao_2b_moe_pretrain.py index 2cfc374..eb0edf2 100755 --- a/training/megatron_bridge/laoyao_2b_moe_pretrain.py +++ b/training/megatron_bridge/laoyao_2b_moe_pretrain.py @@ -3,12 +3,14 @@ from __future__ import annotations import argparse import json +from functools import partial from pathlib import Path import torch import torch.nn.functional as F from megatron.bridge.models.gpt_provider import GPTModelProvider +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec from megatron.bridge.recipes.common import _pretrain_common from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.gpt_step import forward_step @@ -17,19 +19,84 @@ from megatron.bridge.training.pretrain import pretrain REPO_ROOT = Path(__file__).resolve().parents[2] TOKENIZER_DIR = REPO_ROOT / "tokenizer/glm5.2" +DEFAULT_VALIDATION_PREFIX = REPO_ROOT / "dataset/val/megatron_8192_glm52/heldout_2p8k_text_document" -def load_data_blend(args: argparse.Namespace) -> list[tuple[str, float]] | None: +TARGET_CATEGORY_MIX = { + "english_web": 0.40, + "english_edu": 0.20, + "chinese_clean": 0.10, + "science": 0.10, + "logic": 0.10, + "math": 0.05, + "code": 0.05, +} + + +def _normalize_weights(raw_weights: list[float]) -> list[float]: + total = sum(raw_weights) + if total <= 0: + raise ValueError("Data blend weights sum to zero.") + return [weight / total for weight in raw_weights] + + +def _token_weights_from_manifest(prefix_list: list[str], manifest: dict) -> list[float]: + token_by_prefix = {} + for result in manifest.get("results") or []: + output_prefix = result.get("output_prefix") + tokens = result.get("tokens_estimated_from_int32_bin") + if output_prefix and tokens: + token_by_prefix[str(output_prefix) + "_text_document"] = float(tokens) + + weights = [token_by_prefix.get(prefix, 0.0) for prefix in prefix_list] + if any(weight <= 0 for weight in weights): + missing = [prefix for prefix, weight in zip(prefix_list, weights) if weight <= 0] + raise ValueError(f"Missing token counts for {len(missing)} data prefixes; first={missing[:3]}") + return _normalize_weights(weights) + + +def _target_mix_weights_from_category_stats( + prefix_list: list[str], + category_stats_path: Path, +) -> list[float] | None: + if not category_stats_path.exists(): + return None + + stats = json.loads(category_stats_path.read_text(encoding="utf-8")) + per_prefix = stats.get("prefix_category_tokens") or {} + category_totals = stats.get("category_tokens") or {} + raw_weights = [] + for prefix in prefix_list: + category_tokens = per_prefix.get(prefix) + if not category_tokens: + return None + + shard_weight = 0.0 + for category, tokens in category_tokens.items(): + target_ratio = TARGET_CATEGORY_MIX.get(category, 0.0) + category_total = float(category_totals.get(category, 0.0)) + if target_ratio > 0 and category_total > 0 and tokens: + shard_weight += target_ratio * (float(tokens) / category_total) + raw_weights.append(shard_weight) + + return _normalize_weights(raw_weights) + + +def load_data_blend(args: argparse.Namespace) -> tuple[list[str], list[float] | None] | None: if args.data_manifest: manifest_path = Path(args.data_manifest) manifest = json.loads(manifest_path.read_text(encoding="utf-8")) prefixes = manifest.get("ok_prefixes") or [] if not prefixes: raise ValueError(f"{manifest_path} has no ok_prefixes") - weight = 1.0 / len(prefixes) - return [(str(prefix), weight) for prefix in 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) + if weights is None: + weights = _token_weights_from_manifest(prefix_list, manifest) + return (prefix_list, weights) if args.data_prefix: - return [(args.data_prefix, 1.0)] + return ([args.data_prefix], None) return None @@ -43,13 +110,14 @@ def build_config(args: argparse.Namespace) -> ConfigContainer: num_query_groups=4, ffn_hidden_size=4608, seq_length=args.seq_length, - vocab_size=154820, + vocab_size=154856, should_pad_vocab=True, share_embeddings_and_output_weights=True, position_embedding_type="rope", normalization="RMSNorm", gated_linear_unit=True, activation_func=F.silu, + add_bias_linear=False, num_moe_experts=12, moe_layer_freq=[1 if idx in {2, 4, 6, 8, 10} else 0 for idx in range(12)], moe_ffn_hidden_size=6144, @@ -59,9 +127,9 @@ def build_config(args: argparse.Namespace) -> ConfigContainer: moe_z_loss_coeff=0.001, moe_token_dispatcher_type="alltoall", moe_expert_capacity_factor=1.25, - moe_router_enable_expert_bias=True, + moe_router_enable_expert_bias=False, moe_router_bias_update_rate=0.02, - moe_grouped_gemm=True, + moe_grouped_gemm=not args.disable_moe_grouped_gemm, tensor_model_parallel_size=args.tensor_parallel, pipeline_model_parallel_size=args.pipeline_parallel, expert_model_parallel_size=args.expert_parallel, @@ -69,18 +137,34 @@ def build_config(args: argparse.Namespace) -> ConfigContainer: sequence_parallel=args.tensor_parallel > 1, transformer_impl="transformer_engine", attention_backend="flash", + transformer_layer_spec=partial(get_gpt_decoder_block_spec, use_transformer_engine=True), init_method_std=0.02, + recompute_granularity=args.recompute_granularity, + recompute_method=args.recompute_method, + recompute_num_layers=args.recompute_num_layers, + recompute_modules=args.recompute_modules, ) cfg.tokenizer.tokenizer_type = "HuggingFaceTokenizer" cfg.tokenizer.tokenizer_model = str(TOKENIZER_DIR) cfg.dataset.seq_length = args.seq_length cfg.dataset.split = args.split - cfg.dataset.num_workers = args.dataset_workers + cfg.dataset.mid_level_dataset_surplus = 0.05 data_blend = load_data_blend(args) if data_blend: - cfg.dataset.data_path = data_blend[0][0] - cfg.dataset.blend = data_blend + validation_prefix = Path(args.validation_prefix) if args.validation_prefix else None + if validation_prefix and validation_prefix.with_suffix(".bin").exists() and validation_prefix.with_suffix(".idx").exists(): + cfg.dataset.data_path = None + cfg.dataset.split = None + cfg.dataset.blend = None + cfg.dataset.blend_per_split = [ + data_blend, + ([str(validation_prefix)], None), + None, + ] + else: + cfg.dataset.data_path = data_blend[0][0] + cfg.dataset.blend = data_blend else: cfg.dataset.blend = None @@ -91,21 +175,31 @@ def build_config(args: argparse.Namespace) -> ConfigContainer: cfg.optimizer.min_lr = args.min_lr cfg.optimizer.weight_decay = args.weight_decay cfg.scheduler.lr_warmup_fraction = args.warmup_fraction + cfg.scheduler.lr_warmup_iters = 0 + cfg.scheduler.lr_warmup_samples = 0 cfg.checkpoint.save = args.save_dir cfg.checkpoint.load = args.load_dir cfg.checkpoint.save_interval = args.save_interval + cfg.checkpoint.most_recent_k = args.keep_last_checkpoints cfg.logger.tensorboard_dir = args.tensorboard_dir cfg.logger.log_interval = args.log_interval + cfg.logger.log_progress = True + cfg.logger.log_runtime_to_tensorboard = True + cfg.logger.log_throughput = True + cfg.logger.log_throughput_to_tensorboard = True + cfg.logger.log_validation_ppl_to_tensorboard = True cfg.validation.eval_interval = args.eval_interval cfg.validation.eval_iters = args.eval_iters cfg.ddp.use_megatron_fsdp = False - cfg.ddp.overlap_grad_reduce = True - cfg.ddp.overlap_param_gather = True + cfg.ddp.overlap_grad_reduce = args.overlap_grad_reduce + cfg.ddp.overlap_param_gather = args.overlap_param_gather cfg.ddp.check_for_nan_in_grad = True - cfg.ddp.use_distributed_optimizer = True + cfg.ddp.use_distributed_optimizer = args.use_distributed_optimizer + cfg.optimizer.use_distributed_optimizer = args.use_distributed_optimizer + cfg.optimizer.overlap_param_gather = args.overlap_param_gather return cfg @@ -114,6 +208,11 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Laoyao 2B MoE Megatron-Bridge pretrain launcher.") parser.add_argument("--data-prefix", default=None, help="Megatron indexed dataset prefix without .bin/.idx.") parser.add_argument("--data-manifest", default=None, help="Manifest produced by convert_pretrain_parquet_to_megatron.py.") + parser.add_argument( + "--validation-prefix", + default=str(DEFAULT_VALIDATION_PREFIX), + help="Optional Megatron indexed dataset prefix used as the validation split.", + ) parser.add_argument("--seq-length", type=int, default=8192) parser.add_argument("--train-iters", type=int, default=10) parser.add_argument("--micro-batch-size", type=int, default=1) @@ -122,6 +221,14 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--pipeline-parallel", type=int, default=1) parser.add_argument("--expert-parallel", type=int, default=1) parser.add_argument("--context-parallel", type=int, default=1) + parser.add_argument("--disable-moe-grouped-gemm", action="store_true", help="Disable grouped GEMM for MoE experts.") + parser.add_argument("--recompute-granularity", choices=["full", "selective"], default=None) + parser.add_argument("--recompute-method", choices=["uniform", "block"], default=None) + parser.add_argument("--recompute-num-layers", type=int, default=None) + parser.add_argument("--recompute-modules", nargs="*", default=None) + parser.add_argument("--use-distributed-optimizer", action="store_true", help="Enable Megatron distributed optimizer.") + parser.add_argument("--overlap-grad-reduce", action="store_true", help="Overlap DDP gradient reduction.") + parser.add_argument("--overlap-param-gather", action="store_true", help="Overlap distributed optimizer parameter gather.") parser.add_argument("--dataset-workers", type=int, default=8) parser.add_argument("--split", default="9999,8,2") parser.add_argument("--lr", type=float, default=3.0e-4) @@ -132,6 +239,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--load-dir", default=None) parser.add_argument("--tensorboard-dir", default="/mnt/beegfs/yi/laoyao_2b_moe/runs/megatron_bridge/tensorboard") parser.add_argument("--save-interval", type=int, default=1000) + parser.add_argument("--keep-last-checkpoints", type=int, default=10) parser.add_argument("--log-interval", type=int, default=10) parser.add_argument("--eval-interval", type=int, default=1000) parser.add_argument("--eval-iters", type=int, default=10) @@ -151,10 +259,19 @@ def main() -> None: print(f"moe_layer_freq={cfg.model.moe_layer_freq}") print(f"moe_router_load_balancing_type={cfg.model.moe_router_load_balancing_type}") print(f"moe_aux_loss_coeff={cfg.model.moe_aux_loss_coeff}") + print(f"moe_grouped_gemm={cfg.model.moe_grouped_gemm}") + print(f"recompute_granularity={cfg.model.recompute_granularity}") + print(f"recompute_method={cfg.model.recompute_method}") + print(f"recompute_num_layers={cfg.model.recompute_num_layers}") + print(f"recompute_modules={cfg.model.recompute_modules}") + print(f"use_distributed_optimizer={cfg.optimizer.use_distributed_optimizer}") + print(f"overlap_grad_reduce={cfg.ddp.overlap_grad_reduce}") + print(f"overlap_param_gather={cfg.ddp.overlap_param_gather}") print(f"tokenizer_model={cfg.tokenizer.tokenizer_model}") print(f"dataset_prefix={args.data_prefix or 'mock'}") print(f"dataset_manifest={args.data_manifest or 'none'}") print(f"dataset_blend_size={len(cfg.dataset.blend or [])}") + print(f"dataset_blend_per_split={cfg.dataset.blend_per_split is not None}") return pretrain(config=cfg, forward_step_func=forward_step)