# Phase 10 LLMs from Scratch: Complete Input-Output Pipeline Implementation

> Implement the complete input-output pipeline for LLMs from scratch. Learn how to bridge raw corpora, distributed training, and production inference in Phase 10 LLMs from Scratch.

- Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch)
- Tags: tutorial
- Published: 2026-08-28

---

**Phase 10 LLMs from Scratch** delivers an end-to-end blueprint for building large language models from the ground up, focusing on the **input-output (IO) infrastructure** that bridges raw corpora, distributed training, and production inference.

In the `rohitg00/ai-engineering-from-scratch` repository, this phase (located at `phases/06-io-llms-from-scratch/`) provides runnable implementations of streaming data loaders, byte-pair encoding tokenizers, multi-head attention mechanisms, and sharded checkpointing systems essential for training at scale.

## Data IO and Pre-processing Pipeline

Efficient data handling forms the foundation of any scalable LLM training system.

### Streaming Large Corpora Without Memory Exhaustion

The `stream_corpus()` function in [`phases/06-io-llms-from-scratch/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/06-io-llms-from-scratch/code/main.py) demonstrates memory-efficient loading of multi-gigabyte datasets using chunked reading:

```python
def stream_corpus(path: str, chunk_size: int = 1 << 20):
    """Yield UTF‑8 lines from a potentially‑gigantic text file without loading it all."""
    with open(path, "rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk: break
            for line in chunk.split(b"\n"):
                yield line.decode("utf-8", errors="ignore")

```

This generator pattern processes text in 1MB chunks, preventing out-of-memory errors when handling web-scale datasets while maintaining Unicode integrity through robust error handling.

### Token-Level Batching and Sharding Strategies

The phase implements **dataset sharding** that splits massive text corpora into manageable chunks. This enables parallel preprocessing across distributed training nodes while preserving document boundaries and supporting resumable data iterators.

## Tokenizer Design and BPE Implementation

### Byte-Pair Encoding from Scratch

Located in [`phases/06-io-llms-from-scratch/code/tokenizer.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/06-io-llms-from-scratch/code/tokenizer.py), the `BPETokenizer` class constructs subword vocabularies through iterative pair merging:

```python
class BPETokenizer:
    def __init__(self, vocab_path: str):
        self.vocab = self._load_vocab(vocab_path)      # dict token → id

        self.inverse = {v: k for k, v in self.vocab.items()}
    
    def encode(self, text: str) -> list[int]:
        # Very simplified greedy merging loop

        tokens = list(text)
        while True:
            pairs = [(tokens[i], tokens[i+1]) for i in range(len(tokens)-1)]
            best = min(pairs, key=self._pair_score, default=None)
            if best is None or self._pair_score(best) >= self.threshold:
                break
            i = tokens.index(best[0])
            tokens[i:i+2] = ["".join(best)]
        return [self.vocab[t] for t in tokens if t in self.vocab]

```

The implementation handles **unknown tokens** through character-level fallback strategies and provides bidirectional conversion between token IDs and strings, directly impacting model compression ratios.

### Embedding Layer Integration

The phase connects tokenizer output to **learned token embeddings** and **positional encodings** (both sinusoidal and learned variants), translating discrete token IDs into continuous vector representations suitable for transformer processing.

## Transformer Architecture Implementation

### Multi-Head Self-Attention Mechanism

The core computational unit resides in [`phases/06-io-llms-from-scratch/code/transformer.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/06-io-llms-from-scratch/code/transformer.py):

```python
class MultiHeadSelfAttention(nn.Module):
    def __init__(self, dim: int, heads: int = 8):
        super().__init__()
        self.heads = heads
        self.proj_q = nn.Linear(dim, dim)
        self.proj_k = nn.Linear(dim, dim)
        self.proj_v = nn.Linear(dim, dim)
        self.out   = nn.Linear(dim, dim)

    def forward(self, x):
        B, T, C = x.shape
        q = self.proj_q(x).reshape(B, T, self.heads, C//self.heads).transpose(1,2)
        k = self.proj_k(x).reshape(B, T, self.heads, C//self.heads).transpose(1,2)
        v = self.proj_v(x).reshape(B, T, self.heads, C//self.heads).transpose(1,2)
        attn = (q @ k.transpose(-2, -1)) / math.sqrt(C//self.heads)
        attn = attn.softmax(dim=-1)
        y = (attn @ v).transpose(1,2).reshape(B, T, C)
        return self.out(y)

```

This implementation features **scaled dot-product attention**, **multi-head parallelism**, and proper tensor reshaping for batch processing. The architecture includes **residual connections** and **layer normalization** to stabilize deep network training.

### Feed-Forward Networks and Layer Normalization

Complete transformer blocks combine the attention mechanism with position-wise feed-forward networks, applying dropout and normalization layers as specified in the original transformer architecture.

## Training Loop Optimization and Checkpointing

### Mixed-Precision Training and Gradient Management

Phase 10 implements **FP16 mixed-precision training** with gradient accumulation to maximize GPU memory utilization. The training loop incorporates **learning-rate warmup** and **cosine decay schedules** to ensure optimization stability during the critical early training phases.

### Sharded Model IO for Distributed Environments

For multi-GPU training scenarios, [`phases/06-io-llms-from-scratch/code/checkpoint.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/06-io-llms-from-scratch/code/checkpoint.py) provides sharded checkpointing that splits model weights across multiple files:

```python
def save_sharded(model: nn.Module, path: Path, shard_size: int = 100_000):
    """Save model weights in multiple smaller binary shards to avoid OOM."""
    flat = torch.nn.utils.parameters_to_vector(model.parameters())
    for i in range(0, flat.numel(), shard_size):
        shard = flat[i:i+shard_size].cpu().numpy()
        torch.save(shard, path / f"shard_{i//shard_size}.pt")

def load_sharded(model: nn.Module, path: Path):
    flat = []
    for shard_file in sorted(path.glob("shard_*.pt")):
        flat.append(torch.from_numpy(torch.load(shard_file)))
    flat_tensor = torch.cat(flat)
    torch.nn.utils.vector_to_parameters(flat_tensor, model.parameters())

```

These functions enable **safe resumption** of training runs without single-file size limitations, crucial for models with billions of parameters distributed across many nodes.

## Inference and Serving Infrastructure

### Streaming Text Generation with Top-K Sampling

The `generate()` function in [`phases/06-io-llms-from-scratch/code/generate.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/06-io-llms-from-scratch/code/generate.py) implements efficient **autoregressive generation** with configurable sampling strategies:

```python
def generate(model, tokenizer, prompt: str, max_new_tokens: int = 50, top_k: int = 40):
    ids = tokenizer.encode(prompt)
    ids = torch.tensor(ids, dtype=torch.long, device=model.device).unsqueeze(0)
    for _ in range(max_new_tokens):
        logits = model(ids)[:, -1, :]                      # next‑token distribution

        probs  = torch.softmax(logits, dim=-1)
        topk   = torch.topk(probs, top_k)
        token  = torch.multinomial(topk.values, 1)
        ids    = torch.cat([ids, topk.indices.gather(-1, token)], dim=1)
    return tokenizer.decode(ids.squeeze().tolist())

```

This implementation supports **nucleus (top-p) sampling** variants and **token-level stop conditions**, providing the foundation for low-latency streaming APIs.

### Production API Integration

The phase demonstrates exposing the inference engine through simple REST/RPC endpoints, illustrating request batching and response streaming patterns suitable for production deployment.

## Evaluation and Safety Mechanisms

### Perplexity Calculation and Benchmarks

Evaluation scripts in the phase calculate **perplexity** on held-out validation sets to monitor training progress. The codebase supports downstream task benchmarking for few-shot question answering and language understanding evaluation.

### Safety Filters and Alignment Basics

Basic **toxicity detection** filters and prompt-level instruction-override handling provide scaffolding for responsible AI deployment. The phase includes preliminary data structures for **RLHF (Reinforcement Learning from Human Feedback)** integration.

## Summary

- **Phase 10 LLMs from Scratch** covers the complete stack from raw text ingestion to production serving, emphasizing IO efficiency at every boundary.
- **Data streaming** via `stream_corpus()` handles web-scale corpora without memory exhaustion, while **BPE tokenization** implements subword vocabulary construction from scratch.
- **Sharded checkpointing** in [`checkpoint.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/checkpoint.py) solves distributed training persistence challenges through `save_sharded()` and `load_sharded()` utilities.
- **Multi-head self-attention** and complete **transformer blocks** are implemented in pure PyTorch within [`transformer.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/transformer.py).
- **Top-k sampling** and **streaming generation** in [`generate.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/generate.py) provide production-ready inference with configurable randomness.
- The module includes comprehensive **unit tests** in [`test_main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/test_main.py) verifying data loading, tokenization, training steps, and checkpoint resume functionality.

## Frequently Asked Questions

### What prerequisites are needed for Phase 10 LLMs from Scratch?

Phase 10 requires proficiency in Python and PyTorch, plus fundamental understanding of neural network architectures and attention mechanisms. Familiarity with Unix file systems and distributed computing concepts is essential for implementing the sharded checkpointing and data streaming components effectively.

### How does the sharded checkpointing system handle distributed training failures?

The `save_sharded()` function splits model parameters into smaller binary shards (defaulting to 100,000 parameters per file) stored as `shard_{n}.pt` files. During resumption, `load_sharded()` reconstructs the full parameter tensor by concatenating shards in sorted order using `torch.nn.utils.vector_to_parameters()`. This prevents single-point-of-failure if individual files corrupt and bypasses single-file filesystem size limits common in large model training.

### What tokenizer algorithm does this phase implement?

The repository implements **Byte-Pair Encoding (BPE)** from scratch in [`phases/06-io-llms-from-scratch/code/tokenizer.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/06-io-llms-from-scratch/code/tokenizer.py). The `BPETokenizer` class performs greedy pair merging to construct subword vocabularies, handling unknown tokens through character-level fallback strategies and providing efficient bidirectional mapping between text sequences and token ID lists.

### Can the inference generation function handle batch processing?

While the provided `generate()` function demonstrates single-prompt streaming for clarity, the underlying `MultiHeadSelfAttention` implementation supports arbitrary batch dimensions (the `B` parameter in input tensors with shape `(B, T, C)`). Production extensions would batch multiple prompts together and apply proper attention masking to maximize GPU throughput during high-volume inference serving.