polish

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

Author: h-guo18

examples/speculative_decoding/distill_trainer.py

@@ -49,7 +49,8 @@
 
 class BaseDistillTrainer:
     """
-    Base class for distillation trainer. Initalized and called on every rank.
+    Base distill trainer with basic training loop and overlapped teacher and student steps.
+    Initalized and called on every rank.
     Args:
         rank: rank of the current process
         args: arguments
@@ -252,6 +253,8 @@ def train(self):
 
 
 class EagleTPTrainer(BaseDistillTrainer):
+    """A subclass of BaseDistillTrainer for online eagle training, with base model TP and student DDP."""
+
     def __init__(self, rank, args, tokenizer, dataloader):
         # Load eagle config
         args.eagle_config = EAGLE3_DEFAULT_CFG["config"]

examples/speculative_decoding/train.py

@@ -92,7 +92,7 @@ def main():
     # TODO: add sanity check for args
 
     def set_ranks(args):
-        # TODO(hg): add "no-parallel", "MP", "FSDP".
+        # TODO(hg): This is for TP-DDP setting only. Add "no-parallel", "MP", "FSDP".
         args.world_size = len(args.teacher_devices) + len(args.student_devices)
         args.student_ranks = list(range(len(args.student_devices)))
         args.teacher_ranks = list(