fareedkhan-dev--train-llm-from-scratch
wehub-resource-sync
93 行
4.4 KiB
Markdown
93 行
4.4 KiB
Markdown
<!-- omit in toc -->
|
|
# Inference & Chat
|
|
|
|
Training is only satisfying if you can actually *talk* to the result. The original
|
|
[`generate_text.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/generate_text.py) does raw continuation for the base model, but it's
|
|
hard-wired to the legacy config and has no chat template — so I added a small inference layer that loads
|
|
**any** stage checkpoint (base / SFT / DPO / PPO / GRPO) and talks to it correctly.
|
|
|
|
For the underlying decoding loop, context cropping, temperature, and stop-token behavior, read
|
|
[Generation & Sampling](foundations/generation.md).
|
|
|
|
|
|
|
|
<details>
|
|
<summary>Mermaid source (live, editable)</summary>
|
|
|
|
```mermaid
|
|
flowchart LR
|
|
CK[(any checkpoint)]:::ckpt --> LD[load_model_from_ckpt<br/>dims from stored cfg]:::proc
|
|
LD --> MODE{chat or raw?}:::proc
|
|
MODE -->|instruction model| CT[wrap in chat template]:::proc
|
|
MODE -->|base model| RAW[raw prefix]:::proc
|
|
CT --> GEN{{generate<br/>temperature / top-p / greedy}}:::model
|
|
RAW --> GEN
|
|
GEN --> DEC([decode → reply]):::eval
|
|
classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222;
|
|
classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52;
|
|
classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00;
|
|
classDef eval fill:#e8d6ff,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
|
```
|
|
|
|
</details>
|
|
|
|
## Load any checkpoint by its stored config
|
|
|
|
[`load_model_from_ckpt`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py) reads the model dimensions from the
|
|
checkpoint's saved `cfg`, so you never re-specify `n_embed`/`n_blocks`, and it tolerates DDP /
|
|
reward-head key prefixes:
|
|
|
|
```python
|
|
ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
|
|
cfg = {**(ck.get("cfg") or {}), **(overrides or {})}
|
|
model = Transformer(n_head=cfg["n_head"], n_embed=cfg["n_embed"], ...)
|
|
state = {k.removeprefix("module.").removeprefix("transformer."): v for k, v in state.items()}
|
|
```
|
|
|
|
## Chat vs raw
|
|
|
|
[`generate_reply`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py#L37) has two modes, reusing the same tested
|
|
generation core as training/eval ([`batched_generate`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/evaluation.py#L24)):
|
|
|
|
- **chat** (default) — wraps your text in the chat template (optionally with a `system` message) and
|
|
returns the decoded assistant turn. Use this for SFT/DPO/PPO/GRPO checkpoints.
|
|
- **raw** (`--raw`) — treats your text as a prefix and returns the base model's continuation (no
|
|
template). Use this for `base_pretrained.pt`.
|
|
|
|
```python
|
|
if raw:
|
|
ids = get_tokenizer().encode_ordinary(user_text)
|
|
else:
|
|
ids = encode_prompt([{"role": "user", "content": user_text}]) # ...ends at <|assistant|>
|
|
out = batched_generate(model, [ids], max_new_tokens, device=device,
|
|
temperature=temperature, top_k=top_k, top_p=top_p, greedy=greedy)
|
|
```
|
|
|
|
Decoding is defensive — [`decode`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py) drops the EOT terminator and
|
|
any padding-vocab ids (the model's vocab is padded to 50304 but r50k_base only decodes 0–50255).
|
|
|
|
## The CLI
|
|
|
|
[`scripts/chat.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/chat.py) is one-shot or an interactive REPL:
|
|
|
|
```bash
|
|
# instruction-tuned models (chat template applied automatically)
|
|
PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt --prompt "What is 13 + 29?"
|
|
PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/grpo.pt --prompt "..." --greedy
|
|
# base-model continuation
|
|
PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/base_pretrained.pt --raw --prompt "Once upon a time"
|
|
# interactive REPL (omit --prompt)
|
|
PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt
|
|
```
|
|
|
|
Sampling controls: `--temperature`, `--top_p`, `--top_k`, or `--greedy` for deterministic argmax. Runs
|
|
on `--device cuda` or `cpu` (both verified).
|
|
|
|
## Sampling knobs, briefly
|
|
|
|
- **greedy** — reproducible, best for eval / math (`--greedy`).
|
|
- **temperature** — higher = more random; ~`0.7–1.0` for open-ended chat.
|
|
- **top_p / top_k** — nucleus / top-k truncation to cut the long tail of unlikely tokens.
|
|
|
|
That's the full loop: pretrain → align → reason → measure → chat. Back to the [overview](README.md).
|