laoyao_2b_moe/AGENTS.md

AGENTS.md

This repository is the reproducible Laoyao 2B MoE pretraining project. It refactors the original hand-written Laoyao training experiment into a Megatron/NVIDIA NeMo style workflow with explicit data construction, tokenizer choice, model definition, training launchers, checkpoint handling, and inference/export tooling.

Use this file as the first entry point when operating the repo as an agent.

Repository Purpose

The project trains a small MoE language model for practice and infrastructure validation:

• Model family: Laoyao 2B MoE.
• Training backend: Megatron-Bridge / Megatron-Core inside a NeMo 26.06 based Docker image.
• Tokenizer: GLM-5.2 tokenizer from zai-org/GLM-5.2.
• Main training length: seq_length=8192.
• Main data target: approximately 200B tokens with rebalanced category mix.
• Main validation: heldout 2.8k public task/dialogue set, converted to Megatron indexed validation data on the training host.

The repo is designed to keep source code, manifests, configs, and docs in git while keeping generated corpora, checkpoints, tensorboard logs, and model exports out of git.

Directory Structure

• dataset/
  • pretrain/: pretraining data construction scripts, configs, manifests, and README.
  • val/: heldout 2.8k validation data in source JSONL forms and docs.
• model/
  • Megatron/NeMo model architecture config and architecture documentation.
• training/
  • Megatron-Bridge pretraining recipe and training docs.
• docker/
  • NeMo/Megatron Docker image definition and backend notes.
• scripts/
  • Shell entrypoints for data setup, tokenization, Docker image build, training, resume, inference server, and query smoke tests.
• tools/
  • Checkpoint inspection, Megatron DCP probes, HF export, and HF custom model generation tools.
• docs/
  • Case studies and operational notes, especially Megatron-to-HF export pitfalls.
• runs/
  • Runtime outputs, checkpoints, logs, tensorboard files, and exports. Ignored by git.

Important Hosts and Paths

Current g0050 training repo:

/ssd/workspace/yi/laoyao_2b_moe

Current g0033 mirror/dev repo:

/mnt/beegfs/yi/laoyao_2b_moe

Current g0050 pretokenized data root:

/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset

Current main run:

/ssd/workspace/yi/laoyao_2b_moe/runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000

Data Construction

The pretraining target mix is:

Category	Target Ratio	Notes
english_web	40%	English web text, FineWeb/Ultra-FineWeb style sources.
english_edu	20%	Educational English text; also fills science shortfall when needed.
chinese_clean	10%	Clean Chinese web/text data.
science	10%	Science QA/reasoning sources; may be much smaller in available raw data.
logic	10%	Proof-Pile-2/OpenWebMath/arXiv/AlgebraicStack style high-reasoning text.
math	5%	Math data from cleaned/scored sources.
code	5%	Cleaned multilingual code data.

Core source scripts:

• dataset/pretrain/scripts/download_rebalanced_sources.py
• dataset/pretrain/scripts/build_rebalanced_pretrain_dataset.py
• dataset/pretrain/scripts/wait_and_build_rebalanced_pretrain_200b.sh
• dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py

Current g0050 tokenization entrypoint:

bash scripts/g0050_wait_and_tokenize_glm52_8192.sh

This waits for ModelScope parquet downloads, checks a minimum parquet count, then calls:

bash scripts/preprocess_megatron_bridge_pretrain_direct.sh

Important conversion behavior:

• Converts parquet directly to Megatron indexed dataset .bin/.idx without intermediate JSONL.
• Uses GLM-5.2 tokenizer from tokenizer/glm5.2.
• Uses MAX_SEQ_LEN=8192 by default.
• Writes manifest.json with ok_prefixes that already include the correct _text_document prefix.
• Writes or consumes prefix_category_stats.json when available so training can weight shards by category target mix.

Generated Megatron indexed datasets are not committed to git.

Validation Data

Source validation data lives in:

dataset/val/data/

It contains a 2,800-example heldout set with seven categories:

