squash: new trainer with HF and SGL backend

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

Author: h-guo18

examples/speculative_decoding/ar_validate.py

@@ -26,11 +26,14 @@
 mto.enable_huggingface_checkpointing()
 
 
-def validate_ar(model, tokenizer, ds, steps=3, osl=20, num_samples=80, device=None):
+def validate_ar(
+    model, tokenizer, ds, steps=3, osl=20, num_samples=80, device=None, disable_pbar=False
+):
     validator = HFARValidation(model, tokenizer)
     num_samples = min(num_samples, len(ds))
     ars = []
-    for i in tqdm(range(num_samples), desc="Validating AR"):
+    print("validating AR...")
+    for i in tqdm(range(num_samples), disable=disable_pbar):
         prompt = ds[i]["prompt"][0]
         input_ids = tokenizer(prompt, return_tensors="pt").input_ids
         # Apply chat template to the prompt, continuing with assistant response

examples/speculative_decoding/eagle_utils.py

@@ -518,7 +518,7 @@ def compute_loss(self, *args, **kwargs):
         kwargs.pop("num_items_in_batch", None)
         loss, outputs = super().compute_loss(return_outputs=True, *args, **kwargs)
         if hasattr(outputs, "train_acc"):
-            self.state.training_accs.append(outputs.train_acc)
+            self.state.training_accs.append([acc.item() for acc in outputs.train_acc])
         return loss
 
 

examples/speculative_decoding/train.py

@@ -19,11 +19,10 @@
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
-from distill_trainer import EagleTPTrainer
 from eagle_utils import DataCollatorWithPadding, make_eagle_supervised_data_module
+from trainer.distill_trainer import EagleSGLTrainer, EagleTPTrainer
 from transformers import AutoTokenizer
 
-# Hyperparameters for profiling
 torch.manual_seed(0)
 
 
@@ -34,10 +33,10 @@ def _setup_distributed(rank, args, backend="nccl"):
     os.environ["LOCAL_RANK"] = str(rank)
     # Initialize process group
     dist.init_process_group(backend, rank=rank, world_size=args.world_size)
-    if rank in args.student_ranks:
-        torch.cuda.set_device(args.student_devices[rank])
+    if rank in args.teacher_ranks:
+        torch.cuda.set_device(args.teacher_devices[rank])
     else:
-        torch.cuda.set_device(args.teacher_devices[rank - len(args.student_ranks)])
+        torch.cuda.set_device(args.student_devices[rank - len(args.teacher_ranks)])
     print(
         f"Starting process rank={rank}, device={torch.cuda.current_device()}, world_size={args.world_size}"
     )
@@ -51,7 +50,10 @@ def train(rank, args):
     tokenizer = AutoTokenizer.from_pretrained(
         args.model_path, model_max_length=args.training_seq_len
     )
-    data_module = make_eagle_supervised_data_module(tokenizer, args, use_offline_training=False)
+    args.use_offline_training = False
+    args.vlm_processor = None
+    args.offline_data_path = None
+    data_module = make_eagle_supervised_data_module(tokenizer, args)
 
     train_dataloader = torch.utils.data.DataLoader(
         data_module["train_dataset"],
@@ -61,42 +63,56 @@ def train(rank, args):
         collate_fn=DataCollatorWithPadding(max_length=args.training_seq_len),
         drop_last=True,
     )
-
-    trainer = EagleTPTrainer(rank, args, tokenizer, train_dataloader)
+    trainer_cls = {
+        "sglang": EagleSGLTrainer,
+        "hf": EagleTPTrainer,
+    }[args.teacher_backend]
+    trainer = trainer_cls(rank, args, tokenizer, train_dataloader)
     trainer.train()
-    trainer.save_pretrained(args.out_path)
+    trainer.save(args.out_path)
 
 
 def main():
     parser = argparse.ArgumentParser(description="Multi-GPU distributed two-stage forward example")
+
+    # Training args
     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, 6, 7])
-    parser.add_argument(
-        "--data_path", type=str, default="data/magpie_llama3.2_1b_generated/data.cleaned.jsonl"
-    )
+    parser.add_argument("--data_path", type=str, required=True, help="Training dataset.")
     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")
     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=8, help="Total bs across all ranks.")
+
+    # Trainer args
+    parser.add_argument("--teacher_backend", type=str, choices=["sglang", "hf"], default="sglang")
+    parser.add_argument(
+        "--teacher_ep_size",
+        type=int,
+        default=1,
+        help="Teacher EP size, only used for sglang backend.",
+    )
+    parser.add_argument("--teacher_devices", type=list, default=[0, 1, 2, 3])
+    parser.add_argument("--student_devices", type=list, default=[4, 5, 6, 7])
+    parser.add_argument(
+        "--lazy_preprocess", type=bool, default=True, help="Whether to use lazy preprocessing."
+    )
+    parser.add_argument("--log_interval", type=int, default=50)
+    parser.add_argument("--save_interval", type=int, default=20000)
     parser.add_argument(
-        "--batch_size", type=int, default=4, help="Total batch size across all parallel ranks."
+        "--total_steps", type=int, default=60000, help="Total number of steps for debugging."
     )
     parser.add_argument("--master_port", type=str, default="12357")
 
     args = parser.parse_args()
     # TODO: add sanity check for args
 
     def set_ranks(args):
-        # 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(
-            range(len(args.student_devices), len(args.student_devices) + len(args.teacher_devices))
+        args.teacher_ranks = list(range(len(args.teacher_devices)))
+        args.student_ranks = list(
+            range(len(args.teacher_devices), len(args.teacher_devices) + len(args.student_devices))
         )
 
     set_ranks(args)