feat: add nanochat_exp experiment module with training tools and documentation

- Add tokenizer training and data processing utilities
- Add training monitoring and diagnosis tools
- Add parquet inspection and viewing tools
- Add documentation for data processing flow and known issues
- Update dataloader, dataset, and training scripts

Author: phoenixdong

examples/nanochat_exp/DATA_PROCESSING_FLOW.md

@@ -0,0 +1,223 @@
+# 数据处理流程说明
+
+## 数据处理流程概览
+
+整个数据处理流程分为以下几个阶段：
+
+```
+原始数据 → 数据加载 → Parquet 转换 → Tokenization → 训练批次
+```
+
+## 详细流程
+
+### 1. 原始数据阶段（dataset.py）
+
+**处理内容**：从 HuggingFace 数据集或本地 JSONL 文件加载原始数据
+
+**数据格式**：
+- HuggingFace 数据集：`{"text": "文档内容", ...}`
+- JSONL 文件：每行一个 JSON 对象，包含文本字段
+
+**处理步骤**：
+- 从 HuggingFace 下载或从本地加载数据集
+- 提取 `text` 字段（或自动识别文本字段）
+- 统一数据格式，只保留文本内容
+
+**输出**：包含 `text` 列的 Dataset 对象
+
+---
+
+### 2. Parquet 转换阶段（dataset.py - `convert_to_parquet()`）
+
+**处理内容**：将文本数据转换为 Parquet 格式，便于高效读取
+
+**数据格式**：
+- 输入：文本列表 `["文档1", "文档2", ...]`
+- 输出：Parquet 文件，每行一个文档的文本内容
+
+**处理步骤**：
+- 将文本按字符数分片（默认每片 ~250M 字符）
+- 写入 Parquet 文件（使用 Snappy 压缩）
+- 文件命名：`shard_00000.parquet`, `shard_00001.parquet`, ...
+
+**输出**：Parquet 文件目录 `parquet_shards/`
+
+**关键代码位置**：
+```python
+# dataset.py:358-457
+def convert_to_parquet(...):
+    # 提取 text 字段
+    text = example.get(text_column, "")
+    # 写入 Parquet
+    table = pa.Table.from_arrays([pa.array(texts)], names=["text"])
+    pq.write_table(table, filepath, compression="snappy")
+```
+
+---
+
+### 3. 数据加载阶段（dataloader.py - `parquets_iter_batched()`）
+
+**处理内容**：从 Parquet 文件读取文本数据
+
+**数据格式**：
+- 输入：Parquet 文件
+- 输出：文本批次 `["文档1", "文档2", ..., "文档128"]`
+
+**处理步骤**：
+- 按 DDP rank 分片读取（分布式训练）
+- 从 Parquet 的 `text` 列读取文本
+- 按 `tokenizer_batch_size`（默认 128）分批返回
+
+**输出**：文本批次迭代器
+
+**关键代码位置**：
+```python
+# dataset.py:475-504
+def parquets_iter_batched(...):
+    rg = pf.read_row_group(rg_idx)
+    texts = rg.column('text').to_pylist()  # 读取 text 列
+    yield texts
+```
+
+---
+
+### 4. Tokenization 阶段（dataloader.py + tokenizer.py）⭐
+
+**处理内容**：将文本转换为 token IDs（这是 tokenizer 处理的部分）
+
+**数据格式**：
+- 输入：文本批次 `["文档1", "文档2", ...]`
+- 输出：Token ID 列表 `[[1, 234, 567, ...], [1, 890, 123, ...], ...]`
+
+**处理步骤**：
+1. **获取 tokenizer**：`tokenizer = get_tokenizer()`
+2. **添加 BOS token**：在每个文档前添加 `<|bos|>` token
+3. **编码文本**：`tokenizer.encode(doc_batch, prepend=bos_token)`
+4. **返回 token IDs**：每个文档转换为一个 token ID 列表
+
+**关键代码位置**：
+```python
+# dataloader.py:108-114
+token_lists = tokenizer.encode(
+    doc_batch,              # 文本批次：["文档1", "文档2", ...]
+    prepend=bos_token,       # 在每个文档前添加 BOS token
+    num_threads=tokenizer_threads
+)
+# 返回：[[1, 234, 567, ...], [1, 890, 123, ...], ...]
+```
+
+**Tokenizer 处理的具体内容**：
+- ✅ **文本字符串** → **Token IDs**
+- ✅ 使用 BPE（Byte Pair Encoding）算法
+- ✅ 添加特殊 token（BOS、EOS 等）
+- ✅ 处理 Unicode 字符和多语言文本
+
+---
+
+### 5. 批次构建阶段（dataloader.py）
+
+**处理内容**：将 token IDs 组织成训练批次
+
+**数据格式**：
+- 输入：Token ID 流 `[1, 234, 567, 890, ...]`
+- 输出：训练批次 `(inputs, targets)` tensors
+
+**处理步骤**：
+1. **累积 tokens**：从 token buffer 中取出 `B * T + 1` 个 tokens
+2. **创建 inputs/targets**：
+   - `inputs = tokens[:-1]` （前 B*T 个 tokens）
+   - `targets = tokens[1:]` （后 B*T 个 tokens，用于预测）
+3. **重塑形状**：`(B, T)` - B 个样本，每个 T 个 tokens
+4. **移动到设备**：CPU → GPU（如果使用 CUDA）
+
+**关键代码位置**：
+```python
+# dataloader.py:117-143
+tokens = [token_buffer.popleft() for _ in range(needed_tokens)]  # B*T+1 个 tokens
+inputs_cpu = scratch[:-1].to(dtype=torch.int32)  # 前 B*T 个
+targets_cpu = scratch[1:]                         # 后 B*T 个
+inputs = inputs_cpu.view(B, T).to(device=device)  # 重塑为 (B, T)
+targets = targets_cpu.view(B, T).to(device=device)
+```
+
+---
+
+## Tokenizer 处理的数据部分总结
+
+### ✅ Tokenizer 处理的内容：
+
+1. **文本字符串** → **Token IDs**
+   - 输入：纯文本字符串（如 "Hello world"）
+   - 输出：整数列表（如 `[15496, 1917]`）
+
+2. **处理位置**：
+   - 在 `dataloader.py` 的 `tokenizing_distributed_data_loader()` 函数中
+   - 调用 `tokenizer.encode()` 方法
+
+3. **处理时机**：
+   - 从 Parquet 文件读取文本后
+   - 在构建训练批次之前
+
+4. **处理方式**：
+   - 批量处理（默认 128 个文档一批）
+   - 每个文档前添加 BOS token
+   - 使用多线程加速（默认 4 个线程）
+
+### ❌ Tokenizer 不处理的内容：
+
+- ❌ 数据下载和加载（dataset.py）
+- ❌ Parquet 文件转换（dataset.py）
+- ❌ 批次构建和 tensor 操作（dataloader.py）
+- ❌ GPU 内存管理（dataloader.py）
+
+---
+
+## 数据流示例
+
+```
+原始数据：
+{"text": "这是一个测试文档。"}
+
+↓ (dataset.py - convert_to_parquet)
+
+Parquet 文件：
+text: "这是一个测试文档。"
+
+↓ (dataloader.py - parquets_iter_batched)
+
+文本批次：
+["这是一个测试文档。"]
+
+↓ (dataloader.py + tokenizer.py - encode) ⭐ TOKENIZER 处理这里
+
+Token IDs：
+[[1, 234, 567, 890, 1234, 5678]]  # 1 是 BOS token
+
+↓ (dataloader.py - 批次构建)
+
+训练批次：
+inputs:  [[1, 234, 567, 890, 1234]]      # shape: (B, T)
+targets: [[234, 567, 890, 1234, 5678]]  # shape: (B, T)
+```
+
+---
+
+## 关键文件说明
+
+| 文件 | 职责 | 处理的数据部分 |
+|------|------|---------------|
+| `dataset.py` | 数据加载和转换 | 原始数据 → Parquet 文件 |
+| `dataloader.py` | 数据加载和批次构建 | Parquet 文件 → 训练批次 |
+| `tokenizer.py` | Tokenization | **文本字符串 → Token IDs** ⭐ |
+
+---
+
+## 总结
+
+**Tokenizer 处理的是数据的 Tokenization 阶段**，具体来说：
+
+1. **输入**：从 Parquet 文件读取的**文本字符串**
+2. **处理**：使用 BPE 算法将文本转换为 **Token IDs**
+3. **输出**：**Token ID 列表**，供后续批次构建使用
+
+这是整个数据处理流程中的**关键步骤**，将人类可读的文本转换为模型可以处理的数字序列。

