Add agent operating guide
This commit is contained in:
yi_lu
@@ -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 <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.
|
||||
Reference in New Issue
Block a user