• science_reasoning
• logic
• code
• chinese_exam
• math
• chinese_dialogue
• english_dialogue

The training job expects a generated Megatron indexed validation prefix on g0050:

dataset/val/megatron_8192_glm52/heldout_2p8k_text_document

This generated directory is ignored by git. If missing, rebuild it from dataset/val/data/heldout_2p8k_packed_text.jsonl using the same GLM-5.2 tokenizer and Megatron indexed dataset writer.

Model Definition

The active model implementation is in:

training/megatron_bridge/laoyao_2b_moe_pretrain.py

Key architecture settings:

• num_layers=12
• hidden_size=1536
• num_attention_heads=24
• num_query_groups=4
• GQA attention.
• ffn_hidden_size=4608 for dense MLP layers.
• MoE layers at layer indices {2, 4, 6, 8, 10}.
• num_moe_experts=12
• moe_router_topk=4
• moe_ffn_hidden_size=6144
• 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=True unless disabled by CLI.
• share_embeddings_and_output_weights=True.
• vocab_size=154856.
• normalization="RMSNorm", rms_norm_eps is Megatron default/runtime aligned.
• RoPE is Megatron non-interleaved / NeoX half-split style.
• Attention backend is configured as flash through Transformer Engine.

The static YAML config is documented under:

model/nemo_megatron/

However, the currently exercised training path is the Python Megatron-Bridge recipe.

Training

Current primary resume/training entrypoint on g0050:

cd /ssd/workspace/yi/laoyao_2b_moe
bash scripts/resume_pretrain_8192_8gpu_mbs14.sh

Main long-run defaults:

• Image: laoyao/nemo-megatron:26.06-flashattn4
• Container: laoyao_pretrain_resume
• 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
• tensor_parallel=1
• pipeline_parallel=1
• expert_parallel=1
• context_parallel=1
• dataset_workers=4
• save_interval=2500
• keep_last_checkpoints=10
• eval_interval=15000
• eval_iters=10
• distributed optimizer enabled.
• overlap_grad_reduce enabled.
• overlap_param_gather enabled.
• full activation recompute enabled with recompute_method=uniform and recompute_num_layers=1.

The training recipe logs throughput/MFU-like TFLOP metrics, loss, load balancing loss, z-loss, grad norm, and validation perplexity/loss when validation runs.

Docker Environment

Docker files and docs live under:

docker/nemo/

Build entrypoint:

bash scripts/build_nemo_megatron_image.sh

The current image is:

laoyao/nemo-megatron:26.06-flashattn4

The image is based on NVIDIA NeMo 26.06 / Megatron-Bridge 0.5.0 style stack and includes flash-attn4/cutlass-related fixes for B300/H200 compatibility. Avoid changing CUDA/NVIDIA stack casually during an active training run.

Resume, Kill, and Status

Status check:

docker ps --format '{{.Names}} {{.Status}} {{.Image}}' | grep -E 'laoyao_pretrain|tensorboard' || true
nvidia-smi --query-gpu=index,memory.used,utilization.gpu,power.draw --format=csv,noheader,nounits
awk '/^[[:space:]]*\[[0-9][0-9][0-9][0-9]-.* iteration[[:space:]]+[0-9]+\// {print}' \
  /tmp/laoyao_pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000_resume.log | tail -12

Resume training:

bash scripts/resume_pretrain_8192_8gpu_mbs14.sh

Stop training container:

docker rm -f laoyao_pretrain_resume

Before killing, prefer checking the latest checkpoint:

cat runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints/latest_checkpointed_iteration.txt
ls -1dt runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints/iter_* | head

The separate infra handoff scripts may also exist at:

/ssd/workspace/yi/infra_handoff/laoyao_infra_scripts

If those exist, read their AGENTS.md before interrupting a long run for infra users.

Evaluation

During training, validation runs every 15000 iterations by default using the heldout 2.8k Megatron indexed validation prefix.

Expected metrics:

