wehub-resource-sync
162 lines
7.1 KiB
Python
162 lines
7.1 KiB
Python
"""
|
||
Generate every README diagram as a Mermaid .mmd file, then render to PNG with the hand-drawn
|
||
look. Reproducible: edit a graph below, re-run, re-render.
|
||
|
||
Render step (run from this folder, needs Node.js for npx):
|
||
for f in *.mmd; do
|
||
npx -y @mermaid-js/mermaid-cli -i "$f" -o "${f%.mmd}.png" -c mmdc.json -p puppeteer.json -b white -s 3
|
||
done
|
||
"""
|
||
import os
|
||
|
||
OUT = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
# Shared palette: soft fill + strong border per role (data, store, proc, model, rl, loss, ckpt, eval, out, op).
|
||
PALETTE = """
|
||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||
classDef proc fill:#cfe3fb,stroke:#1565c0,stroke-width:2px,color:#0d2c52;
|
||
classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00;
|
||
classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500;
|
||
classDef loss fill:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||
"""
|
||
|
||
DIAGRAMS = {
|
||
"00_pipeline": """flowchart LR
|
||
PILE([The Pile<br/>raw text]):::data --> PRE{{Pretrain<br/>base ~400M}}:::model
|
||
PRE --> SFT{{SFT<br/>instruct}}:::model
|
||
SFT --> RM{{Reward Model<br/>Bradley-Terry}}:::rl
|
||
SFT --> DPO{{DPO / ORPO / KTO}}:::rl
|
||
SFT --> GRPO{{GRPO / RLVR}}:::rl
|
||
RM -->|reward signal| PPO{{PPO}}:::rl
|
||
PPO --> EVAL([GSM8K eval<br/>+ chat]):::eval
|
||
DPO --> EVAL
|
||
GRPO --> EVAL
|
||
""",
|
||
|
||
"01_data": """flowchart LR
|
||
A1([Pretrain: The Pile]):::data --> A2[tiktoken r50k_base<br/>stream + encode]:::proc --> A3[(pile_train.h5<br/>flat int32 tokens)]:::store
|
||
B1([SFT: Alpaca · Dolly · GSM8K]):::data --> B2[chat template<br/>+ mask the prompt]:::proc --> B3[(sft_packed.h5<br/>tokens + loss_mask)]:::store
|
||
C1([Prefs: HH-RLHF · UltraFeedback]):::data --> C2[split prompt /<br/>chosen / rejected]:::proc --> C3[(preferences.jsonl)]:::store
|
||
D1([RL: GSM8K · arithmetic]):::data --> D2[keep the gold answer]:::proc --> D3[(rl_prompts.jsonl)]:::store
|
||
""",
|
||
|
||
"02_tokenization": """flowchart LR
|
||
T1["raw text<br/>'Once upon a time'"]:::data --> T2["tiktoken<br/>r50k_base"]:::proc
|
||
T2 --> T3["token ids<br/>[7454, 2402, 257, 640]"]:::store
|
||
T3 --> T4["append <|endoftext|><br/>(id 50256)"]:::proc
|
||
T4 --> T5["fixed windows of<br/>context_length + 1"]:::store
|
||
T5 --> T6["x = window[:-1]<br/>y = window[1:]"]:::out
|
||
""",
|
||
|
||
"03_mlp": """flowchart LR
|
||
M1["x (B, T, n_embed)"]:::data --> M2["Linear<br/>n_embed -> 4 · n_embed"]:::proc
|
||
M2 --> M3["ReLU"]:::proc
|
||
M3 --> M4["Linear<br/>4 · n_embed -> n_embed"]:::proc
|
||
M4 --> M5["out (B, T, n_embed)"]:::out
|
||
""",
|
||
|
||
"04_attention_head": """flowchart LR
|
||
X["x : token vectors<br/>(B, T, C)"]:::data --> Q["Q = x Wq"]:::proc
|
||
X --> K["K = x Wk"]:::proc
|
||
X --> V["V = x Wv"]:::proc
|
||
Q --> S["scores = Q Kᵀ / √d"]:::rl
|
||
K --> S
|
||
S --> M["causal mask<br/>+ softmax"]:::rl
|
||
M --> O["weights · V"]:::rl
|
||
V --> O
|
||
O --> H["head output<br/>(B, T, head_size)"]:::out
|
||
""",
|
||
|
||
"05_multi_head_attention": """flowchart LR
|
||
X["x (B, T, C)"]:::data --> H1["head 1"]:::proc
|
||
X --> H2["head 2"]:::proc
|
||
X --> H3["... (n_head heads)"]:::proc
|
||
H1 --> CC["concat heads<br/>(B, T, C)"]:::rl
|
||
H2 --> CC
|
||
H3 --> CC
|
||
CC --> PJ["output projection<br/>Linear C -> C"]:::proc
|
||
PJ --> HO["multi-head output"]:::out
|
||
""",
|
||
|
||
"06_transformer_block": """flowchart LR
|
||
X["x"]:::data --> LN1["LayerNorm"]:::proc --> AT["Multi-Head Attention"]:::rl --> A1(("+")):::op
|
||
X --> A1
|
||
A1 --> LN2["LayerNorm"]:::proc --> ML["MLP"]:::proc --> A2(("+")):::op
|
||
A1 --> A2
|
||
A2 --> BO["block output"]:::out
|
||
""",
|
||
|
||
"07_transformer": """flowchart LR
|
||
I["token ids (B, T)"]:::data --> TE["token embedding"]:::proc
|
||
I --> PE["position embedding"]:::proc
|
||
TE --> SUM(("+")):::op
|
||
PE --> SUM
|
||
SUM --> BLK["N_BLOCKS ×<br/>Transformer Block"]:::model
|
||
BLK --> LNF["final LayerNorm"]:::proc
|
||
LNF --> LMH["lm_head<br/>Linear -> vocab_size"]:::proc
|
||
LMH --> LOG["logits (B, T, vocab)"]:::out
|
||
LOG -. targets present .-> CE["cross-entropy loss"]:::loss
|
||
""",
|
||
|
||
"08_training_loop": """flowchart LR
|
||
D[(pile_train.h5)]:::store --> IT["get_batch_iterator<br/>random windows"]:::proc --> FW["forward<br/>(bf16 autocast)"]:::model --> CE["cross-entropy"]:::loss --> BW["backward<br/>× grad_accum"]:::proc --> CL["clip grad norm 1.0"]:::proc --> ST["AdamW step<br/>cosine LR + warmup"]:::model --> CK[(checkpoint)]:::ckpt
|
||
ST -. next step .-> IT
|
||
""",
|
||
|
||
"09_sft": """flowchart LR
|
||
S0[(sft_packed.h5<br/>tokens + loss_mask)]:::store --> S1["Transformer forward"]:::model --> S2["shift: predict t+1"]:::proc --> S3["per-token cross-entropy"]:::loss
|
||
S0 --> S4["loss_mask = 1 only on<br/>assistant tokens"]:::proc --> S5["mean over masked tokens"]:::loss
|
||
S3 --> S5 --> S6["AdamW step"]:::model
|
||
""",
|
||
|
||
"10_reward_model": """flowchart LR
|
||
R0["prompt + chosen / rejected"]:::data --> R1["SFT backbone<br/>forward_hidden"]:::model --> R2["take last real token"]:::proc --> R3["reward head<br/>Linear -> 1"]:::proc
|
||
R3 --> R4["r_chosen , r_rejected"]:::rl --> R5["Bradley-Terry loss<br/>-log σ(r_chosen - r_rejected)"]:::loss --> R6["AdamW step"]:::model
|
||
""",
|
||
|
||
"11_dpo": """flowchart LR
|
||
P0["chosen / rejected pair"]:::data --> P1["policy (trainable)<br/>sequence log-probs"]:::model
|
||
P0 --> P2["reference (frozen SFT)<br/>sequence log-probs"]:::ckpt
|
||
P1 --> P3["DPO loss<br/>-log σ(β · Δ)"]:::loss
|
||
P2 --> P3 --> P4["AdamW step"]:::model
|
||
""",
|
||
|
||
"12_ppo": """flowchart LR
|
||
Q0([GSM8K prompts]):::data --> Q1["rollout<br/>generate + log-probs"]:::model --> Q2["score<br/>verifier or reward model"]:::rl --> Q3["+ per-token KL to ref"]:::rl --> Q4["GAE<br/>advantages + returns"]:::proc --> Q5["clipped policy + value<br/>(K epochs)"]:::model --> Q6[(ppo.pt)]:::ckpt
|
||
Q5 -. sync old policy .-> Q1
|
||
""",
|
||
|
||
"13_grpo": """flowchart LR
|
||
G0([prompt]):::data --> G1["sample a GROUP<br/>of G answers"]:::model --> G2["verifier reward<br/>per answer"]:::rl --> G3["group advantage<br/>(r - mean) / std"]:::proc --> G4["clipped surrogate<br/>+ k3 KL to ref"]:::loss --> G5["policy update"]:::model
|
||
G5 -. next prompt .-> G0
|
||
""",
|
||
|
||
"14_evaluation": """flowchart LR
|
||
E0[(stage checkpoint)]:::ckpt --> E1["greedy generate"]:::model --> E2["extract_answer<br/><answer>N</answer>"]:::proc --> E3{"== gold?"}:::rl --> E4["GSM8K accuracy<br/>Base -> SFT -> ... -> GRPO"]:::eval
|
||
""",
|
||
|
||
"15_inference": """flowchart LR
|
||
C0[(any checkpoint)]:::ckpt --> C1["load dims from cfg"]:::proc --> C2{"chat or raw?"}:::rl
|
||
C2 -->|instruct| C3["wrap in chat template"]:::proc
|
||
C2 -->|base| C4["raw prefix"]:::proc
|
||
C3 --> C5["generate<br/>temperature / top-p / greedy"]:::model
|
||
C4 --> C5 --> C6["decode -> reply"]:::out
|
||
""",
|
||
}
|
||
|
||
def main():
|
||
for name, body in DIAGRAMS.items():
|
||
path = os.path.join(OUT, name + ".mmd")
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
f.write(body.rstrip() + "\n" + PALETTE)
|
||
print("wrote", name + ".mmd")
|
||
print("total:", len(DIAGRAMS), "diagrams")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|