DEV Community

Abhishek Banerjee
Abhishek Banerjee

Posted on Originally published at Medium on

Deconstructing the Black Box: A Comprehensive Guide to Building LLMs from Scratch in PyTorch

#aiarchitecture #softwareengineering #machinelearning #generativeaiusecases

The Strategic Case for “From Scratch” Implementation

For a Senior Staff Engineer, the transition from consuming black-box APIs to first-principles implementation is the difference between a practitioner and an architect. While high-level abstractions like Hugging Face are sufficient for rapid prototyping, they obscure the hardware-level bottlenecks and memory constraints that define production-grade systems. Implementing a Large Language Model (LLM) from scratch in PyTorch guided by the Raschka framework provides the granular visibility required to optimize GPU utilization and avoid common pitfalls like DataLoader starvation.

A deep understanding of the “humble byte” allows for sophisticated optimizations such as memory-efficient weight loading, which is critical when serving models on consumer-grade hardware. Furthermore, first-principles knowledge is the only defense against subtle hardware-specific discrepancies, such as the floating-point precision variations observed between CPU and MPS (Metal Performance Shaders) devices (Source: Ch 5, bonus 19). Every modern frontier model from Llama 3.2 to Qwen 3.5 is built upon these foundational tensor operations. By mastering the assembly of these components, we gain the capability to debug architectural bottlenecks at the kernel level rather than guessing behind an API wall.

Part I: Text Ingestion and Embedding Dynamics

The sensory interface of an LLM begins with the conversion of discrete linguistic tokens into continuous vector spaces. This stage defines the model’s representational capacity and its initial memory footprint.

Technical Deep Dive Tokenisation

Modern architectures utilize Byte Pair Encoding (BPE) rather than simple character or word-level tokenization. BPE is superior because it balances vocabulary size and sequence length. By iteratively merging frequent byte pairs, BPE creates a vocabulary that represents common words as single units while decomposing rare words into sub-word components. This prevents “out-of-vocabulary” errors and maintains high information density in the input tensor x \in \mathbb{R}^{B \times T}.

PyTorch Data Architecture

We utilize DataLoader objects to manage the context window (T). The relationship between input_ids and target_ids is a shifted-by-one mapping, which is fundamental for autoregressive training.

ASCII Tensor Flow: Input to Target Mapping

Sample Sequence: "The cat sat on"
------------------------------------------------------------
Batch Dimension [B=1, T=4]
Input IDs [input_ids]: [245, 102, 56, 89] (x)
                              \ \ \ \
Target IDs [target_ids]: [102, 56, 89, 412] (y)
------------------------------------------------------------
Shift Logic: target_ids = input_ids[1:] + [next_token]
This enables the model to predict the next token at every position.
Enter fullscreen mode Exit fullscreen mode

Vector Embeddings

Tokens are mapped to a high-dimensional space where semantic relationships are captured. We combine Token Embeddings with Positional Embeddings to preserve sequence order.

Part II: The Attention Engine The Heart of the Transformer

The Attention mechanism enables the model to resolve long-range dependencies by assigning dynamic weights to different parts of the input sequence, overcoming the “forgetting” issues of RNNs.

The Mechanism of Scaled Dot-Product Attention

Attention relies on three linear projections: Query (Q), Key (K), and Value (V).

  1. Calculate the scores: Compute the dot product QK^\top to determine token relevance.
  2. Apply the Scaling Factor: Divide scores by \sqrt{d_k} (where d_k is the head dimension). This prevents the softmax function from entering regions with vanishingly small gradients as the dimensionality grows.
  3. Apply the Softmax: Convert scaled scores into attention weights A \in \mathbb{R}^{B \times H \times T \times T} where weights sum to 1.
  4. Compute the output: Multiply A by V to produce the final context-aware representation.

Causal Masking & Multi-Head Scaling

To ensure the model only predicts based on the past (autoregression), we apply a causal mask to the attention score matrix before the softmax step.

ASCII Causal Mask Matrix (T=4)

[0, -inf, -inf, -inf] (Token 1 sees Token 1)
[0, 0, -inf, -inf] (Token 2 sees 1, 2)
[0, 0, 0, -inf] (Token 3 sees 1, 2, 3)
[0, 0, 0, 0] (Token 4 sees 1, 2, 3, 4)
Enter fullscreen mode Exit fullscreen mode

The “So What?” Layer

While Single-Head Attention computes a single focus, Multi-Head Attention (MHA) allows the model to attend to multiple semantic subspaces in parallel (e.g., grammatical structure in Head 1, factual context in Head 2). However, MHA has high memory overhead during inference, leading to the adoption of Grouped-Query Attention (GQA) in modern models like Llama 3 to reduce KV cache size.

Part III: The Transformer Block and GPT Network Assembly

The GPT architecture is a stack of modular Transformer Blocks. Each block refines the token representations through self-attention and non-linear transformations.