examples/nanochat_exp/README.md

@@ -4,7 +4,7 @@ This directory provides integration between OpenSeek datasets and the [nanochat]
 
 ## Overview
 
-This module adapts OpenSeek datasets (specifically OpenSeek-Pretrain-100B) to work with nanochat's training pipeline. It provides:
+This module adapts OpenSeek datasets (specifically OpenSeek-Pretrain-Data-Examples) to work with nanochat's training pipeline. It provides:
 
 - **Dataset conversion**: Converts OpenSeek datasets from HuggingFace format to parquet format compatible with nanochat
 - **Data loader**: Provides a data loader compatible with nanochat's training interface
@@ -23,12 +23,25 @@ This module adapts OpenSeek datasets (specifically OpenSeek-Pretrain-100B) to wo
 
 2. **Python dependencies**（推荐 Python≥3.10、PyTorch≥2.1，与 nanochat 官方示例保持一致）：
    ```bash
-   pip install pyarrow datasets huggingface_hub
+   # 推荐使用虚拟环境（避免权限警告）
+   python -m venv venv
+   source venv/bin/activate  # Linux/Mac
+   # 或 venv\Scripts\activate  # Windows
+   
+   # OpenSeek 使用 HuggingFace tokenizers 库（易于安装，无需 Rust 编译）
+   pip install pyarrow datasets huggingface_hub tokenizers>=0.22.0
    ```
