fixed the rmsnorm calculation

configs/model/mpgpt.yaml

@@ -14,4 +14,4 @@ scheduler:
   _target_: src.schedulers.linear_warmup_cosine_decay
   _partial_: True
   warmup_steps: 2
-  total_steps: 10_000
+  total_steps: 20_000

scripts/generate.py

@@ -10,6 +10,7 @@
 
 
 from src.models.gpt import GPT
+from src.models.mpgpt import MPGPT
 
 T.set_grad_enabled(False)
 
@@ -19,17 +20,17 @@ def main():
     enc = tiktoken.get_encoding("gpt2")
 
     # Load the model from the lightning checkpoint.
-    model_path = "/srv/beegfs/scratch/groups/rodem/nlp/gpt/GPT/checkpoints/last.ckpt"
-    model = GPT.load_from_checkpoint(model_path)
+    model_path = "/srv/beegfs/scratch/groups/rodem/nlp/gpt/MP-GPT/checkpoints/last.ckpt"
+    model = MPGPT.load_from_checkpoint(model_path)
     model.eval()
     model.requires_grad_(False)
 
     # Get the maximum input tokens
     max_len = 1024
     max_new_tokens = 100
-    temp = 0.7
-    top_k = 100
-    text = ""
+    temp = 0.9
+    top_k = None
+    text = "The Mystery of the Hidden Garden. Leon Bozianu was no ordinary boy; he was a young adventurer with an insatiable curiosity. Living in the quaint village of St. Eloi, Leon spent his days exploring the picturesque landscape that surrounded his home. His favorite spot was the old, abandoned mansion at the edge of town, a place filled with secrets and wonders. One sunny afternoon, as Leon was wandering through the overgrown garden of the mansion, he stumbled upon an ancient key buried beneath a bed of ivy. The key was intricate, with delicate engravings and a hint of rust. Intrigued, Leon slipped the key into his pocket, determined to uncover its secrets. One sunny afternoon, as Leon was wandering through the overgrown garden of the mansion, he stumbled upon an ancient key buried beneath a bed of ivy. The key was intricate, with delicate engravings and a hint of rust. Intrigued, Leon slipped the key into his pocket, determined to uncover its secrets."
 
     # Do a loop to generate the new tokens
     while True:

src/torch_utils.py

@@ -48,6 +48,12 @@ def rms_norm(x, dim: int | tuple = -1, eps: float = 1e-4) -> T.Tensor:
     return x / n.to(x.dtype)
 
 
+def rms(x, dim: int | tuple = -1, eps: float = 1e-6) -> T.Tensor:
+    """Calculate the RMS of the vector."""
+    n = T.linalg.vector_norm(x.float(), dim=dim, keepdim=True, dtype=T.float32)
+    return T.add(eps, n, alpha=math.sqrt(n.numel() / x.numel()))
+
+
 def get_activations(
     model: nn.Module,
     activation_dict: dict,
@@ -59,7 +65,7 @@ def get_activations(
 
     def hook(name) -> callable:
         def forward_hook(_module: nn.Module, _input: T.Tensor, output: T.Tensor):
-            activation_dict[name] = output.detach().std().cpu().item()
+            activation_dict[name] = rms(output.detach()).mean().cpu().item()
 
         return forward_hook