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:

```bash
/ssd/workspace/yi/laoyao_2b_moe
```

Current g0033 mirror/dev repo:

```bash
/mnt/beegfs/yi/laoyao_2b_moe
```

Current g0050 pretokenized data root:

```bash
/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset
```

Current main run:

```bash
/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
bash scripts/g0050_wait_and_tokenize_glm52_8192.sh
```

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

```bash
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:

```bash
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:

```bash
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:

```bash
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:

```bash
model/nemo_megatron/
```

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

## Training

Current primary resume/training entrypoint on g0050:

```bash
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:

```bash
docker/nemo/
```

Build entrypoint:

```bash
bash scripts/build_nemo_megatron_image.sh
```

The current image is:

```text
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:

```bash
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
bash scripts/resume_pretrain_8192_8gpu_mbs14.sh
```

Stop training container:

```bash
docker rm -f laoyao_pretrain_resume
```

Before killing, prefer checking the latest checkpoint:

```bash
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:

```bash
/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:

```bash
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:

```bash
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:

```bash
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:

```bash
tools/hf_laoyao_moe/
```

Export command template:

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

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:

```bash
docs/megatron_to_hf_export_notes.md
```

## Inference

Megatron-native inference server scripts:

```bash
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:

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

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:

```bash
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:

```bash
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:

```bash
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
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.