From c5c87fc4e7523aaaea61a8167ddb028f829bc533 Mon Sep 17 00:00:00 2001
From: Aleksa Gordic <gordicaleksa@gmail.com>
Date: Thu, 8 Aug 2024 20:54:17 +0200
Subject: [PATCH] Add comments, fix stop tokens

---
 train_llama3.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/train_llama3.py b/train_llama3.py
index 5cb7b97..4013db7 100644
--- a/train_llama3.py
+++ b/train_llama3.py
@@ -321,6 +321,8 @@ class LLaMA(nn.Module):
 
     @staticmethod
     def adapt_llama_state_dict_keys(checkpoint, config: LlamaConfig):
+        # Modify key names from Meta's LLaMA to our LLaMA
+        # our key names are derived from GPT-2's key names
         checkpoint['transformer.wte.weight'] = checkpoint.pop('tok_embeddings.weight')
 
         for i in range(config.n_layer):
@@ -355,6 +357,8 @@ class LLaMA(nn.Module):
 
     @staticmethod
     def adapt_llama_state_dict_keys_hf(checkpoint, config: LlamaConfig):
+        # Modify key names from HuggingFace's LLaMA to our LLaMA
+        # our key names are derived from GPT-2's key names
         checkpoint['transformer.wte.weight'] = checkpoint.pop('model.embed_tokens.weight')
 
         # We need to unpermute K and V because HF script permuted the original Meta-LLaMA weights
@@ -432,9 +436,6 @@ class LLaMA(nn.Module):
         torch.set_default_tensor_type(torch.tensor([], dtype=original_default_type, device="cpu").type())  # restore default type
 
         tokenizer = Tokenizer(model_path=tokenizer_path)
-        # add <|end_of_text|> as the stop token for base model - this is an omission in the reference code
-        # the reference code only adds instruct model stop tokens...
-        tokenizer.stop_tokens = tokenizer.stop_tokens + [128001]
         model.tokenizer = tokenizer
         return model
 
@@ -676,9 +677,10 @@ class Tokenizer:
         self.eom_id: int = self.special_tokens["<|eom_id|>"]
         self.python_tag_id = self.special_tokens["<|python_tag|>"]
         self.pad_id: int = self.special_tokens["<|finetune_right_pad_id|>"]
+        # hardcoded stop tokens for the base model
         self.stop_tokens = [
-            self.special_tokens["<|eom_id|>"],
-            self.special_tokens["<|eot_id|>"],
+            self.special_tokens["<|begin_of_text|>"],
+            self.special_tokens["<|end_of_text|>"],
         ]
 
     def encode(