-
-3. **OpenSeek dataset**: Download the OpenSeek-Pretrain-100B dataset:
+   
+   > **重要**: OpenSeek 使用 HuggingFace `tokenizers` 库替代 nanochat 的 `rustbpe` 模块。
+   > - **优势**: `tokenizers` 库更容易安装，只需 `pip install tokenizers`，无需 Rust 编译
+   > - **兼容性**: 完全兼容 nanochat 的 tokenizer 接口
+   > - **性能**: HuggingFace tokenizers 库性能优秀，基于 Rust 实现但提供预编译的 Python 包
+   > 
+   > 如果以 root 用户运行 pip，会收到警告。建议使用虚拟环境，或使用 `--root-user-action=ignore` 选项（仅在明确知道自己在做什么时使用）。
+
+3. **OpenSeek dataset**: Download the OpenSeek-Pretrain-Data-Examples dataset:
    - Option 1: Download from HuggingFace (automatic)
-   - Option 2: Download manually to `OpenSeek-Pretrain-100B/` directory in the project root
+   - Option 2: Download manually to `OpenSeek-Pretrain-Data-Examples/` directory in the project root
 
 ## Quick Start
 
@@ -49,37 +62,62 @@ First, convert the OpenSeek dataset to the parquet format expected by nanochat:
 
 ```bash
 # From OpenSeek root directory
+# 对于示例数据集，默认使用 -1（处理所有数据，不限制 shards 数量）
 python -m examples.nanochat_exp.dataset \
-    --dataset "BAAI/OpenSeek-Pretrain-100B" \
-    --num-shards 240 \
+    --dataset "BAAI/OpenSeek-Pretrain-Data-Examples" \
+    --num-shards -1 \
     --streaming
 ```
 
 Or if you have the dataset locally:
 
 ```bash
 python -m examples.nanochat_exp.dataset \
-    --dataset ./OpenSeek-Pretrain-100B \
-    --num-shards 240
+    --dataset ./OpenSeek-Pretrain-Data-Examples \
+    --num-shards -1
 ```
 
 This will create parquet shards in `~/.cache/openseek_nanochat/parquet_shards/` (or the directory specified by `OPENSEEK_NANOCHAT_DATA_DIR`).
 
-### 2. Modify nanochat Training Scripts
+### 2. Train Tokenizer (Optional but Recommended)
 
-To use OpenSeek data with nanochat, you need to modify nanochat's training scripts to use our data loader. In `nanochat/scripts/base_train.py`, change:
+Train a BPE tokenizer from your data using HuggingFace tokenizers (no rustbpe needed):
 
