Training Llama 3 from scratch Aman Gupta 26 min read Jan 10, 2025 If you’ve been interested in training LLMs, sooner or later, you would want to understand how to train a model from random weights. That’s the curiosity that made me excited to do this exercise after having spent last 2 years focusing on post-training. In this article, we will start from scratch and write all the code needed to pre-train a small Llama architecture model. It’s very useful to have a target model about which some details are known — in this case, that target is the SmolLM series of models by Huggingface. Specifically, we will target the 135M base model from that series. We will also borrow some information from the Llama 2 technical report and a plethora of other available information. What makes pre-training different from post-training? Essentially, dataset size. Typically, post-training datasets are in the millions or 10s of millions of tokens in size. Compare that to 600B tokens that SmolLM-135M was trained on and 9T tokens that Llama-3.2 –1B was trained on! That’s six orders of magnitude larger! If it takes an hour to run the post-training (SFT) run on 9M tokens, then it would take 114 years(!) to pre-train 9T tokens on the same hardware/software stack. The other practical difference is that parameter-efficient methods, like LoRA, are quite popular because they help achieve 90% (figuratively) of the gains while effectively reducing the number of trainable parameters down to less than 1% of the model size. It helps make the training run on a hardware stack with lower memory and run faster. For us, this means that we have to find a balance between model size, dataset size, and the hardware available to us (which is 8 A100s). Secondly, this means that every ounce of performance gain matters. Pre-requisites There exist quite amazing tutorials on how to pre-train models from scratch with very similar constraints, and they’ve done an excellent job of explaining the concepts. For this exercise, Andrej Karpathy’s video on reproducing GPT-2 was a significant influence. To follow the rest of this article, I believe it’s essential to have seen and understood two of his videos: Let’s build GPT and Let’s reproduce GPT-2 (124M). However, if you understand the transformer architecture in general, you’ll still find value in just carrying on. LLama architecture While Andrej’s videos focus on reproducing one of the most well-known architectures — GPT-2, that architecture is now 5 years old (in 2024). There are some fundamental updates to the architecture that recent open models contain, and the Llama series of models is the flagship of the open-source armada. The SmolLM models adopt the Llama architecture, and so do we. What are the updates? Here are the main ones: 1. Grouped Query Attention (GQA)— an implementation of Self-attention that promises to be more efficient than regular Multi-Head Attention. 2. Rotary Position Embeddings (RoPE) — an updated mechanism to convey positional information to the attention mechanism that doesn’t require positional embeddings and allows you to keep the maximum sequence length of the model theoretically unbounded. 3. SwiGLU non-linearity in the MLP/FFN layer Let’s do a quick review of these, and I will provide resources to dive deeper if you want to. One thing to note is that it’s unlikely that these improvements will cause a significant jump in the model quality after training on the same data. In fact, GQA and RoPE both have other practical purposes — GQA speeds up inference, and RoPE allows increasing context length windows of the models. Grouped Query Attention Multi Head Attention (MHA) is a severe bottleneck for model inference. Fundamentally, the attention computation shown below is O(n^2) in the number of operations, where n is the input sequence length. To optimize inference, inference engines like VLLM use a KV cache to avoid recomputations, which brings the time complexity down to O(n). But, this shifts the bottleneck to memory bandwidth — with large sequence lengths, the cache size bloats, and a significant amount of time is spent loading the cached computations from memory. This excellent post (and its sequel) explains the phenomenon in detail. Reducing the number of KV tensors participating in the computation can significantly reduce that overhead, and Grouped Query Attention (GQA) achieves exactly that. Essentially, it has fewer Key and Value heads and duplicates them to match the number of Query heads. In code, it means that the Linear layers tensors for k_proj and v_proj are smaller: # Example config: # n_attn_heads = 9 # n_kv_heads = 3 # d_model = 576 # d_head = 64 self. q_proj = nn. Linear (config. d_model , config. n_attn_heads * config. d_head , bias=False) self. k_proj = nn. Linear (config. d_model , config. n_kv_heads * config. d_head , bias=False) self. v_proj = nn. Linear (config. d_model , config. n_kv_heads * config. d_head , bias=False) self. o_proj = nn. Linear (config. d_model , config. d_model , bias=False) You can see that there each Key and Value head serves 3 Query attention heads (n_attn_heads/n_kv_heads == 3). During the pass forward, we copy the keys and values using torch. repeat_interleave : # GQA # k. shape == v. shape == (batch_size, n_kv_heads, seq_len, d_head) # for k, v, match size of dim=-3 to be equal to n_attn_heads (up from n_kv_heads) k = k. repeat_interleave (self. config. n_attn_heads / self. config. n_kv_heads, -3) v = v. repeat_interleave (self. config. n_attn_heads / self. config. n_kv_heads, -3) Note that the torch’s fused kernel for attention (FlashAttention 2) supports GQA, so we can use that when it’s available. Rotary Position Embeddings Rotary Position Embeddings, or RoPE, is one of the most fundamental architectural improvements that essentially has very little downside (small performance loss) and a lot of upsides. Remember that the original transformer architecture requires you to add position embeddings to the input embeddings to encode positional information before entering the residual stream so that the attention mechanism can know the distance between two tokens in the residual stream. In GPT-2, these embeddings are learned during the optimization (from Andrej Karpathy’s nanoGPT repo): # in __init__() wpe = nn. Embedding (config. block_size , config. n_embd ) # in forward() tok_emb = self. transformer. wte(idx) # token embeddings of shape (b, t, n_embd) pos_emb = self. transformer. wpe(pos) # position embeddings of shape (t, n_embd) RoPE essentially gets away from using positional embeddings, and instead, it encodes position into the query and key embeddings before computing the attention weights: # q and k don't have positional information encoded in them here. q = self. q_proj (x) k = self. k_proj (x) #... # RoPE: encode positional information q, k = self._apply_rotary_pos_emb (q, k, cos, sin) # compute attention weight matrix out = F. scaled_dot_product_attention (q, k, v, is_causal=True, enable_gqa=True) Because we don’t have a fixed-size tensor for positional embeddings, the maximum sequence length supported by a model is determined by the sequence length at train time. Moreover, you can adapt existing models to a longer sequence length with fine-tuning (this is how Llama 3.1 models increased their sequence lengths to 128K tokens, from 8K of Llama 3.0 models). We will not go into details of how RoPE actually works in this article because here is an article that does a fantastic job of explaining it. It’s worth taking a pause right now and reading that. Our implementation is effectively taken from Eleuther’s blog here, which has another, more mathematically driven explanation and also includes some experimental data on loss and performance metrics. SwiGLU The original Transformers paper suggested the ReLU non-linearity for the MLP (feed-forward) network. Since then, people have experimented with other activation functions, and SwiGLU, introduced here, is the current choice for most architectures (in late 2024/early 2025). As before, there are excellent explanations of SwiGLU — and here is a blog post about it, along with an interactive graph showing possible shapes SwiGLU can take. To understand the implementation, let’s look at the formula (omitting the bias terms): Thus, instead of applying ReLU on the linear transformation xW , we have two linear transformations — xW and xV , and we then apply swish on the first one, implemented in pytorch as torch. nn. SiLU. This is how it manifests in code: class GatedMlp(nn. Module ): def __init__(self, config): super(GatedMlp, self).__init__() self. up_proj = nn. Linear (config. d_model , config. d_mlp_proj , bias=False) self. gate_proj = nn. Linear (config. d_model , config. d_mlp_proj , bias=False) self. down_proj = nn. Linear (config. d_mlp_proj , config. d_model , bias=False) self. silu = nn. SiLU () def forward(self, x): up = self. silu (self. gate_proj (x)) * self. up_proj (x) return self. down_proj (up) The parallel path with the gate_proj layer and SiLU non-linearity acts as a more learnable form of activation rather than a fixed function like ReLU or GELU. Empirically, SwiGLU has been known to perform well in helping models converge to a lower loss (with the author famously attributing the improvement to “divine benevolence”). More recently, a talk called The Physics of Language Modeling studied the impact of MLP design on memorization (as an aside, that video is a must-watch for anyone in the field). Their experiments show that SwiGLU reduces the model’s ability to memorize rarely occurring knowledge in the dataset. Now, let’s get into the implementation! First, a note about abstractions If you have written software, you have been in discussions about abstractions. Should we use a library for training, or should we write our own training loops? What about running evals? Use pure PyTorch or use libraries like deepspeed? There is an attraction towads removing abstractions because then you have control over the implementation. This reminds me of this quote by Carl Sagan: If you wish to make an apple pie from scratch, you must first invent the universe. There is no such thing as a world without abstractions. We experience the world through an abstraction — we all depend on our brain interpreting the signals from our sensory organs before informing us what we have sensed. We never actually experience reality as it is. In computer programming, that has also been very true since none of us write code in 0s and 1s. The choice of abstraction level is a decision that’s really based on the needs of the individual situation where this decision is being made. In my case, I wanted control over: 1. Model Architecture: There are many possible experiments that we can do by adapting the architecture. We could add multiple LM heads or implement a different attention mechanism. That’s one of the major points of doing this exercise in the first place, to try new or even whacky stuff and see what kind of model we get. 2. Training Loop: This is another area of exploration. We could try some variation of the loss function (for example, I’m curious about multi-token prediction loss vs. only next-token prediction loss). Or different optimizer and LR scheduler implementations. Things I wanted to re-use existing implementations for: 1. Model weight storage and distribution: Huggingface is the obvious choice, and it makes life so much easier if your model is compatible with HF implementation. Uou can use tons of utilities that work with their infra (like our next point, the evaluation harness). 2. Evaluation harness: EleutherAI harness has a mature library that runs most of the well-known evals. You can just provide the huggingface model ID, the benchmarks you want computed, and it computes the scores. 3. Tokenization and data loading: When working with large datasets, tokenization, and data loading can be a bit of work. I use Huggingface’s tokenizers implementation for the tokenizers themselves and the Datatrove library to run parallel fault-tolerant jobs for large-scale tokenization, which also provides a neat utility to load the token data and provide it to the model. Code The GitHub repo smol-llama is here. Here are its contents: • model. py — an implementation of the Llama model. • train. py — the training loop and dataloaders. • hf_utils. py — tools to interact with HF, like loading weights into our model and pushing our model weights to HF • tokenize_fineweb. py and run_ddp_train. py — actual scripts to run the tokenization and training. • train_shakespeare. ipynb , other notebooks— code for smaller scale runs to test code. Let’s dive a little deeper. Implementing the model The model implementation doesn’t hold many surprises. One of the key parts is that the module names, like embed_tokens and lm_head are designed to match the name of these modules in Huggingface’s implementation. The code structure has a config class ModelConfig that contains all the model hyperparameters like embedding dimension and number of layers (similar to HF’s LlamaConfig, but a subset). We copy-paste the generation implementation from Karpathy’s nanoGPT implementation. It’s also somewhat incomplete in that it doesn't implement RoPE scaling (that enhances performance at longer context lengths). The SmolLM models from HF, which we are trying to imitate, don’t use RoPE scaling. It’s important to test any model implementation, so to do that, there is a method called load_from_pretrained in hf_utils. py to load weights of a compatible model from HF. Here is what we get when we try the model HuggingFaceTB/SmolLM2–135M hf_checkpoint = "HuggingFaceTB/SmolLM2-135M" #... input_ids = tokenizer(["Gravity is", "Dark Matter is"], return_tensors="pt").to(device)['input_ids'] model = load_from_pretrained(hf_checkpoint).to(device) idx = model. generate (input_ids, temperature=0.25 , top_k=25, max_new_tokens=16) print(tokenizer. batch_decode (idx)) # Outputs # ['Gravity is a force that pulls things towards each other. The gravitational pull of the', # 'Dark Matter is a form of energy that is not made up of particles. It is a form'] Tokenization and data loader For tokenization, we use the tokenizer trained by Huggingface specifically for SmolLM series of models. In their article introducing the models, they only say, “We used a tokenizer trained on the Smollm Corpus with a vocab size of 49152.”. Tokenizing a large dataset with 10 or 100BT tokens takes a while, and you want to parallelize that work and make it fault-tolerant. Huggingface’s Datatrove library is designed to run multi-step workflows in parallel, and it’s very easy to restart from failures. Here is the pipeline that I wrote to run the tokenization: pipeline = [ HuggingFaceDatasetReader( dataset=hf_dataset_id, dataset_options={ "split": 'train', "name": name, }, text_key=text_column, ), DocumentTokenizer( output_folder=output_folder, tokenizer_name_or_path=tokenizer_id, eos_token=eos_token, batch_size=10000, max_tokens_per_file=int(1e8), shuffle=shuffle, seed=1998 ) ] The first step fetches the dataset, and the second step runs the tokenizer. In practice, it works reasonably well. I did face two problems: • It does not log progress %. I had to look at the timestamps of the files it is writing to make sure it’s still working. • When I enabled shuffle, it seemed to get stuck (the file timestamps didn’t update for a long time), and I resorted to interrupting the job and running again. The library needs some love, but it saved me a bunch of time to write this myself and the tokenization worked well. The tokenization output is stored in binary files where every 2 or 4 bytes represent the token integer (2 bytes if vocabulary size < 65535, 4 bytes otherwise). To read the token data, Datatrove also includes a utility to read those files (it actually implements them as a PyTorch Dataset, which you can use with Torch data loaders). We simply use the indexing that the Dataset object provides, as shown below: self. dataset = DatatroveFolderDataset( folder_path=config. tokens_folder , filename_pattern=os. path. join(config. tokens_folder , "*.ds"), seq_len=config. max_seq_len , token_size=(2 if tokenizer. vocab_size < 65535 else 4), recursive=False ) x, y = zip(*[(self. dataset [idx]['input_ids'][:-1], self. dataset [idx]['input_ids'][1:]) for idx in range(start, end)]) x_t, y_t = torch. stack (list(x)), torch. stack (list(y)) In train. py there are two DataLoader implementations: • SimpleDataLoader — it just takes the whole text as a string and provides a dataloader for training and validation batches. • FileDataLoader — it expects the path of tokenized files and uses Datatrove’s loading utility. For working with multiple GPUs, the FileDataLoader shards the data. It shards. It does so in a way that the sequences for each rank are contiguous segment in the ordered list of sequences in the whole dataset. self. total_train_seqs = math. ceil ((1-config. val_size ) * self. num_seqs ) shard_size = self. total_train_seqs // world_size self. train_start_idx = rank * shard_size self. train_end_idx = (rank+1) * shard_size self. train_seqs = self. train_end_idx - self. train_start_idx print(f"Shard range rank:{rank:<13} | ({self. train_start_idx },{self. train_end_idx })") # Current index self. train_index = self. train_start_idx Finally, we keep a separate shard for validation loss: self. val_seqs = self. num_seqs - self. total_train_seqs # Current index self. val_index = self. total_train_seqs Training Loop The training process implemented in train. py is implemented in two classes — TrainerConfig and Trainer. The config class takes train time hyperparams like the number of epochs and learning rate. The Trainer implements the initialization (like device selection, DDP, and torch. compile ), the training loop, the evaluation, and the checkpoint saving. Here is how one can run training (copy-pasted from train_shakespeare. ipynb ): tokenizer_id = "HuggingFaceTB/SmolLM2-135M" tokenizer = AutoTokenizer. from_pretrained (tokenizer_id) tokenizer. pad_token = tokenizer. eos_token # Same config as SmolLM2-135M, but no weight tying model_config = ModelConfig( vocab_size=tokenizer. vocab_size , d_model=576, d_head=64, d_mlp_proj=1536, n_layers=30, n_kv_heads=3, n_attn_heads=9, rms_norm_eps=1e-5, initializer_range=0.041666666666666664 , rope_theta=100000.0 , padding_idx=tokenizer. pad_token_id ) train_config = TrainerConfig( per_device_train_batch_size=32, max_seq_len=128, num_epochs=12, eval_interval_steps=25, learning_rate=1e-4, grad_clip_norm=1.0 , val_size=0.2 , log_dir="runs/shakespeare", warmup_ratio=0.1 ) with open("data/tiny_shakespeare. txt ") as f: text = f. read () model = LlamaModel(model_config) dataloader = SimpleDataLoader(train_config, tokenizer, text=text) trainer = Trainer(train_config, model) # Total tokens | 341,120 # Num Trainable Params | 162,826,560 # Train device | cuda, NVIDIA GeForce RTX 3090, N=1 # Training precision | torch. bfloat16 # Flash Attention | True # torch. compile () | True # DistributedDataParallel | False trainer. train (dataloader) # Training steps | 804 # Step: 0, Training Loss: 11.30751 , LR: 0.0000050 , Tokens/sec: 152.02 # Step: 1, Training Loss: 11.31026 , LR: 0.0000062 , Tokens/sec: 174.16 # Step: 2, Training Loss: 11.30679 , LR: 0.0000074 , Tokens/sec: 78161.30 # Step: 3, Training Loss: 11.21072 , LR: 0.0000086 , Tokens/sec: 88787.07 # Computing Eval loss, steps: 17 # Step: 3, Eval Loss: 11.20814 # Step: 4, Training Loss: 11.17848 , LR: 0.0000098 , Tokens/sec: 91534.29 The training loop implements a cosine LR scheduler with warmup: warmup_steps = math. floor (self. config. warmup_ratio * num_steps) warmup_factor = lambda st: 0.05 + 0.95 *(st / max(warmup_steps, 1)) warmup_scheduler = torch. optim. lr_scheduler. LambdaLR (optimizer, warmup_factor) cos_scheduler = torch. optim. lr_scheduler. CosineAnnealingLR ( optimizer, T_max=num_steps-warmup_steps, eta_min=0.1 *self. config. learning_rate ) scheduler = torch. optim. lr_scheduler. SequentialLR (optimizer, schedulers=[warmup_scheduler, cos_scheduler], milestones=[warmup_steps]) Furthermore, the training script: • Writes logs to display metrics on tensorboard using torch’s SummaryWriter. • Has a method to save the checkpoint • And doesn’t have partial checkpointing and resume features! I’ve been YOLOing through my longer training runs! Model hyperparameters The model hyperparameters are identical to the ones for Huggingface’s SmolLM 135M model, with one exception. Before we discuss the exception, let’s discuss the other params, and compare how these differ from GPT-2 135M. model_config = ModelConfig( vocab_size=tokenizer. vocab_size , # 49152 d_model=576, d_head=64, # 576/9 d_mlp_proj=1536, n_layers=30, n_kv_heads=3, # 3 K,V heads serve each Q head n_attn_heads=9, rms_norm_eps=1e-5, initializer_range=0.041666666666666664 , rope_theta=100000.0 , padding_idx=tokenizer. pad_token_id ) # For GPT2, it would be this gpt2_model_config = ModelConfig( vocab_size=50257 d_model=768, d_head=64, # 768/12 d_mlp_proj=3072, # 4*768 n_layers=12, n_kv_heads=12, # this just means no GQA, just pure MHA n_attn_heads=12, rms_norm_eps=1e-5, initializer_range=0.02 , #rope_theta=100000.0 , - instead we have n_positions=1024 ) There is a stark difference between the two models even though both SmolLM2 and GPT2 end up with 135M params — the number of layers. With 30 layers, the SmolLM architecture has way more than the 12 layers in GPT2. Their SmolLM article discusses this decision: For the architecture of our 135M and 360M parameter models, we adopted a design similar to MobileLLM, incorporating Grouped-Query Attention (GQA) and prioritizing depth over width The Mobile LLM paper discusses their approach: We conducted an extensive study involving the training of 19 models, including 9 models with ∼125M parameters and 10 models with ∼350M parameters. Each model is designed with a similar size but varied in terms of depth and width And here are their results; In the first graph, if we compare 12 layers (light green) vs 30 layers (light orange) plots, the 30 layer model performs better on arc-easy, PiQA, HellaSwag, OBQA (OpenBookQA) and WinoGrande, ending up with one of the highest average scores with no significant weaknesses (as opposed to 62 layers which performs worse on OpenBookQA). The second and third graphs isolate results for TQA (TriviaQA) and RACE, both reading comprehension benchmarks, which show even more of a difference between shallow and deep model architectures. All of these are language and common sense reasoning benchmarks, relevant to some of the most common practical usage of LLMs (like answering questions from text). The MobileLLM paper has some more studies on model architecture and here are some relevant excerpts: On SwiGLU By changing the vanilla FFN (FC → ReLU → C) to SwiGLU, The average performance on zero-shot reasoning tasks is boosted from 42.6 to 43.9 for the 125M model On the number of attention heads and using GQA Results in Figure 5 show that using 16 query heads produces the best results. Additionally, reducing the number of kv-heads from 16 to 4 resulted in comparable accuracy for the 125M model In the figure above, measuring the average accuracy on reasoning tasks, you can see that using GQA with a ratio=4 (16 Attn heads, 4 KV heads) results in the same performance as MHA with 16 heads (ratio=1). Thus, you can effectively reduce the parameter count with no loss in performance (and improvement in inference performance). Now, let’s get to the exception. One thing I changed from all of these architectures (GPT2, MobileLLM, SmolLM) is the tying of weights for input and output embeddings. Remember that the input embedding (called embed_tokens in our code) is a Linear layer with in_features=49152 and out_features=576 , and the output embedding (called lm_head in our code) is just the transpose of that. That’s 28M params, a pretty significant portion of the number of params for a model in the 135M class. It could make sense to reduce the model size noticeably, and tie the weights — since essentially both are mappings between latent vector spaces and tokens. But, it puts a constraint on the model — for example if the model wants to output “chocolate” (assume it’s one token for now), it has to recover effectively the same embedding vector after progressing through all the decoder layers as the input embedding for “chocolate”. This constraint implies that the model must dedicate some of its processing bandwidth through the decoder blocks to make sure the decoder stack output maps to the same embedding space. For small models like these, it could be argued that saving the number of params is more important (given the %age impact), and that’s what MobileLLM’s authors think. Their studies show that tying these weights and adding more depth is a win-win approach, as shown in the table below: So why didn’t I tie the embeddings? One, I was watching Neel Nanda’s video on model interpretability, where he discusses the “Zero-layer transformer”, how that is effectively a bigram model and that input and output embeddings do fundamentally different things. In a zero layer transformer, the model is meaningless if you make them the same. And in reality, most larger models don’t tie weights, especially as the impact to model size is much less. I wanted to train models both ways — with and without tying embeddings. I started with the models without weight tying, hence the size of 162M as opposed to 135M for SmolLM. But, as of writing this post, I haven’t gotten to training the 135M versions with weight tying. Keep an eye out for my observations on this debate. Training hyperparameters For training, we use a mix of hyperparameters used by SmolLM and by Llama2. In the SmolLM article, they use the learning rate of 3e-3, batch size of 1M (in tokens), and a “a trapezoidal learning rate scheduler with a cooldown phase equal to 20% of the total training time”. They use a context length of 2048 tokens. Llama 2 paper details their of hyperparameterss Given these, I ended up with values as shown below. The trainer implements uses a cosine LR scheduler that falls down to 10% of the peak learning rate. The warmup_ratio is set to get approx 2000 steps of warmup. The learning rate value is the same as what Andrej Karpathy used in his videos, and is somewhat in the middle of SmolLM’s values and Llama’s values. Varying LR is another thing I want to test, and see if we can get more juice from the same amount of data. # only showing training hyperparameters train_config = TrainerConfig( per_device_train_batch_size=32, max_seq_len=2048, num_epochs=1, learning_rate=1e-3, grad_clip_norm=1.0 , warmup_ratio=0.01 , ) The effective batch size was 0.5M , lower than the 1M for SmolLM. I chose this because it seemed sufficient, given that Karpathy also uses a batch size of 0.5M. Performance optimization For large scale training, optimizing the performance is critical — and most frontier model builders have large teams dedicated to solving this problem, as you can see in the Pre-Training -> Infrastructure section of the Llama3 paper. I had a personal machine with a 3090, and also access to a box with 8xA100 (80GB) GPUs. I largely retraced the same journey as Andrej Karpathy in his videos. Here is the summary: I am curious what performance I would get if I used an existing library like nanotron (though they don’t have torch. compile support as of Jan 2025) or torchtitan. Now, let’s get to actually training some models! Definitely some interesting observations and discussions to be had through these sections below. Experiment 1: Train on a paragraph To test the code, I trained the model on a single paragraph — which is the Luthen’s famous monologue in the TV Series Andor. Calm. Kindness. Kinship. Love. I’ve given up all chance at inner peace. I’ve made my mind a sunless space. I share my dreams with ghosts. I wake up every day to an equation I wrote 15 years ago from which there’s only one conclusion, I’m damned for what I do. My anger, my ego, my unwillingness to yield, my eagerness to fight, they’ve set me on a path from which there is no escape. I yearned to be a savior against injustice without contemplating the cost and by the time I looked down there was no longer any ground beneath my feet. What is my sacrifice? I’m condemned to use the tools of my enemy to defeat them. I burn my decency for someone else’s future. I burn my life to make a sunrise that I know I’ll never see. And the ego that started this fight will never have a mirror or an audience or the light of gratitude. So what do I sacrifice? Everything! You’ll stay with me, Lonni. I need all the heroes I can get. I ran this for 64 epochs with a sequence length of 256 (longer than the paragraph) and wanted to check if the model can memorize this. This is implemented in the main() of train. py. The final train loss was pretty much 0:... Step: 60, Training Loss: 0.00148 , LR: 0.0000993 , Tokens/sec: 14006.74 Step: 61, Training Loss: 0.00142 , LR: 0.0000995 , Tokens/sec: 13764.64 Step: 62, Training Loss: 0.00137 , LR: 0.0000997 , Tokens/sec: 13765.32 Step: 63, Training Loss: 0.00133 , LR: 0.0000999 , Tokens/sec: 13742.13 With a prompt of just Calm. Kindness. I get the full paragraph back: input_ids = tokenizer(["Calm. Kindness."], return_tensors="pt")['input_ids'].cuda() idx = model. generate (input_ids, temperature=0.01 , top_k=5, max_new_tokens=240) print(tokenizer. batch_decode (idx)) # Outputs: # ["Calm. Kindness. Kinship. Love. I've given up all chance at inner peace. I've made my mind a sunless space. I share my dreams with ghosts. I wake up every day to an equation I wrote 15 years ago from which there's only one conclusion, I'm damned for what I do. My anger, my ego, my unwillingness to yield, my eagerness to fight, they've set me on a path from which there is no escape. I yearned to be a savior against injustice without contemplating the cost and by the time I looked down there was no longer any ground beneath my feet. What is my sacrifice? I'm condemned to use the tools of my enemy to defeat them. I burn my decency for someone else's future. I burn my life to make a sunrise that I know I'll never see. And the ego that started this fight will never have a mirror or an audience or the light of gratitude. So what do I sacrifice? Everything! You'll stay with me, Lonni. I need all the heroes I can get.<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>"] This surprised me. I’ve always seen LLMs as incapable of 100% accurate memorization. I don’t know if someone knows the answer to this question — what is the size of the dataset that an LLM (of a given size) can memorize its dataset sequences perfectly? It also suggests pretraining a model on a specialized dataset has value. There is a perception that only extreme-scale datasets can lead to meaningful outputs from a model. If we don’t want the full generalized ability of GPT-like models, then models trained without the noise of the web’s data can potentially do much better on specialized tasks. Experiment 2: Train on tiny Shakespeare As a right of passage for anyone who’s seen Karpathy’s videos, I ran the training on the Tiny Shakespeare dataset. I did try to attempt memorization, but after some attempts, I am unable to reproduce even a similar response to a test prompt prefix I have been trying (the train loss was 1.55 ): input_text = """ All: Content, content. MENENIUS: O sir, you are not right: have you not known The worthiest men have done't? CORIOLANUS: """.strip() # Expected continuation # # CORIOLANUS: # What must I say? # 'I Pray, sir'--Plague upon't! I cannot bring # My tongue to such a pace:--'Look, sir, my wounds! # I got them in my country's service, when # Some certain of your brethren roar'd and ran # From the noise of our own drums.' input_ids = tokenizer([input_text], return_tensors="pt")['input_ids'].to(trainer. device ) idx = model. generate (input_ids, temperature=0.01 , top_k=5, max_new_tokens=64) print(tokenizer. batch_decode (idx)[0]) # Outputs: # CORIOLANUS: # I am, no more. # MENENIUS: # A thousandly he dost! # CORIOLANUS: # I am, I have more by the people, # To see my good my master. # MENENIUS: # I am it more! That said, it is not identical to the last example because when the model must have seen this text, it must have seen a different prefix to the expected generation. I have some ideas to improve the test setup here and I might write another post on LLMs memorizing Shakespeare. Experiment 3: Train on Hindi Wikipedia Next we move on to somewhat larger datasets, and start using the larger scale tokenization and data loader workflows. The Hindi section of Wikipedia consists of 240M tokens. For this training run, I didn’t try to drive down the loss since I just wanted to test the pipeline, but I do find that the generations make some valid sentences, if not very sensible ones. It did surprise me how bad the tokenization is for Hindi - 64 tokens is barely a sentence. I should get back to non-Latin language training with a different tokenizer and potentially a larger dataset. Experiment 4: Train with Fineweb-edu, 10 Billion token sample This is one of the two runs the whole work is building towards. Fineweb-Edu is a series of datasets with 1.3T and 5.4T Tokens sized datasets. It’s part of Huggingface’s FineWeb initiative to curate datasets for pretraining. They use the Datatrove library to write a custom data curation workflow that includes many steps that filter out data based on different quality criteria. Pre-training data curation is a big area of research. Both Llama 2 and Llama 3 papers mention significant work done to improve and enlarge the pre-training datasets. From the Llama 3 paper: Llama 3 uses a standard, dense Transformer architecture. It does not deviate significantly from Llama and Llama 2 in terms of model architecture; our performance gains are primarily driven by improvements in data quality and diversity as well as by increased training scale Here are the tensorboard graphs from this training run. The graphs look normal, with no noticeable issues. The train and eval loss are 2.7763 and 2.8795 , respectively. Compared to Andrej Karpathy’s loss values (taken from his video), they are slightly better than the 2.9220 and 3.0726 that his training run achieved with the same dataset. On hellaswag and arc, we get these results: | Tasks |Version|Filter|n-shot| Metric | |Value | |Stderr| |-------------|------:|------|-----:|--------|---|-----:|---|-----:| |arc_challenge| 1|none | 0|acc |↑ |0.2270 |± |0.0122 | | | |none | 0|acc_norm|↑ |0.2637 |± |0.0129 | |arc_easy | 1|none | 0|acc |↑ |0.5518 |± |0.0102 | | | |none | 0|acc_norm|↑ |0.4975 |± |0.0103 | |hellaswag | 1|none | 0|acc |↑ |0.3030 |± |0.0046 | | | |none | 0|acc_norm|↑ |0.3425 |± |0.0047 | Andrej Karpathy’s video reports a value of 0.3068 for arc, and our value of 0.3030 +- 0.0046 is pretty much identical. The checkpoint is available here. It is compatible with the transformers library and can be used like any other text generation model on Huggingface. Experiment 5: Train with Fineweb-edu, 100 Billion token sample We repeat the exact same exercise, but now with the 100B token sample of Fineweb-edu. Here are the curves from that run: They look fairly similar, other than the train/grad_norm curve, which increases consistently for most of the training. Normally, this is counter-intuitive, as an increasing gradient norm typically means that the optimization hasn’t converged. But, as Bengio writes in his Deep Learning book (Ch8: Optimization), this is fairly common in real-life deep learning optimization. The train and eval loss values are 2.7568 and 2.6753 , respectively — showing some improvement, especially in the eval loss. The hellaswag and arc eval metrics show a far more noticeable jump. | Tasks |Version|Filter|n-shot| Metric | |Value | |Stderr| |-------------|------:|------|-----:|--------|---|-----:|---|-----:| |arc_challenge| 1|none | 0|acc |↑ |0.2662 |± |0.0129 | | | |none | 0|acc_norm|↑ |0.2875 |± |0.0132 | |arc_easy | 1|none | 0|acc |↑ |0.5791 |± |0.0101 | | | |none | 0|acc_norm|↑ |0.5135 |± |0.0103 | |hellaswag | 1|none | 0|acc |↑ |0.3401 |± |0.0047 | | | |none | 0|acc_norm|↑ |0.4028 |± |0.0049 | The SmolLM-135M achieves an accuracy of 42.30 on hellaswag and 43.99 on arc averaged (compared to 34.01 and 42.26 from our model). Of course, their dataset smollm-corpus is a more complex mix, including a significant amount of synthetic data, totaling 600M tokens. For this run, I did output sample generations at every eval step. The prompt for them is “The world is”. Here are some generations at different steps Step 3 >>>The world is vibrharv renting Chingarman+= Electrooln Hand labelling Breheetsésano pulpiarstatistics Amendment goosearer weakened catching sqlite Waters Poor ans SurgThese crazy� assembliesribly >>>The world is exit influences stirositantage Beet dubbed zero coach evidencesTodayatha�biased whenever profoundly wolf bareploid Expressions crazytldbahs exertscansuda excretedaser Redist schem hymn >>>The world isCHANTABILITYsplitCompuuidulpholip bru nightmareSoviet Affairs nozzles Patience tranqu erythe opinionDefineilinear trustworthy illumin add Rece BPDotype modernSherantly Monica viol actress Ontariosf uplifting >>>The world issideredConstruction inoc segregation pilots L Bosnia orchid Nervous Apostlesorneromyalgia Criticism sket scrub suggestive PubMed prism newslettersuts======== infectious showcasedNL Bergernez Males Scalally theologians PubMed jackets Step 950 >>>The world is a very difficult place in the 1960s. The first two years of the 1960s were the 196 >>>The world is the most common of the 20th century. The 20th century was the first major period of the 20th century. It >>>The world is a term of the 19th century. The 19th century was the first to be the first to be the first to be the first >>>The world is a very important factor to the research of the Department of the Interior. The National Institute of the Interior has been a major effort to the agency. The agency It’s starting to form sentences pretty quickly. Step 3800 >>>The world is a world of change, and the world is a world of change. The world is a world of change. The world of change is a world of >>>The world is now in a state of crisis. The world is now in a state of crisis. The world is in a state of crisis. The world is now in >>>The world is not the same as the world. The world is not the same as the world. The world is not the same as the world. The world >>>The world is in a dark place. It is not a place where the sun shines. The sun is the sun, and the sun is the sun. The sun Looks like it’s in a pessimistic place! Step 19,950 (about 10B tokens) >>>The world is full of opportunities for the young to learn and grow. The United States is one of the most successful countries in the world. It is a country where children >>>The world is becoming more and more dependent on the internet. The internet is a great resource for students and teachers. It is a great tool for students to learn and share >>>The world is a much more complex place than it was in the 19th century. The world is not only a vast and vast expanse of water, but it is >>>The world is a very different place. We are all different. We are all different. We are all different. We are all different. We are all different. We are Way more optimistic in this set of generations! Step 193800 (end of training) >>>The world is changing fast. The world is changing fast. The world is changing fast. The world is changing fast. The world is changing fast. The world is changing fast >>>The world is changing, and we need to adapt to it. We need to be more creative in our thinking, and we need to be more flexible in our thinking. >>>The world is a big place, and we all have a responsibility to be a part of the solution. The World Health Organization (WHO) has identified 10 key >>>The world is moving towards a more sustainable future, and the global community is working to reduce greenhouse gas emissions. One way to do this is through the use of renewable energy sources Here, it's attempting to spur some action! LLMs in general have a tendency to repeat sentences, and that’s something I’ve seen even when doing Continual Pretraining with 70B Instruct models! I’m not sure what causes this behavior, though. This checkpoint is available here. What’s next First of all, thank you for reading all the way! My hope is that it provides value to people learning more about this area of technology, and maybe also inculcate the spirit of tinkering and getting your hands dirty. There are tons of ideas for future attempts — simple architecture changes, significant architecture changes, discussing more evals than the couple we’ve focused on, and some more whackier ideas that I’m just ruminating. As the title suggests, this is just the beginning of the journey.