Update README.md

README.md

@@ -26,7 +26,7 @@ Original [Dataset](https://huggingface.co/datasets/roneneldan/TinyStories) conta
 
 ## Model description
 
-A very very small model (3M params) tarined on a very small dataset
+A very very small model (8M params) tarined on a very small dataset
 
 A [sample inference script](https://huggingface.co/Norod78/TinyStories-3M-val-Hebrew/blob/main/TinyStories-3M-val-Hebrew-inference.py) is available
 
@@ -47,4 +47,31 @@ The following hyperparameters were used during training:
 - Transformers 4.31.0.dev0
 - Pytorch 2.0.0
 - Datasets 2.13.1
-- Tokenizers 0.13.3
+- Tokenizers 0.13.3
+
+- ### Parameter calculation
+
+```
+  def gpt_params(seq_len, vocab_size, d_model, num_heads, num_layers):
+    """ Given GPT config calculate total number of parameters """
+    ffw_size = 4*d_model # in GPT the number of intermediate features is always 4*d_model
+    # token and position embeddings
+    embeddings = d_model * vocab_size + d_model * seq_len
+    # transformer blocks
+    attention = 3*d_model**2 + 3*d_model # weights and biases
+    attproj = d_model**2 + d_model
+    ffw = d_model*(ffw_size) + ffw_size
+    ffwproj = ffw_size*d_model + d_model
+    layernorms = 2*2*d_model
+    # dense
+    ln_f = 2*d_model
+    dense = d_model*vocab_size # note: no bias here
+    # note: embeddings are not included in the param count!
+    total_params = num_layers*(attention + attproj + ffw + ffwproj + layernorms) + ln_f + dense
+    return total_params
+
+#gpt2 = dict(seq_len = 1024, vocab_size = 50257, d_model = 768, num_heads = 12, num_layers = 12)
+gpt2 = dict(seq_len = 256, vocab_size = 50259, d_model = 128, num_heads = 16, num_layers = 8)
+result = gpt_params(**gpt2)/1e6
+print(result) #Prints 8.019584
+```