The character-level model built in the walkthrough spends its capacity learning to spell: every word must be assembled one letter at a time, and a context of 256 characters holds only a few lines of verse. Large language models instead read subword tokens, frequent words and pieces of words, so that common words cost one token and rare ones a few. A later version of the program in the same repository, train_gpt_tokenized.py, swaps the character vocabulary for a byte-level byte-pair-encoding (BPE) tokenizer trained on the same text and keeps the Transformer unchanged. This article describes what changes and what the change implies. The walkthrough itself does not cover this version, so everything here is read from the code; no model was trained or tokenizer run for this article.
The code is in the same folder as the character-level program, gpt_text_model. The code blocks are excerpts of the unified diff from train_gpt.py to train_gpt_tokenized.py, generated from the files.
Byte-pair encoding in brief
Byte-pair encoding builds a vocabulary by merging. It starts from single symbols, counts which adjacent pair occurs most often in the training text, adds that pair to the vocabulary as a new symbol, and repeats until the vocabulary reaches the requested size. Frequent words end up as single tokens; rare words remain split into several. The method, its origin as a compression algorithm, and its adaptation to neural machine translation by Sennrich, Haddow and Birch (2016) are explained in the subword section of Tokenization.
GPT-2 made the method byte-level: the starting symbols are the 256 possible byte values rather than Unicode characters, so that any string at all can be encoded without an "unknown" token, from a base vocabulary of only 256 (Radford et al., 2019, Section 2.2). GPT-2 also prevented merges across character categories, so that variants such as "dog." and "dog!" do not each consume a vocabulary slot, with an exception for spaces; its vocabulary has 50,257 tokens (Sections 2.2 and 2.3). In GPT-2's tokenizer a space is attached to the beginning of the word that follows it: the pattern that splits text before merging includes ?\p{L}+, an optional space followed by letters, in its released encoder.py. The tokenizer used here is the ByteLevelBPETokenizer of the Hugging Face tokenizers library, an implementation of the same byte-level scheme.
Training the tokenizer on the text
The character vocabulary and the stoi/itos dictionaries are replaced by a tokenizer that is trained on the dataset the first time the program runs and then saved, as vocab.json and merges.txt, so later runs reload it:
--- train_gpt.py
+++ train_gpt_tokenized.py
@@ -8,10 +8,16 @@
+# Import the ByteLevelBPETokenizer from Hugging Face's tokenizers library
+try:
+ from tokenizers import ByteLevelBPETokenizer
+except ImportError:
+ raise ImportError("Please install the tokenizers library with: pip install tokenizers")
+
# Updated Hyperparameters
-batch_size = 64 # Increased batch size
-block_size = 256 # Increased block size for more context
-C = 384 # Increased embedding dimension
-learning_rate = 3e-4 # Lowered learning rate for stable training
-n_layer = 6 # Increased number of layers
-num_heads = 6 # Increased number of attention heads
-dropout_rate = 0.2 # Dropout rate for regularization
+batch_size = 64 # Increased batch size
+block_size = 256 # Increased block size for more context (block_size now counts tokens)
+C = 384 # Increased embedding dimension
+learning_rate = 3e-4 # Lowered learning rate for stable training
+n_layer = 6 # Increased number of layers
+num_heads = 6 # Increased number of attention heads
+dropout_rate = 0.2 # Dropout rate for regularization
eval_interval = 100
@@ -69,22 +67,39 @@
####################################################################################################
-# Load the Dataset and Create Vocabulary
-
-with open(data_file, 'r') as f:
+# GPT-2 Style BPE Tokenization Setup
+
+# Define a directory to store (or load) the BPE tokenizer files
+bpe_dir = os.path.join(data_dir, "bpe")
+vocab_file = os.path.join(bpe_dir, "vocab.json")
+merges_file = os.path.join(bpe_dir, "merges.txt")
+
+# If the BPE model files do not exist, train a new tokenizer on the dataset.
+if not os.path.exists(vocab_file) or not os.path.exists(merges_file):
+ os.makedirs(bpe_dir, exist_ok=True)
+ print("Training new BPE tokenizer...")
+ tokenizer = ByteLevelBPETokenizer()
+ # You can adjust vocab_size as needed; GPT-2 used 50,257 tokens by default.
+ tokenizer.train(files=[data_file], vocab_size=50257, min_frequency=2, special_tokens=["<|pad|>", "<|unk|>", "<|bos|>", "<|eos|>"])
+ tokenizer.save_model(bpe_dir)
+ print(f"Tokenizer saved to {bpe_dir}.")
+else:
+ print("Loading existing BPE tokenizer...")
+ tokenizer = ByteLevelBPETokenizer(vocab_file, merges_file)
+
+# Get the vocabulary size from the tokenizer
+vocab_size = tokenizer.get_vocab_size()
+
+# Define encode and decode functions that work with tokens rather than characters
+def encode(s):
+ # Returns a list of token IDs
+ return tokenizer.encode(s).ids
+
+def decode(token_ids):
+ # Converts a list of token IDs back into a string
+ return tokenizer.decode(token_ids)
+
+# Read the dataset and encode it into token IDs
+with open(data_file, 'r', encoding='utf-8') as f:
text = f.read()
-chars = sorted(list(set(text)))
-vocab_size = len(chars)
-
-####################################################################################################
-# Encode and Decode Text
-
-stoi = {ch: i for i, ch in enumerate(chars)}
-itos = {i: ch for i, ch in enumerate(chars)}
-
-def encode(s):
- return [stoi[c] for c in s]
-
-def decode(l):
- return ''.join([itos[i] for i in l])
-
+# Encode the entire text; note that this produces a long list of token IDs.
data = torch.tensor(encode(text), dtype=torch.long)
The training call asks for vocab_size=50257, GPT-2's size, with min_frequency=2, so that a pair must occur at least twice to be merged, and four special tokens, <|pad|>, <|unk|>, <|bos|> and <|eos|>. The requested size is an upper limit rather than a guarantee: merges stop when no remaining pair is frequent enough, and on a text of about a million characters that can happen before 50,257 is reached. The program does not print the resulting vocabulary size, but vocab_size = tokenizer.get_vocab_size() passes whatever it is to the model.
encode and decode keep their names and roles, so the rest of the program is unaffected: the whole text is still encoded into one long tensor, now of token IDs, and split 90/10 into training and validation by position, now counted in tokens. The download also gains encoding='utf-8', which matters on Windows for any text that is not pure ASCII.
What the vocabulary size does to the model
The Transformer is untouched; only vocab_size changes. But two layers scale with it: the token embedding table ($V \times C$ weights) and the LM head ($C \times V$ weights plus $V$ biases). With $C = 384$, every other layer together has 10,738,944 parameters, so the model has
$$769\,V + 10{,}738{,}944$$
parameters, where 769 = 384 + 384 + 1. With the 65-character vocabulary this gives the familiar 10,788,929. The counts for some other vocabulary sizes, computed with tools/experiments/gpt/param_counts.py in this site's repository:
| Vocabulary size $V$ | Parameters | Share in embedding and LM head |
|---|---|---|
| 65 (characters) | 10,788,929 | 0.5% |
| 5,000 | 14,583,944 | 26% |
| 20,000 | 26,118,944 | 59% |
| 50,257 (if reached) | 49,386,577 | 78% |
At GPT-2's full vocabulary size, more than three quarters of the parameters would sit in the two vocabulary-sized layers, and the model would have about 49 million parameters to fit on about a megabyte of text. Each token's embedding row is trained only when that token appears, so rows for tokens that occur a handful of times receive few updates. The program prints the actual parameter count when training starts, which reveals the vocabulary size indirectly. This is one reason models trained on small corpora use smaller vocabularies, and why vocab_size is the first setting worth reducing in experiments with this program.
What the context length now means
block_size is still 256, but it now counts tokens. Since a typical English word is one or a few BPE tokens, 256 tokens cover several times as much text as 256 characters, so each prediction can draw on a longer stretch of the play. The cost of attention depends on the number of positions, not on how much text they represent, so this longer reach comes at no extra cost per step. The trade is that each prediction is harder: the model must choose among thousands of possible tokens instead of 65 characters.
Comparing losses across tokenizations
The training loop reports the same cross-entropy loss as before, but the loss is now per token, not per character, and the two are not comparable. A model that predicts whole words faces a harder choice at each step and has fewer steps per line, so a per-token loss is naturally higher than a per-character loss for the same quality of prediction. Vaswani et al. (2017) make the same caution about their own results: the perplexities in their Table 3 are "per-wordpiece, according to our byte-pair encoding, and should not be compared to per-word perplexities" (Table 3, caption).
The fair comparison converts both to the same unit. The total loss over a validation text, in nats, is the per-token loss times the number of tokens; dividing by the number of characters gives nats per character, comparable with the character-level model's 1.48 from the run described in Scaling the Shakespeare GPT:
$$\text{loss per character} = \text{loss per token} \times \frac{\text{number of tokens}}{\text{number of characters}}.$$
For example, if the tokenizer turned the validation text into a quarter as many tokens as it has characters, a per-token loss of 5.0 would correspond to 1.25 nats per character, better than the character model, while a per-token loss of 6.4 would correspond to 1.6, worse. (These figures illustrate the conversion; they are not results from this program.) The same issue arises with perplexity, discussed for topic models in Evaluating Topic Models: perplexity per token and perplexity per word or per character measure uncertainty over different units.
Prompts for generation
The sample prompts change to suit tokens:
--- train_gpt.py
+++ train_gpt_tokenized.py
@@ -320,7 +331,6 @@
scenarios = {
- "Start with 'A'": "A",
- "Start with 'The'": "The",
+ "Start with '<|bos|>'": "<|bos|> ",
+ "Start with 'The'": "The ",
"Start with 'To be, or not to be, that is the '": "To be, or not to be, that is the "
}
-
for description, start_text in scenarios.items():
Two details deserve care. The prompt <|bos|> uses one of the special tokens as a "beginning of sequence" marker, but the training text contains no such marker anywhere, so the model has never seen it in context and has no learned statistics for what follows it; it starts from a token sequence unlike anything in its training data. And the prompts The and To be, or not to be, that is the end with a space. In GPT-2's scheme the space belongs to the start of the next word's token, so a prompt that ends in a bare space ends with a token that, in the training text, almost always occurs fused with a following word; dropping the trailing space usually gives a more natural continuation. (As Scaling the Shakespeare GPT notes, Hamlet is not in tiny-shakespeare, so the third prompt is unfamiliar in any tokenization.)
Smaller fixes in this version
The tokenized program also cleans up several rough edges of the character-level one.
A clearer local name. Inside the attention head and compute_loss, B, T, C = x.shape reused the name of the global embedding size C. That was harmless, since the local value equals the global one in the head, but confusing; the new version calls it C_:
--- train_gpt.py
+++ train_gpt_tokenized.py
@@ -114,3 +129,3 @@
def forward(self, x):
- B, T, C = x.shape
+ B, T, C_ = x.shape
Loading a GPU checkpoint on a CPU. torch.load now maps the saved tensors to the current device, so a checkpoint written on a GPU machine loads on one without a GPU:
--- train_gpt.py
+++ train_gpt_tokenized.py
@@ -306,4 +316,5 @@
# Load the Model
+
def load_checkpoint(model, optimizer, checkpoint_path):
- checkpoint = torch.load(checkpoint_path)
+ checkpoint = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(checkpoint['model_state_dict'])
The first run without checkpoints. The character-level main called os.path.exists(None) when no checkpoint folder existed, which raises a TypeError; the new version checks for None first:
--- train_gpt.py
+++ train_gpt_tokenized.py
@@ -418,11 +432,9 @@
print(generated_text)
- pass
def main():
- # Attempt to load from a checkpoint
+ # Attempt to load from the latest checkpoint if one exists
start_epoch = 0
checkpoint_path = get_path_to_latest_checkpoint()
- if os.path.exists(checkpoint_path):
+ if checkpoint_path and os.path.exists(checkpoint_path):
start_epoch = load_checkpoint(model, optimizer, checkpoint_path)
-
run()
Whole prompts. The generate command now passes everything after the command word as the prompt, not just the first word:
--- train_gpt.py
+++ train_gpt_tokenized.py
@@ -397,3 +412,4 @@
elif command.startswith("generate"):
- prompt = command.split()[1]
+ # Everything after the command is taken as the prompt.
+ prompt = command[len("generate "):]
generate(prompt)
Clearer progress reports. Training prints the parameter count and configuration at the start, and formats each evaluation as a progress report. The final checkpoint after the last step, which the character-level version saved, is no longer written; the last saved checkpoint is the one at step 9,900:
--- train_gpt.py
+++ train_gpt_tokenized.py
@@ -338,4 +347,12 @@
-# Update the training loop to call the generate_text_scenarios function at each checkpoint
def train_model(start_step=0):
+ print(f"\n[+] Starting training from step {start_step}")
+ print(f"[+] Model parameters: {sum(p.numel() for p in model.parameters()):,}")
+ print(f"[+] Training configuration:")
+ print(f" - Batch size: {batch_size}")
+ print(f" - Learning rate: {learning_rate}")
+ print(f" - Number of layers: {n_layer}")
+ print(f" - Number of heads: {num_heads}")
+ print(f" - Embedding dimension: {C}")
+
for step in range(start_step, num_train_steps):
@@ -350,13 +367,13 @@
losses = estimate_loss()
- print(f"Step {step}, Train Loss: {losses['train']}, Val Loss: {losses['val']}")
- # Save checkpoint at each evaluation interval
+ print(f"\n[+] Step {step}/{num_train_steps} ({step/num_train_steps*100:.1f}%)")
+ print(f" - Train Loss: {losses['train']:.4f}")
+ print(f" - Val Loss: {losses['val']:.4f}")
+ print(f" - Learning rate: {optimizer.param_groups[0]['lr']:.2e}")
+
+ # Save checkpoint and generate sample
save_checkpoint(model, optimizer, step)
- # Generate text after saving the checkpoint
+ print("\n[+] Generating sample text...")
generate_text()
- # Save the final model after training
- save_checkpoint(model, optimizer, num_train_steps)
-
def get_path_to_latest_checkpoint():
- """Return the path to the latest checkpoint file"""
checkpoint_dir = 'checkpoints'
One caution for anyone running both programs: they save to the same checkpoints folder with the same file names, and the tokenized program loads the latest checkpoint it finds at start-up. A character-level checkpoint has embedding and LM-head matrices of the wrong shape for the BPE model, so loading it fails. Running each version in its own working directory, or giving each its own checkpoint folder, avoids this.
What the switch does and does not change
Switching from characters to subword tokens changes what the model has to learn, not how it learns. The Transformer, the training loop, the loss and the sampling are the same code; only the mapping between text and integers is different. The model no longer has to learn spelling, and a context of 256 tokens reaches much further back. In exchange, it must learn a far larger output distribution from the same small text, with most of its parameters in vocabulary-sized tables. Whether that trade pays off on tiny-shakespeare is an empirical question that the per-character comparison above is designed to answer.
References
- Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. (2019). Language models are unsupervised multitask learners. OpenAI technical report. Cited: Section 2.2 (byte-level BPE, base vocabulary of 256, no merges across character categories), Section 2.3 (vocabulary of 50,257). Tokenizer code: https://github.com/openai/gpt-2/blob/master/src/encoder.py
- Sennrich, R., Haddow, B., and Birch, A. (2016). Neural machine translation of rare words with subword units. In Proceedings of ACL 2016, 1715–1725.
- Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017). Attention is all you need. In Advances in Neural Information Processing Systems 30. arXiv:1706.03762. Cited: Table 3 caption (per-wordpiece perplexities).
- Code for this article: train_gpt_tokenized.py and train_gpt.py in github.com/mdb42/ml-demos.