Architectural Components

Modern implementations use the Pre-LayerNorm architecture, which places normalization before the attention and feed-forward layers to improve training stability in deep networks.

# GPTBlock Implementation with Pre-LayerNorm Residual Connections
class GPTBlock(nn.Module):
    def __init__ (self, cfg):
        super(). __init__ ()
        self.ln1 = LayerNorm(cfg["emb_dim"])
        self.attn = MultiHeadAttention(cfg)
        self.ln2 = LayerNorm(cfg["emb_dim"])
        self.ffn = FeedForward(cfg)
        self.drop = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
        # x is [Batch, Seq_Len, Emb_Dim]
        x = x + self.attn(self.ln1(x)) # Residual connection 1
        x = x + self.ffn(self.ln2(x)) # Residual connection 2
        return x
Enter fullscreen mode Exit fullscreen mode

Network Synthesis

The GPTModel consolidates these blocks into a unified pipeline:

  1. Input Embedding Layer: Combines Token + Positional embeddings (B, T, D).
  2. Dropout: Standard regularization to prevent overfitting.
  3. Transformer Blocks: A sequence of N stackable modules.
  4. Final Layer Norm: Standardizing the output of the final block.
  5. Linear Head: Projects the latent space back to [B, T, Vocab_Size] for probability distribution calculation.

Part IV: Autoregressive Pretraining and Decoding Strategies

Pretraining transforms a randomly initialized architecture into a statistical model of language via self-supervised learning on massive datasets.

Loss and Optimization

The model is trained using Cross-Entropy Loss , minimizing the negative log-likelihood of the correct next token. Optimization involves balancing the learning rate and weight decay to ensure the weights converge without exploding gradients.

Weight Loading and Transfer Learning

We can load pretrained weights from OpenAI’s GPT-2 or Llama variants. A Senior Staff approach utilizes memory-efficient state dict loading (Source: Ch 5, bonus 8), which maps weights directly to the model architecture without doubling the memory footprint during the transfer process.

Inference and Decoding

At inference time, we use different strategies to control token generation:

  • Temperature Sampling: Scales the logits before softmax. T < 1.0 makes the distribution “sharper” (more logical), while T > 1.0 makes it “flatter” (more creative).
  • Top-k Sampling: Filters the top k candidates, ensuring the model does not sample from the “long tail” of low-probability, nonsensical tokens.

Part V: Downstream Fine-Tuning and Preference Alignment

Base models generate general text but require alignment for specific utility. This is the stage where the model becomes a specialized tool.

Classification and Instruction Tuning

By replacing the final linear head with a task-specific head, we can transform the GPT architecture into a classifier (e.g., Spam Detection). For conversational agents, Instruction Finetuning is followed by Direct Preference Optimization (DPO), which directly optimizes the model to prefer “chosen” over “rejected” responses based on human feedback data.

Efficiency with LoRA

Full-parameter fine-tuning is often impractical. Low-Rank Adaptation (LoRA) freezes the original weights W_0 and adds a pair of low-rank matrices A and B, such that the update \Delta W = B \times A. Top 3 Advantages of LoRA:

  1. Compute Efficiency: Only a fraction (e.g., 1%) of parameters are updated.
  2. Storage Scalability: Task-specific adapters are only a few megabytes.
  3. VRAM Conservation: Enables fine-tuning of 7B+ models on consumer GPUs.

Part VI: Modern SOTA Extensions and Reasoning Architectures

The field is rapidly moving toward more efficient and reasoning-capable architectures.

Inference and Architectural Innovations

Modern models like Olmo 3 and Tiny Aya (Source: Ch 5, bonus 13/15) implement several SOTA enhancements:

  • KV Caching: Storing previous K and V tensors to avoid O(T²) recomputation during generation.
  • Multi-Head Latent Attention (MLA): Compressed KV projections to further optimize cache memory.
  • Sliding Window Attention (SWA): Limiting attention to a fixed local window to handle extremely long contexts (Source: Ch 4 bonus).
  • Mixture-of-Experts (MoE): Models like Qwen 3.5 MoE activate only specific “experts” (sub-networks) for each token, allowing for high parameter counts with low active compute costs.

The Reasoning Frontier

The latest evolution involves Reasoning Models and Reinforcement Learning (GRPO). These models utilize inference-time scaling allowing the model to “think” longer via chain-of-thought and verifier-based evaluation to solve complex logical and mathematical problems. This marks the transition from purely predictive text to verifiable reasoning.

Need High-Impact Technical Content for Your Engineering Team?

I partner with developer-tooling startups, SaaS platforms, and engineering teams to translate complex infrastructure, agentic systems, and backend architecture into publication-grade technical writing.

Whether you need deep-dive architecture essays, hands-on developer tutorials, or technical counter-narratives:

Top comments (0)