• validation LM loss
• validation PPL
• training loss
• load balancing loss
• z-loss
• grad norm
• throughput / TFLOP/s/GPU

TensorBoard output is under the active run directory:

runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/tensorboard

Do not commit TensorBoard logs.

HF Heldout 2.8k Evaluation

After exporting a checkpoint to HF format, run:

python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \
  --model-dir runs/hf_exports/iter_0107500 \
  --data dataset/val/data/heldout_2p8k_sft_prompt_completion.jsonl \
  --out-dir runs/hf_eval/heldout_2p8k/iter_0107500 \
  --device cuda \
  --max-length 2048

For CPU smoke tests, use a very small sample count:

python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \
  --model-dir runs/hf_exports/iter_0107500 \
  --max-items 4 \
  --device cpu \
  --dtype float32 \
  --max-length 512

This evaluation follows the tokenizer-swap heldout logic:

• Compute prompt-conditioned completion NLL/PPL for all 2.8k examples.
• Parse MCQ choices only where the prompt has explicit A./B./... choices and the gold completion starts with the answer label.
• Score each MCQ choice as a continuation.
• Main accuracy is mcq_acc_avg_norm, based on average token logprob.
• mcq_acc_sum is auxiliary and can favor shorter choices.

Do not treat all 2.8k examples as MCQ. The set also includes GSM8K, HumanEval/MBPP style code, Chinese dialogue, and English dialogue examples.

Model Weight Exporting

The HF export tools are under:

tools/hf_laoyao_moe/

Export command template:

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

The export writes:

• model.safetensors
• config.json
• tokenizer files
• HF custom model files

Important export/inference alignment details:

• HF custom model must use use_cache=false; KV cache is not implemented.
• RMSNorm epsilon must match Megatron runtime, currently 1e-5.
• RoPE must match Megatron non-interleaved / NeoX half-split layout.
• MoE router should account for moe_expert_capacity_factor=1.25 token dropping behavior.
• HF dynamic module imports must be relative imports.

Read the case study before changing export logic:

docs/megatron_to_hf_export_notes.md

Inference

Megatron-native inference server scripts:

bash scripts/check_laoyao_megatron_inference_ready.sh
bash scripts/serve_laoyao_megatron.sh
bash scripts/query_laoyao_megatron_server.sh
bash scripts/stop_laoyao_megatron_server.sh

The Megatron server is the more trustworthy inference baseline because it uses the native Megatron model stack. HF export is useful for portability and debugging but is not the highest-performance inference path.

HF custom model generation after export:

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"

The HF generation script disables KV cache by default. Do not remove that unless KV cache support has been implemented and tested.

Git and Artifact Policy

Commit and push:

git status --short
git add <source/docs/scripts only>
git commit -m "..."
git push origin main

Do not commit:

• runs/
• checkpoints
• TensorBoard logs
• generated .bin/.idx pretraining datasets
• generated validation Megatron indexed data
• HF exported weights
• local caches or __pycache__

When pushing from g0050, export the machine proxy if direct access fails:

export http_proxy=http://100.72.0.101:8888
export https_proxy=http://100.72.0.101:8888
export HTTP_PROXY=http://100.72.0.101:8888
export HTTPS_PROXY=http://100.72.0.101:8888

When pulling on g0033, use its proxy configuration if needed:

export http_proxy=http://10.20.34.2:3128
export https_proxy=http://10.20.34.2:3128
export HTTP_PROXY=http://10.20.34.2:3128
export HTTPS_PROXY=http://10.20.34.2:3128

Quick Sanity Checks Before Handoff

Run these before claiming the repo is ready:

bash -n scripts/*.sh
PYTHONPYCACHEPREFIX=/tmp/laoyao_pycache_check python3 -m py_compile \
  dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py \
  training/megatron_bridge/laoyao_2b_moe_pretrain.py \
  tools/*.py tools/hf_laoyao_moe/*.py

If py_compile fails due to permission denied under __pycache__, set PYTHONPYCACHEPREFIX as above instead of deleting unrelated cache files.