# How to Configure Different LLM Sizes Using Hyperparameters: n_embed, n_head, and n_blocks

> Learn to configure LLM sizes using n_embed, n_head, and n_blocks hyperparameters. Adjust model dimensions and layers via Python, JSON, or command line in train LLM from scratch.

- Repository: [Fareed Khan/train-llm-from-scratch](https://github.com/FareedKhan-dev/train-llm-from-scratch)
- Tags: how-to-guide
- Published: 2026-06-11

---

**You can configure LLM model sizes in the train-llm-from-scratch repository by adjusting three core hyperparameters—`n_embed` (embedding dimension), `n_head` (attention heads), and `n_blocks` (transformer layers)—through Python dictionaries, JSON configuration files, or command-line overrides.**

The FareedKhan-dev/train-llm-from-scratch repository separates pre-training and post-training pipelines, but both stages expose identical architectural size knobs. Understanding how to manipulate `n_embed`, `n_head`, and `n_blocks` allows you to scale models from tiny CPU-testable configurations to large production-scale transformers without modifying the underlying architecture code.

## Understanding the Core Hyperparameters

Three values control the transformer architecture's capacity and computational requirements. Each parameter directly impacts the model's parameter count and memory footprint.

### n_embed (Embedding Dimension)

**`n_embed`** defines the hidden size of the model, controlling the width of token embeddings and the dimensionality of each transformer's feed-forward network. In [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py), this is defined as `N_EMBED` (default `2048` for pre-training), while [`config/post_training_config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/post_training_config.py) uses `n_embed` (default `1024` for post-training stages). Larger values increase model capacity but require significantly more memory and compute.

### n_head (Number of Attention Heads)

**`n_head`** specifies the number of parallel attention mechanisms operating in each transformer block. The repository defaults to `16` heads across both configurations. A critical constraint is that `n_embed` must be divisible by `n_head` to ensure even splitting of the embedding dimension across heads (e.g., 1024 dimensions with 16 heads yields 64 dimensions per head).

### n_blocks (Number of Transformer Blocks)

**`n_blocks`** determines the model's depth by controlling how many transformer layers are stacked. Pre-training defaults to `24` blocks in [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py), while post-training configurations inherit this value through `BaseModelConfig.n_blocks`. Increasing this value improves the model's ability to capture long-range dependencies but linearly increases training time.

## Where These Hyperparameters Live in the Codebase

The repository maintains separate configuration systems for pre-training and post-training stages, though both expose the same three size parameters.

### Pre-training Configuration (config/config.py)

For pre-training, hyperparameters reside in a Python dictionary within [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py). The [`scripts/train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py) file instantiates the model by passing these values directly to the `Transformer` class:

```python

# From scripts/train_transformer.py

model = Transformer(
    n_head=config['n_head'],
    n_embed=config['n_embed'],
    context_length=config['context_length'],
    vocab_size=config['vocab_size'],
    N_BLOCKS=config['n_blocks']
)

```

The `default_config` dictionary in [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py) contains the default assignments for `N_EMBED`, `N_HEAD`, and `N_BLOCKS`.

### Post-training Configuration (post_training_config.py)

Post-training stages (SFT, DPO, PPO, GRPO) use a dataclass-based configuration system defined in [`config/post_training_config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/post_training_config.py). The `BaseModelConfig` class holds the architectural parameters:

```python
@dataclass
class BaseModelConfig:
    n_embed: int = 1024
    n_head: int = 16
    n_blocks: int = 24
    # ... other fields

```

Configuration resolution follows a hierarchy: dataclass defaults merge with [`configs/base.json`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/configs/base.json), then stage-specific JSON files (e.g., [`configs/sft.json`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/configs/sft.json)), and finally CLI overrides.

## How to Modify LLM Sizes

You can adjust model sizes through four distinct methods depending on your workflow stage.

### Editing Pre-training Parameters

To change the model size for pre-training, modify the constants in [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py):

```python

# config/config.py

N_EMBED = 1536      # Must be divisible by N_HEAD

N_HEAD = 12         # Number of attention heads

N_BLOCKS = 18       # Number of transformer layers

```

These changes take effect immediately when running [`scripts/train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py).

### Configuring Post-training via JSON

For post-training stages, create or edit JSON files in the `configs/` directory. The `loader.load_config` function merges these with base settings:

```json
{
  "n_embed": 1536,
  "n_head": 12,
  "n_blocks": 18,
  "lr": 2e-5,
  "batch_size": 32
}

```

Save this as [`configs/sft.json`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/configs/sft.json) or similar, and the training script will resolve these values through the configuration inheritance chain.

### Using CLI Overrides

When launching post-training scripts, override any architectural parameter using the `--field` flag:

```bash
python -m src.post_training.cli train_sft \
    --config configs/sft.json \
    --field n_embed=2048 \
    --field n_head=16 \
    --field n_blocks=24

```

The [`src/post_training/cli.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/cli.py) entry point passes these overrides to `loader.load_config`, which gives precedence to command-line values over JSON configurations.

### Programmatic Configuration with the Smoke Helper

For rapid experimentation or CPU testing, use the `smoke` helper function to generate minimal configurations:

```python
from config.post_training_config import SFTConfig, smoke

# Creates a tiny model: n_embed=64, n_head=4, n_blocks=2

tiny_cfg = smoke(SFTConfig)
print(tiny_cfg.n_embed)  # Output: 64

```

This approach is ideal for debugging pipeline logic without loading multi-billion-parameter models.

## Parameter Count Estimation

The total parameter count scales approximately according to the formula:

```

params ≈ n_blocks * (4 * n_embed^2 + 2 * n_head * (n_embed / n_head)^2) + vocab * n_embed

```

This quadratic relationship with `n_embed` and linear relationship with `n_blocks` explains why increasing embedding dimensions yields more powerful models but requires substantially more GPU memory. Always ensure `n_embed` is divisible by `n_head` to avoid runtime errors during tensor reshaping operations.

## Summary

- **Three core knobs** control LLM size: `n_embed` (width), `n_head` (parallel attention), and `n_blocks` (depth).
- **Pre-training** uses [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py) with dictionary keys `N_EMBED`, `N_HEAD`, and `N_BLOCKS` consumed by [`scripts/train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py).
- **Post-training** uses dataclass-based configs in [`config/post_training_config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/post_training_config.py), resolved via JSON files and CLI overrides through `loader.load_config`.
- **CLI overrides** allow rapid experimentation using `--field parameter=value` syntax when launching training scripts.
- **Validation rule**: `n_embed` must be divisible by `n_head` to ensure proper head dimension allocation.

## Frequently Asked Questions

### What happens if n_embed is not divisible by n_head?

The model will raise a runtime error during the attention mechanism's forward pass. The `train-llm-from-scratch` implementation assumes even division to split the embedding dimension into `n_head` separate attention heads. For example, with `n_embed=1024`, valid `n_head` values include 2, 4, 8, 16, 32, or 64.

### Can I change model size between pre-training and post-training stages?

Yes, but you must ensure the checkpoint conversion handles dimension mismatches. Typically, you would pre-train with one configuration (e.g., `n_embed=2048`) and then adjust the post-training configuration JSON to match the saved checkpoint dimensions, or implement interpolation/trimming logic if switching sizes between stages.

### How do I quickly test my pipeline with a small model?

Import the `smoke` function from `config.post_training_config` and pass your stage configuration class (e.g., `PPOConfig`, `SFTConfig`). This generates a tiny configuration with `n_embed=64`, `n_head=4`, and `n_blocks=2`, allowing you to verify data loading and training logic on CPU without loading large weights.

### Where are the default values for post-training defined?

Default architectural values for post-training reside in `BaseModelConfig` within [`config/post_training_config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/post_training_config.py). These are then potentially overridden by [`configs/base.json`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/configs/base.json) and stage-specific files like [`configs/sft.json`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/configs/sft.json). The `loader.load_config` function manages this merge hierarchy, with CLI arguments taking final precedence.