polish

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

Author: h-guo18

examples/speculative_decoding/distill_trainer.py

@@ -40,10 +40,8 @@
 mto.enable_huggingface_checkpointing()
 
 # Hyperparameters for profiling
-EPOCHS = 1
 LOG_INTERVAL = 100
 SAVE_INTERVAL = 20000
-# VALIDATE_INTERVAL = 20
 
 # Shape and dtype description of the distillation signal
 DistillMetadata = dict[str, tuple[torch.Size, torch.dtype]]
@@ -61,11 +59,11 @@ class BaseDistillTrainer:
 
     def __init__(self, rank, args, tokenizer, dataloader):
         self.rank = rank
-        args.teacher_pgroup = dist.new_group(ranks=args.teacher_ranks)
-        args.student_pgroup = dist.new_group(ranks=args.student_ranks)
         self.args = args
         self.tokenizer = tokenizer
         self.dataloader = dataloader
+
+        # Prepare models
         if rank in args.student_ranks:
             self.model = self.prepare_student_model()
             self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.args.lr)
@@ -180,7 +178,11 @@ def _get_logging_context(self):
             return wandb.init(
                 entity=os.environ["WANDB_ENTITY"],
                 project=os.environ["WANDB_PROJECT"],
-                config={"epochs": EPOCHS, "lr": self.args.lr, "batch_size": self.args.batch_size},
+                config={
+                    "epochs": self.args.epoch,
+                    "lr": self.args.lr,
+                    "batch_size": self.args.batch_size,
+                },
             )
         return nullcontext()
 
@@ -193,7 +195,7 @@ def train(self):
                 self._init_student_recv_buffer()
 
                 # Student training loop
-                for epoch in range(EPOCHS):
+                for epoch in range(self.args.epoch):
                     pbar = (
                         tqdm(self.dataloader)
                         if self.rank == self.args.student_ranks[0]
@@ -236,7 +238,7 @@ def train(self):
 
         else:
             # Inference Loop
-            for epoch in range(EPOCHS):
+            for epoch in range(self.args.epoch):
                 for i, batch in enumerate(self.dataloader):
                     inputs = {k: v.to(self.model.device) for k, v in batch.items()}
                     with torch.inference_mode():
@@ -390,8 +392,10 @@ def student_step(
     ) -> ModelOutput:
         self.optimizer.zero_grad()
 
-        # Chunk inputs for each student rank.
+        # Chunk input_ids and attention_mask for each student rank.
         inputs = {k: v.chunk(len(self.args.student_ranks))[self.rank] for k, v in inputs.items()}
 
         # Second stage forward with provided base model outputs.
-        return self.model(**inputs, base_model_outputs=distill_msgs)
+        output = self.model(**inputs, base_model_outputs=distill_msgs)
+
+        return output

examples/speculative_decoding/train.py

@@ -41,6 +41,8 @@ def _setup_distributed(rank, args, backend="nccl"):
     print(
         f"Starting process rank={rank}, device={torch.cuda.current_device()}, world_size={args.world_size}"
     )
+    args.teacher_pgroup = dist.new_group(ranks=args.teacher_ranks)
+    args.student_pgroup = dist.new_group(ranks=args.student_ranks)
 
 
 def train(rank, args):
@@ -67,47 +69,24 @@ def train(rank, args):
 
 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(
-        "--model_path",
-        type=str,
-        default="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-        help="Path to the model.",
-    )
-    parser.add_argument(
-        "--student_devices", type=list, default=[0, 1, 2, 3], help="Devices for student model"
-    )
-    parser.add_argument(
-        "--teacher_devices", type=list, default=[4, 5], help="Devices for teacher model"
-    )
-    parser.add_argument(
-        "--data_path",
-        type=str,
-        default="data/magpie_llama3.2_1b_generated/data.cleaned.jsonl",
-        help="Path to the training data.",
-    )
-    parser.add_argument(
-        "--training_seq_len",
-        type=str,
-        default=1024,
-        help="Training sequence length.",
-    )
-    parser.add_argument(
-        "--eagle_config_path",
-        type=str,
-        default="eagle_config.json",
-        help="Path to the eagle config.",
+        "--data_path", type=str, default="data/magpie_llama3.2_1b_generated/data.cleaned.jsonl"
     )
+    parser.add_argument("--training_seq_len", type=str, default=1024)
+    parser.add_argument("--eagle_config_path", type=str, default="eagle_config.json")
     parser.add_argument(
         "--lazy_preprocess", type=bool, default=True, help="Whether to use lazy preprocessing."
     )
-    parser.add_argument(
-        "--out_path", type=str, default="ckpts/fast-trained", help="Path to save the model."
-    )
-    parser.add_argument("--lr", type=float, default=1e-5, help="Learning rate.")
+    parser.add_argument("--out_path", type=str, default="ckpts/fast-trained")
+    parser.add_argument("--lr", type=float, default=1e-5)
+    parser.add_argument("--epoch", type=int, default=1)
     parser.add_argument(
         "--batch_size", type=int, default=4, help="Total batch size across all parallel ranks."
     )
-    parser.add_argument("--master_port", type=str, default="12357", help="Master port.")
+    parser.add_argument("--master_port", type=str, default="12357")
 
     args = parser.parse_args()
     # TODO: add sanity check for args