diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..ed45450 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,406 @@ +# 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. + +## 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 +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.