docs: fix documentation errors across README and assets/docs

- Correct training CLI args: remove non-existent --adamw_beta1/2, fix --weight_decay
- Fix optimizer section: document MuonMix instead of plain AdamW
- Fix inference.md decode phase description (all groups, not largest)
- Fix dataflow.md H5Store description (no share_memory_ in code)
- Fix architecture.md class diagram: remove Task.stream_callback,
  EncoderConfig.use_gated_attention; update stop_ids docs
- Fix --min_rate default value description to match code (0.01)
- Update all document timestamps to 2026-07-09

README.md

@@ -103,9 +103,7 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
+    --weight_decay=0.1 \
-    --adamw_beta2=0.95 \
-    --adamw_weight_decay=0.01 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \

assets/docs/README-zh-CN.md

@@ -109,9 +109,7 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
+    --weight_decay=0.1 \
-    --adamw_beta2=0.95 \
-    --adamw_weight_decay=0.01 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \

assets/docs/architecture.md

@@ -63,7 +63,6 @@ classDiagram
             +Optional[int] n_heads
             +Optional[int] n_kv_heads
             +Optional[bool] use_qk_norm
-            +Optional[bool] use_gated_attention
             +str ffn_type
             +Optional[dict] rope_scaling
             +Optional[str] pooling_type
@@ -743,23 +742,22 @@ classDiagram
             +table_tensor(task_ids, device) Tensor
         }
 
-        class Task {
+    class Task {
-            +str task_id
+        +str task_id
-            +List prompt_ids
+        +List prompt_ids
-            +Optional[int] max_tokens
+        +Optional[int] max_tokens
-            +float temperature
+        +float temperature
-            +float top_p
+        +float top_p
-            +int top_k
+        +int top_k
-            +TaskStatus status
+        +TaskStatus status
-            +List output_ids
+        +List output_ids
-            +int input_tokens
+        +int input_tokens
-            +int output_tokens
+        +int output_tokens
-            +float arrival_time
+        +float arrival_time
-            +Optional[float] finish_time
+        +Optional[float] finish_time
-            +Optional[Callable] stream_callback
+        +int next_pos
-            +int next_pos
+        +is_finished(stop_ids) bool
-            +is_finished(stop_ids) bool
+    }
-        }
 
         class TaskStatus {
             <<enumeration>>
@@ -1249,4 +1247,4 @@ classDiagram
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers
 
-> Document Update Time: 2026-07-05
+> Document Update Time: 2026-07-09

assets/docs/dataflow.md

@@ -61,7 +61,7 @@ StoreFactory.create("bin")   → MmapStore
 StoreFactory.create("jsonl") → JsonlStore
 ```
 
-**H5Store**: Reads HDF5 files, supports `share_memory_()` for multi-process DataLoader workers (copies tensors to shared memory).
+**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage.
 
 **MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`.
 
@@ -109,4 +109,4 @@ DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_typ
 
 Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
 
-> Document Update Time: 2026-07-05
+> Document Update Time: 2026-07-09

assets/docs/inference.md

@@ -56,7 +56,7 @@ PageCache (paged KV cache with prefix sharing, alternative)
 1. Cleanup → Remove finished tasks, free KV cache slots/pages
 2. Refill  → Pop from waiting_queue, task_alloc resources, activate
 3. Prefill → Group by (prompt_len, start_pos), run full forward
-4. Decode  → Pick largest same-position group, single-token forward
+4. Decode  → Run single-token forward for each same-position group
 ```
 
 ## Sampling (Strategy Pattern)
@@ -249,4 +249,4 @@ async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[s
     print(token)
 ```
 
-> Document Update Time: 2026-07-05
+> Document Update Time: 2026-07-09

assets/docs/params.md

@@ -28,13 +28,17 @@
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
 | `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
 
-### Optimizer (AdamW)
+### Optimizer (MuonMix)
+
+Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`fused=True`).
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--adamw_beta1` | AdamW beta1 | 0.9 |
+| `--weight_decay` | Weight decay (applied to Muon matrix params; non-matrix use 0) | 0.1 |
-| `--adamw_beta2` | AdamW beta2 | 0.95 |
+| `--muon_momentum` | Muon momentum factor | 0.95 |
-| `--adamw_weight_decay` | AdamW weight decay | 0.01 |
+| `--muon_nesterov` | Enable Nesterov momentum for Muon | True |
+| `--muon_ns_steps` | Newton-Schulz iteration steps for Muon | 5 |
+| `--muon_adjust_lr` | Muon LR adjustment strategy (`original`, `match_rms_adamw`) | `match_rms_adamw` |
 
 ### Data Loading
 
@@ -104,7 +108,7 @@
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--schedule_type` | LR scheduler type (`cosine`, `sgdr`, `wsd`) | cosine |
-| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default) |
+| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default: 0.01) |
 | `--cycle_length` | SGDR first cycle length in steps | None (total_steps - warmup_steps) |
 | `--t_mult` | SGDR cycle length multiplier per restart | 2 |
 | `--stable_steps` | WSD stable plateau steps | None (required for wsd) |
@@ -126,9 +130,7 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
+    --weight_decay=0.1 \
-    --adamw_beta2=0.95 \
-    --adamw_weight_decay=0.01 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \
@@ -199,4 +201,4 @@ See [Preprocessing Guide](preprocessing.md) for config file format and examples.
 
 ---
 
-> Document Update Time: 2026-07-05
+> Document Update Time: 2026-07-09

assets/docs/preprocessing.md

@@ -361,4 +361,4 @@ Pipeline(
 ).run()
 ```
 
-> Document Update Time: 2026-07-05
+> Document Update Time: 2026-07-09

assets/docs/training.md

@@ -201,9 +201,7 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
+    --weight_decay=0.1 \
-    --adamw_beta2=0.95 \
-    --adamw_weight_decay=0.01 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \
@@ -214,4 +212,4 @@ nohup python scripts/tools/train.py \
 
 Full parameter reference at [params.md](params.md).
 
-> Document Update Time: 2026-07-05
+> Document Update Time: 2026-07-09