-```python
-from nanochat.dataloader import tokenizing_distributed_data_loader
+```bash
+# From OpenSeek root directory
+python -m examples.nanochat_exp.tok_train \
+    --vocab-size 50257 \
+    --data-dir ~/.cache/openseek_nanochat/parquet_shards
+```
+
+This will create `tokenizer.json` in the tokenizer directory. The script uses HuggingFace tokenizers library, which is much easier to install than rustbpe (no Rust compilation needed).
+
+### 3. Modify nanochat Training Scripts
+
+To use OpenSeek data with nanochat, you need to modify nanochat's training scripts to use our data loader and tokenizer. You can use the automated patch script:
+
+```bash
+# From OpenSeek root directory
+python -m examples.nanochat_exp.patch_nanochat --nanochat-path /path/to/nanochat
 ```
 
-to:
+This will automatically modify `nanochat/scripts/base_train.py` to:
+- Use OpenSeek's dataloader (from `examples.nanochat_exp.dataloader`)
+- Use OpenSeek's tokenizer (from `examples.nanochat_exp.tokenizer`, uses HuggingFace tokenizers, no rustbpe)
+
+Alternatively, you can manually modify `nanochat/scripts/base_train.py`:
 
 ```python
+# Change this:
+from nanochat.dataloader import tokenizing_distributed_data_loader
+from nanochat.tokenizer import get_tokenizer
+
+# To this:
 from examples.nanochat_exp.dataloader import tokenizing_distributed_data_loader
+from examples.nanochat_exp.tokenizer import get_tokenizer
 ```
 
-### 3. Run Training
+### 4. Run Training
 
 After modifying the import, you can run nanochat training as usual:
 
@@ -100,6 +138,9 @@ examples/nanochat_exp/
 ├── __init__.py              # Module initialization
 ├── dataset.py               # Dataset conversion and loading utilities
 ├── dataloader.py            # Data loader compatible with nanochat
+├── tokenizer.py             # Tokenizer wrapper using HuggingFace tokenizers (easy install)
+├── tok_train.py             # Tokenizer training script (uses HuggingFace tokenizers, no rustbpe)
+├── patch_nanochat.py        # Script to patch nanochat's base_train.py (replaces dataloader & tokenizer)
 ├── run_openseek_exp.sh      # Experiment runner script
 ├── train_wrapper.py         # Training wrapper script
 └── README.md                # This file
@@ -109,10 +150,37 @@ This module can be imported via `examples.nanochat_exp` when the repository root
 
 ## Design Notes & Configuration
 
+### Tokenizer 说明
+
+OpenSeek 使用 **HuggingFace tokenizers 库**替代 nanochat 的 `rustbpe` 模块：
+
+- **优势**: 
+  - 更容易安装：只需 `pip install tokenizers`，无需 Rust 编译
+  - 完全兼容：提供与 nanochat tokenizer 相同的接口
+  - 性能优秀：基于 Rust 实现，提供预编译的 Python 包
+  
+- **使用方式**: 
+  - `examples.nanochat_exp.tokenizer.get_tokenizer()` 会自动使用 HuggingFace tokenizers
+  - 如果找不到 tokenizer.json，会尝试从标准位置加载，或创建默认 tokenizer
+  - 数据加载器会自动使用新的 tokenizer，无需修改代码
+  - 训练 tokenizer 使用 `python -m examples.nanochat_exp.tok_train`（无需 rustbpe）
+
+- **训练 Tokenizer**: 
+  - 使用 `examples.nanochat_exp.tok_train` 脚本训练 BPE tokenizer
+  - 完全替代 nanochat 的 `tok_train.py`，无需 rustbpe
+  - 支持自定义词汇表大小、最小频率等参数
+  - 输出标准的 `tokenizer.json` 文件，兼容 HuggingFace tokenizers
+
+- **兼容性**: 
+  - 如果系统中已安装 rustbpe，代码会优先尝试使用 HuggingFace tokenizers
+  - 如果 HuggingFace tokenizers 不可用，会回退到 nanochat 的 rustbpe（如果可用）
+  - 训练脚本完全独立，不依赖 rustbpe
+
 ### 数据加载接口说明
 
 - `examples.nanochat_exp.dataloader.tokenizing_distributed_data_loader()` 保持与 nanochat 原生接口兼容，可直接替换 import。
-- 默认 tokenizer 仍由 nanochat 配置决定；若需变更 tokenizer、最大长度或 padding 规则，可在 `train_wrapper.py` 中调整对应函数，再由训练脚本调用。
+- Tokenizer 使用 HuggingFace tokenizers 库（易于安装），完全兼容 nanochat 的 tokenizer 接口。
+- 若需变更 tokenizer、最大长度或 padding 规则，可在 `train_wrapper.py` 中调整对应函数，再由训练脚本调用。
 - 批量大小、梯度累积等策略建议继续放在 nanochat 脚本侧统一配置，确保与数据加载逻辑一致。
 
 ### Environment Variables