Saves layer output stack

Signed-off-by: Flavia Beo <flavia.beo@ibm.com>

Author: flaviabeo

scripts/generate_layers_metrics.py

@@ -1,4 +1,5 @@
 import os
+import json
 import time
 import logging
 import argparse
@@ -205,7 +206,7 @@ def __register_call_layers(model, batch_size, device, seq_length, max_new_tokens
         tokenizer (Tokenizer): The tokenizer used for tokenization.
 
     Returns:
-        list: A list of tuples containing the name and output of each layer in the model.
+        dict: A dict containing the name and output of each layer in the model.
     """
     layer_stack = {}
     pt_compile_model_time = time.time()
@@ -389,6 +390,9 @@ def generate_layers_metrics(model_path, batch_size, seq_length, max_new_tokens):
                                             seq_length=seq_length, max_new_tokens=max_new_tokens, 
                                             tokenizer=tokenizer)
 
+    with open(os.join(output_path,f"{model_path}-layer-output-stack-cpu.json"), 'w') as f:
+        json.dump(layer_stack_cpu, f)
+    
     global generate_iters
     generate_iters = 0
     logger.info(f"Finished registering CPU layers")
@@ -398,6 +402,9 @@ def generate_layers_metrics(model_path, batch_size, seq_length, max_new_tokens):
                                              device="cuda", 
                                              seq_length=seq_length, max_new_tokens=max_new_tokens, 
                                              tokenizer=tokenizer)
+    
+    with open(os.join(output_path,f"{model_path}-layer-output-stack-gpu.json"), 'w') as f:
+        json.dump(layer_stack_cuda, f)
 
     assert len(layer_stack_cuda.keys()) == len(layer_stack_cpu.keys())