profile

Signed-off-by: h-guo18 <[email protected]>

Author: h-guo18

examples/speculative_decoding/distill_trainer.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 import json
 import os
+import time
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 from abc import abstractmethod
@@ -42,6 +43,7 @@
 # Hyperparameters for profiling
 LOG_INTERVAL = 100
 SAVE_INTERVAL = 20000
+TOTAL_STEPS = 500
 
 # Shape and dtype description of the distillation signal
 DistillMetadata = dict[str, tuple[torch.Size, torch.dtype]]
@@ -204,6 +206,10 @@ def train(self):
                     )
                     for i, batch in enumerate(pbar):
                         global_step = epoch * len(self.dataloader) + i
+                        if global_step >= TOTAL_STEPS:
+                            break
+                        if global_step == 50:
+                            self.start_time = time.time()
                         inputs = {k: v.to(self.model.device) for k, v in batch.items()}
 
                         # Receive distill messages from teacher
@@ -241,10 +247,18 @@ def train(self):
             # Inference Loop
             for epoch in range(self.args.epoch):
                 for i, batch in enumerate(self.dataloader):
+                    global_step = epoch * len(self.dataloader) + i
+                    if global_step >= TOTAL_STEPS:
+                        break
+                    if global_step == 50:
+                        self.start_time = time.time()
                     inputs = {k: v.to(self.model.device) for k, v in batch.items()}
                     with torch.inference_mode():
                         self._send_to_student(self.teacher_step(self.model, inputs))
 
+        self.average_step_time = (time.time() - self.start_time) / (TOTAL_STEPS-50)
+        print(f"Rank {self.rank} average step time: {self.average_step_time}")
+
         self._print_mem_stats()
         # Makesure all processes finished before destroy.
         dist.barrier()
@@ -321,6 +335,7 @@ def _prepare_student_model(self):
             process_group=self.args.student_pgroup,
             find_unused_parameters=True,
         )
+        self._print_mem_stats()
         return model
 
     @property

examples/speculative_decoding/eagle_config.json

@@ -6,5 +6,6 @@
         "original_max_position_embeddings": 8192,
         "rope_type": "llama3"
     },
-    "initializer_range": 0.02
+    "initializer_range": 0.02,
+    "head_dim": 64
 }

examples/speculative_decoding/eagle_utils.py

@@ -42,7 +42,10 @@
 
 
 def preprocess(examples, tokenizer):
-    tokenizer.chat_template = tokenizer.chat_template.replace(REMOVE_THINK_CHAT_TEMPLATE, "")
+    if tokenizer.chat_template:
+        tokenizer.chat_template = tokenizer.chat_template.replace(REMOVE_THINK_CHAT_TEMPLATE, "")
+    else:
+        tokenizer.chat_template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}"
     new_examples = {
         "input_ids": [],
         "attention_mask": [],

examples/speculative_decoding/train.py

@@ -71,7 +71,7 @@ def main():
     parser = argparse.ArgumentParser(description="Multi-GPU distributed two-stage forward example")
     parser.add_argument("--model_path", type=str, default="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
     parser.add_argument("--student_devices", type=list, default=[0, 1, 2, 3])
-    parser.add_argument("--teacher_devices", type=list, default=[4, 5])
+    parser.add_argument("--teacher_devices", type=list, default=[4, 5, 6, 7])
     parser.add_argument(
         "--data_path", type=str, default="data/magpie_llama3.2_1b_generated/data.cleaned.jsonl"
     )