update eagle3 example; add export

Author: h-guo18

examples/speculative_decoding/README.md

@@ -26,7 +26,7 @@
 mto.enable_huggingface_checkpointing()
 
 
-def validate_ar(model, tokenizer, ds, steps=3, osl=20, num_samples=20, device=None):
+def validate_ar(model, tokenizer, ds, steps=3, osl=20, num_samples=80, device=None):
     validator = HFARValidation(model, tokenizer)
     num_samples = min(num_samples, len(ds))
     ars = []
@@ -54,12 +54,12 @@ def validate_ar(model, tokenizer, ds, steps=3, osl=20, num_samples=20, device=No
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_path", type=str, required=True, help="Path to model directory")
-    parser.add_argument("--steps", type=int, default=1, help="Steps for AR validation")
+    parser.add_argument("--steps", type=int, default=3, help="Steps for AR validation")
     parser.add_argument(
-        "--osl", type=int, default=100, help="Output sequence length for AR validation"
+        "--osl", type=int, default=32, help="Output sequence length for AR validation"
     )
     parser.add_argument(
-        "--num_samples", type=int, default=20, help="Number of MT-Bench samples to use"
+        "--num_samples", type=int, default=80, help="Number of MT-Bench samples to use"
     )
     parser.add_argument(
         "--ar_lower_bound",

examples/speculative_decoding/ar_validate.py

@@ -28,11 +28,10 @@ def main():
     parser.add_argument("--model", type=str, required=True, help="Model name or path for tokenizer")
     parser.add_argument("--data", type=str, required=True, help="Path to training data (jsonl)")
     parser.add_argument(
-        "--eagle_config",
-        type=str,
+        "--draft_vocab_size",
+        type=int,
         required=True,
-        default="eagle_config.json",
-        help="Path to eagle_config.json",
+        help="Draft vocab size",
     )
     parser.add_argument(
         "--calibrate_size",
@@ -45,12 +44,6 @@ def main():
     )
     args = parser.parse_args()
 
-    with open(args.eagle_config) as f:
-        eagle_config = json.load(f)
-    if "draft_vocab_size" not in eagle_config:
-        print("No draft vocab size specified in eagle_config.json, no need to calibrate for d2t.")
-        return
-
     print("Calibrating vocab...")
     tokenizer = AutoTokenizer.from_pretrained(args.model)
     with open(args.data) as f:
@@ -59,7 +52,7 @@ def main():
             conversations = conversations[: args.calibrate_size]
         conversations = [item for sublist in conversations for item in sublist]
 
-    d2t = calibrate_frequent_vocab(tokenizer, conversations, eagle_config["draft_vocab_size"])
+    d2t = calibrate_frequent_vocab(tokenizer, conversations, args.draft_vocab_size)
     model_name = os.path.basename(os.path.normpath(args.model))
     vocab_path = os.path.join(args.save_dir, model_name, "d2t.pt")
     os.makedirs(os.path.dirname(vocab_path), exist_ok=True)

examples/speculative_decoding/calibrate_draft_vocab.py

@@ -1,3 +1,10 @@
 {
-    "draft_vocab_size": 32000
+    "rope_scaling": {
+        "factor": 32.0,
+        "low_freq_factor": 1.0,
+        "high_freq_factor": 4.0,
+        "original_max_position_embeddings": 8192,
+        "rope_type": "llama3"
+    },
+    "initializer_range": 0.02
 }

examples/speculative_decoding/eagle_config.json

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+import torch
+from transformers import AutoModelForCausalLM
+
+import modelopt.torch.opt as mto
+from modelopt.torch.export import export_hf_checkpoint
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default="")
+    parser.add_argument("--export_path", type=str, default="")
+    return parser.parse_args()
+
+
+mto.enable_huggingface_checkpointing()
+
+args = parse_args()
+model = AutoModelForCausalLM.from_pretrained(args.model_path, torch_dtype="auto")
+with torch.inference_mode():
+    export_hf_checkpoint(
+        model,  # The quantized model.
+        export_dir=args.export_path,  # The directory where the exported files will be stored.
+    )
+print(f"Exported checkpoint to {args.export_path}")

examples/speculative_decoding/export_hf_checkpoint.py

@@ -46,6 +46,7 @@
 import modelopt.torch.opt as mto
 import modelopt.torch.speculative as mtsp
 from modelopt.torch.utils import print_rank_0
+from modelopt.torch.utils.distributed import is_master
 
 torch.manual_seed(0)
 mto.enable_huggingface_checkpointing()
@@ -170,6 +171,8 @@ def train():
                 {
                     "hidden_size": model.config.hidden_size,
                     "vocab_size": model.config.vocab_size,
+                    # we also overwrite max_pos_embedding for deployment compatibility
+                    "max_position_embeddings": model.config.max_position_embeddings,
                     "draft_vocab_size": custom_config["draft_vocab_size"]
                     if eagle_args.eagle_config and "draft_vocab_size" in custom_config
                     else model.config.vocab_size,
@@ -202,6 +205,15 @@ def train():
     class ARValidationCallback(TrainerCallback):
         def __init__(self, ar_validate_steps: int = 500):
             self.ar_validate_steps = ar_validate_steps
+            self.wandb = None
+            if is_master():
+                try:
+                    import wandb
+
+                    self.wandb = wandb
+                    self.wandb.init()
+                except ImportError:
+                    pass
 
         def on_step_end(self, args, state, control, **kwargs):
             if state.global_step % self.ar_validate_steps == 0 and state.global_step > 0:
@@ -213,6 +225,8 @@ def on_step_end(self, args, state, control, **kwargs):
                     device=kwargs["model"].device,
                 )
                 print_rank_0(f"Step {state.global_step} AR: {sum(ars) / len(ars):.4f}")
+                if self.wandb:
+                    self.wandb.log({"validate_ar": sum(ars) / len(ars)}, step=state.global_step)
             return control
 
     trainer = Trainer(

examples/speculative_decoding/launch.sh

@@ -46,7 +46,7 @@
 parser.add_argument(
     "--max_tokens", type=int, default=2048, help="Maximum number of tokens to generate"
 )
-parser.add_argument("--chat", action="store_true", help="Use chat mode")
+parser.add_argument("--chat", action="store_true", default=True, help="Use chat mode")
 parser.add_argument("--model", type=str, default="model", help="Model name")
 parser.add_argument("--url", type=str, default="http://localhost:8000/v1", help="URL of the API")
 parser.add_argument("--api_key", type=str, default="token-abc123", help="API key (if any)")