yi_lu/laoyao_2b_moe
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Sync g0050 training and inference tooling

Browse Source
This commit is contained in:
yi_lu
2026-07-09 08:02:25 +08:00
parent 0e8c2a554d
commit 2b0588ad17
29 changed files with 1726 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
.gitignore vendored
View File

@@ -17,6 +17,10 @@ core.*
dataset/pretrain/data/* dataset/pretrain/data/*
!dataset/pretrain/data/.gitkeep !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 # Python/cache/editor
__pycache__/ __pycache__/
*.py[cod] *.py[cod]

34
README.md
View File

@@ -7,7 +7,9 @@
- `dataset/`: 预训练数据和验证集的构建、manifest、数据落盘位置。 - `dataset/`: 预训练数据和验证集的构建、manifest、数据落盘位置。
- `model/`: 2B MoE 架构定义，优先用 NeMo/Megatron 配置表达。 - `model/`: 2B MoE 架构定义，优先用 NeMo/Megatron 配置表达。
- `training/`: 训练 recipe、评估 recipe、并行和优化超参。 - `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 /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 路径约定 ## g0033 路径约定

8
dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py
View File

@@ -148,6 +148,8 @@ def convert_one_parquet(task: tuple[str, argparse.Namespace]) -> dict[str, objec
"status": "skip", "status": "skip",
"source": str(parquet_path), "source": str(parquet_path),
"output_prefix": str(output_prefix), "output_prefix": str(output_prefix),
"bin_file": str(bin_file),
"idx_file": str(idx_file),
"reason": "existing bin/idx", "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_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_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), "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") manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
print(f"manifest={manifest_path}", flush=True) print(f"manifest={manifest_path}", flush=True)

10
dataset/val/README.md
View File

@@ -19,3 +19,13 @@
- `build_stats.json`: 构建统计和来源分布。 - `build_stats.json`: 构建统计和来源分布。
这些样本来源包括 GSM8K、COIG-CQIA、UltraChat、HumanEval、MBPP、BBH logical deduction、ARC-Challenge、MMLU formal logic/logical fallacies、CEVAL 等。 这些样本来源包括 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。

111
docs/megatron_to_hf_export_notes.md Normal file
View File

@@ -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 仍是更可信的推理基线。

27
scripts/README.md
View File

@@ -7,6 +7,13 @@
- `train_megatron_bridge_2b_moe.sh`: 当前主训练入口，使用 NeMo 26.06 镜像中的 Megatron-Bridge。 - `train_megatron_bridge_2b_moe.sh`: 当前主训练入口，使用 NeMo 26.06 镜像中的 Megatron-Bridge。
- `train_nemo_megatron_2b_moe.sh`: NeMo/Megatron 训练入口占位，包含 image、mount、路径检查。 - `train_nemo_megatron_2b_moe.sh`: NeMo/Megatron 训练入口占位，包含 image、mount、路径检查。
- `g0050_download_and_setup_from_modelscope.sh`: 在 g0050 上一键准备训练环境并从 ModelScope 下载未 tokenize parquet 数据。 - `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 下载与部署 ## g0050 下载与部署
@@ -32,3 +39,23 @@ GIT_REPO_URL=https://yi_lu:<token>@git.deeepseek.net/yi_lu/laoyao_2b_moe.git \
MODELSCOPE_API_TOKEN=ms-... \ MODELSCOPE_API_TOKEN=ms-... \
bash scripts/g0050_download_and_setup_from_modelscope.sh 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`

29
scripts/check_laoyao_megatron_inference_ready.sh Executable file
View File

@@ -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

66
scripts/g0050_wait_and_tokenize_glm52_8192.sh Executable file
View File

@@ -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

72
scripts/laoyao_megatron_inference_setup.sh Executable file
View File

@@ -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."

18
scripts/query_laoyao_megatron_server.sh Executable file
View File

@@ -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
}'

96
scripts/resume_pretrain_8192_8gpu_mbs14.sh Executable file
View File

@@ -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"

43
scripts/serve_laoyao_megatron.sh Executable file
View File

@@ -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

65
scripts/serve_laoyao_megatron_inner.sh Executable file
View File

@@ -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}"

6
scripts/stop_laoyao_megatron_server.sh Executable file
View File

@@ -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}"

2
tokenizer/README.md
View File

@@ -17,7 +17,7 @@ glm5.2/generation_config.json
```text ```text
tokenizer.json sha256: 19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d tokenizer.json sha256: 19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d
vocab_size: 154820 vocab_size used by training config: 154856
config.model_type: glm_moe_dsa config.model_type: glm_moe_dsa
source: https://modelscope.ai/models/zai-org/GLM-5.2/files source: https://modelscope.ai/models/zai-org/GLM-5.2/files
``` ```

35
tools/README.md Normal file
View File

@@ -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`。

33
tools/dump_laoyao_model_keys.py Normal file
View File

@@ -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()

69
tools/hf_laoyao_moe/README.md Normal file
View File

@@ -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`

41
tools/hf_laoyao_moe/configuration_laoyao_moe.py Executable file
View File

@@ -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

115
tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py Executable file
View File

@@ -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()

80
tools/hf_laoyao_moe/generate_laoyao_hf.py Executable file
View File

@@ -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()

38
tools/hf_laoyao_moe/hf_laoyao_moe_configuration.py Normal file
View File

@@ -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

202
tools/hf_laoyao_moe/modeling_laoyao_moe.py Executable file
View File

@@ -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)

26
tools/inspect_laoyao_ckpt_metadata.py Normal file
View File

@@ -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])

185
tools/laoyao_cpu_generate_from_distcp.py Normal file
View File

@@ -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()

128
tools/laoyao_cpu_infer_probe.py Normal file
View File

@@ -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()

43
tools/laoyao_dcp_probe.py Normal file
View File

@@ -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()

15
training/megatron_bridge/README.md
View File

@@ -8,10 +8,12 @@
- 训练上下文先用 `seq_length=8192`，不要一开始上 16K。 - 训练上下文先用 `seq_length=8192`，不要一开始上 16K。
- tokenizer 使用 repo 内已验证的 GLM-5.2 tokenizer：`tokenizer/glm5.2`。 - tokenizer 使用 repo 内已验证的 GLM-5.2 tokenizer：`tokenizer/glm5.2`。
- 从零预训练使用 Megatron indexed dataset，不能直接把 parquet 喂给 Bridge pretrain。 - 从零预训练使用 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/preprocess_megatron_bridge_pretrain.sh`：从 parquet 导出 JSONL，并调用 Megatron-LM `preprocess_data.py` 生成 `.bin/.idx`。
- `../../scripts/train_megatron_bridge_2b_moe.sh`：Docker + torchrun 启动入口。 - `../../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。 注意：`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 ## Dry Run
H200 被占用时可以先跑单进程 dry-run： H200 被占用时可以先跑单进程 dry-run：

143
training/megatron_bridge/laoyao_2b_moe_pretrain.py
View File

@@ -3,12 +3,14 @@ from __future__ import annotations
import argparse import argparse
import json import json
from functools import partial
from pathlib import Path from pathlib import Path
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
from megatron.bridge.models.gpt_provider import GPTModelProvider 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.recipes.common import _pretrain_common
from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.config import ConfigContainer
from megatron.bridge.training.gpt_step import forward_step 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] REPO_ROOT = Path(__file__).resolve().parents[2]
TOKENIZER_DIR = REPO_ROOT / "tokenizer/glm5.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: if args.data_manifest:
manifest_path = Path(args.data_manifest) manifest_path = Path(args.data_manifest)
manifest = json.loads(manifest_path.read_text(encoding="utf-8")) manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
prefixes = manifest.get("ok_prefixes") or [] prefixes = manifest.get("ok_prefixes") or []
if not prefixes: if not prefixes:
raise ValueError(f"{manifest_path} has no ok_prefixes") raise ValueError(f"{manifest_path} has no ok_prefixes")
weight = 1.0 / len(prefixes) prefix_list = [str(prefix) for prefix in prefixes]
return [(str(prefix), weight) 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: if args.data_prefix:
return [(args.data_prefix, 1.0)] return ([args.data_prefix], None)
return None return None
@@ -43,13 +110,14 @@ def build_config(args: argparse.Namespace) -> ConfigContainer:
num_query_groups=4, num_query_groups=4,
ffn_hidden_size=4608, ffn_hidden_size=4608,
seq_length=args.seq_length, seq_length=args.seq_length,
vocab_size=154820, vocab_size=154856,
should_pad_vocab=True, should_pad_vocab=True,
share_embeddings_and_output_weights=True, share_embeddings_and_output_weights=True,
position_embedding_type="rope", position_embedding_type="rope",
normalization="RMSNorm", normalization="RMSNorm",
gated_linear_unit=True, gated_linear_unit=True,
activation_func=F.silu, activation_func=F.silu,
add_bias_linear=False,
num_moe_experts=12, num_moe_experts=12,
moe_layer_freq=[1 if idx in {2, 4, 6, 8, 10} else 0 for idx in range(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_ffn_hidden_size=6144,
@@ -59,9 +127,9 @@ def build_config(args: argparse.Namespace) -> ConfigContainer:
moe_z_loss_coeff=0.001, moe_z_loss_coeff=0.001,
moe_token_dispatcher_type="alltoall", moe_token_dispatcher_type="alltoall",
moe_expert_capacity_factor=1.25, 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_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, tensor_model_parallel_size=args.tensor_parallel,
pipeline_model_parallel_size=args.pipeline_parallel, pipeline_model_parallel_size=args.pipeline_parallel,
expert_model_parallel_size=args.expert_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, sequence_parallel=args.tensor_parallel > 1,
transformer_impl="transformer_engine", transformer_impl="transformer_engine",
attention_backend="flash", attention_backend="flash",
transformer_layer_spec=partial(get_gpt_decoder_block_spec, use_transformer_engine=True),
init_method_std=0.02, 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_type = "HuggingFaceTokenizer"
cfg.tokenizer.tokenizer_model = str(TOKENIZER_DIR) cfg.tokenizer.tokenizer_model = str(TOKENIZER_DIR)
cfg.dataset.seq_length = args.seq_length cfg.dataset.seq_length = args.seq_length
cfg.dataset.split = args.split 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) data_blend = load_data_blend(args)
if data_blend: if data_blend:
cfg.dataset.data_path = data_blend[0][0] validation_prefix = Path(args.validation_prefix) if args.validation_prefix else None
cfg.dataset.blend = data_blend 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: else:
cfg.dataset.blend = None cfg.dataset.blend = None
@@ -91,21 +175,31 @@ def build_config(args: argparse.Namespace) -> ConfigContainer:
cfg.optimizer.min_lr = args.min_lr cfg.optimizer.min_lr = args.min_lr
cfg.optimizer.weight_decay = args.weight_decay cfg.optimizer.weight_decay = args.weight_decay
cfg.scheduler.lr_warmup_fraction = args.warmup_fraction 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.save = args.save_dir
cfg.checkpoint.load = args.load_dir cfg.checkpoint.load = args.load_dir
cfg.checkpoint.save_interval = args.save_interval 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.tensorboard_dir = args.tensorboard_dir
cfg.logger.log_interval = args.log_interval 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_interval = args.eval_interval
cfg.validation.eval_iters = args.eval_iters cfg.validation.eval_iters = args.eval_iters
cfg.ddp.use_megatron_fsdp = False cfg.ddp.use_megatron_fsdp = False
cfg.ddp.overlap_grad_reduce = True cfg.ddp.overlap_grad_reduce = args.overlap_grad_reduce
cfg.ddp.overlap_param_gather = True cfg.ddp.overlap_param_gather = args.overlap_param_gather
cfg.ddp.check_for_nan_in_grad = True 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 return cfg
@@ -114,6 +208,11 @@ def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Laoyao 2B MoE Megatron-Bridge pretrain launcher.") 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-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("--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("--seq-length", type=int, default=8192)
parser.add_argument("--train-iters", type=int, default=10) parser.add_argument("--train-iters", type=int, default=10)
parser.add_argument("--micro-batch-size", type=int, default=1) 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("--pipeline-parallel", type=int, default=1)
parser.add_argument("--expert-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("--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("--dataset-workers", type=int, default=8)
parser.add_argument("--split", default="9999,8,2") parser.add_argument("--split", default="9999,8,2")
parser.add_argument("--lr", type=float, default=3.0e-4) 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("--load-dir", default=None)
parser.add_argument("--tensorboard-dir", default="/mnt/beegfs/yi/laoyao_2b_moe/runs/megatron_bridge/tensorboard") 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("--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("--log-interval", type=int, default=10)
parser.add_argument("--eval-interval", type=int, default=1000) parser.add_argument("--eval-interval", type=int, default=1000)
parser.add_argument("--eval-iters", type=int, default=10) 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_layer_freq={cfg.model.moe_layer_freq}")
print(f"moe_router_load_balancing_type={cfg.model.moe_router_load_balancing_type}") 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_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"tokenizer_model={cfg.tokenizer.tokenizer_model}")
print(f"dataset_prefix={args.data_prefix or 'mock'}") print(f"dataset_prefix={args.data_prefix or 'mock'}")
print(f"dataset_manifest={args.data_manifest or 'none'}") print(f"dataset_manifest={args.data_manifest or 'none'}")
print(f"dataset_blend_size={len(cfg.dataset.blend or [])}") 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 return
pretrain(config=cfg, forward_step_func=forward_step) pretrain(config=cfg, forward_step_func=forward_step)