refactor

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

Author: h-guo18

examples/speculative_decoding/eagle_utils.py

@@ -28,7 +28,8 @@
 from datasets import load_dataset
 from packaging.version import Version
 from scripts.ar_validate import validate_ar
-from torch.distributed.tensor.experimental._attention import _SDPAMerger
+
+# from torch.distributed.tensor.experimental._attention import _SDPAMerger
 from torch.utils.data import Dataset
 from transformers import Trainer, TrainerCallback
 from transformers.trainer_pt_utils import LabelSmoother
@@ -382,11 +383,11 @@ def patch_ring_attention_for_ttt():
     )
 
     # 3. Patch merger to skip the blank shard to avoid difference in output.
-    original_sdpa_merger_step = _SDPAMerger.step
+    original_sdpa_merger_step = torch.distributed.tensor.experimental._attention._SDPAMerger.step
 
     def patched_sdpa_merger_step(self, out: torch.Tensor, lse: torch.Tensor, partial: bool):
         if lse.sum() <= 0:
             return
         return original_sdpa_merger_step(self, out, lse, partial)
 
-    _SDPAMerger.step = patched_sdpa_merger_step
+    torch.distributed.tensor.experimental._attention._SDPAMerger.step = patched_sdpa_merger_step

examples/speculative_decoding/main.py

@@ -76,9 +76,9 @@ class DataArguments:
         },
     )
     lazy_preprocess: bool = True
-    draft_vocab_cache_dir: str = field(
-        default="draft_vocab_cache",
-        metadata={"help": "Path to the d2t cache directory."},
+    draft_vocab_cache: str = field(
+        default=None,
+        metadata={"help": "Path to d2t.pt cache file."},
     )
     vlm_img_dir: str = field(default=None, metadata={"help": "Path to the VLM image directory."})
     vlm_processor: str = field(default=None, metadata={"help": "Path to the VLM processor."})
@@ -97,7 +97,7 @@ class TrainingArguments(transformers.TrainingArguments):
     )
     dataloader_drop_last: bool = field(default=True)
     bf16: bool = field(default=True)
-    mode: Literal["eagle1", "eagle3", "medusa"] = "eagle3"
+    mode: Literal["eagle3", "medusa"] = "eagle3"
     estimate_ar: bool = field(
         default=False, metadata={"help": "Whether to estimate AR during training for logging."}
     )
@@ -147,30 +147,35 @@ def train():
         training_args.parallelism_config.sp_backend = None
     print_rank_0(f"arguments: {model_args}, {training_args}, {medusa_args}, {eagle_args}")
 
-    # Detecting last checkpoint.
-    last_checkpoint = None
-    if os.path.isdir(training_args.output_dir):
-        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+    # Detect checkpoint to resume from
+    last_checkpoint = (
+        get_last_checkpoint(training_args.output_dir)
+        if os.path.isdir(training_args.output_dir)
+        else None
+    )
+    if last_checkpoint:
         print_rank_0(f"Last checkpoint detected: {last_checkpoint}")
 
-    checkpoint = None
-    if training_args.resume_from_checkpoint is not None:
-        checkpoint = training_args.resume_from_checkpoint
-    elif last_checkpoint is not None:
-        checkpoint = last_checkpoint
+    checkpoint = training_args.resume_from_checkpoint or last_checkpoint
 
     use_offline_training = data_args.offline_data_path is not None
 
+    model_config = transformers.AutoConfig.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=True
+    )
+    if "vl" in model_config.model_type.lower():
+        model_cls = transformers.AutoModelForVision2Seq
+    else:
+        model_cls = transformers.AutoModelForCausalLM
+
     if checkpoint:
-        model = transformers.AutoModelForCausalLM.from_pretrained(
-            checkpoint, torch_dtype="auto", trust_remote_code=True
-        )
+        model = model_cls.from_pretrained(checkpoint, torch_dtype="auto", trust_remote_code=True)
         tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
     else:
         # To avoid OOM for large models, we load and convert model on CPU first.
         # Model will be moved to GPU during HF trainer.init().
         offline_kwargs = {"num_hidden_layers": 0} if use_offline_training else {}
-        model = transformers.Qwen3VLForConditionalGeneration.from_pretrained(
+        model = model_cls.from_pretrained(
             model_args.model_name_or_path,
             torch_dtype="auto",
             device_map="cpu",
@@ -180,77 +185,38 @@ def train():
         if use_offline_training:
             # When doing offline training, we need to set num_hidden_layers
             # since we override it when loading the model for space savings
-            model_config = transformers.AutoConfig.from_pretrained(
-                model_args.model_name_or_path, trust_remote_code=True
-            )
             model.config.num_orig_hidden_layers = model_config.num_hidden_layers
         tokenizer = transformers.AutoTokenizer.from_pretrained(
             model_args.model_name_or_path,
             model_max_length=training_args.training_seq_len,
             trust_remote_code=True,
         )
-        if tokenizer.chat_template is None:
-            tokenizer.chat_template = (
-                "{%- for message in messages %}"
-                "{{- '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n' }}"
-                "{%- endfor %}"
-            )
-        if tokenizer.pad_token_id is None:
-            tokenizer.pad_token_id = tokenizer.eos_token_id
-
         if training_args.mode == "medusa":
             config = {
                 "medusa_num_heads": medusa_args.medusa_num_heads,
                 "medusa_num_layers": medusa_args.medusa_num_layers,
             }
             mtsp.convert(model, [("medusa", config)])
-        elif training_args.mode in ["eagle1", "eagle3"]:
-            from modelopt.torch.speculative.config import (
-                default_eagle_config,
-                eagle3_default_config,
-                kimik2_eagle_default_config,
+        elif training_args.mode == "eagle3":
+            custom_config = (
+                json.load(open(eagle_args.eagle_config)) if eagle_args.eagle_config else {}
             )
 
-            if eagle_args.eagle_decoder_type == "kimik2":
-                eagle_architecture_config = kimik2_eagle_default_config
-            else:
-                eagle_architecture_config = {
-                    "eagle1": default_eagle_config,
-                    "eagle3": eagle3_default_config,
-                }[training_args.mode]
-
-            if eagle_args.eagle_config:
-                with open(eagle_args.eagle_config) as f:
-                    custom_config = json.load(f)
-                eagle_architecture_config.update(custom_config)
-
             config = {
                 "eagle_decoder_type": eagle_args.eagle_decoder_type,
                 "eagle_offline": use_offline_training,
-                "eagle_architecture_config": eagle_architecture_config,
+                "eagle_architecture_config": custom_config,
             }
 
             mtsp.convert(model, [("eagle", config)])
 
-            # read draft vocab cache
-            if model.eagle_config.draft_vocab_size < model.eagle_config.vocab_size:
-                try:
-                    model_name = os.path.basename(os.path.normpath(model_args.model_name_or_path))
-                    vocab_cache_path = os.path.join(
-                        data_args.draft_vocab_cache_dir, model_name, "d2t.pt"
-                    )
-                    vocab_cache = torch.load(vocab_cache_path)
-                    model.eagle_module.d2t = vocab_cache
-                    print_rank_0(f"Loaded draft vocab cache from {vocab_cache_path}.")
-                except Exception as e:
-                    raise e
         else:
             raise Exception(f"{training_args.mode} is not supported!")
 
     print_rank_0("Loading dataset...")
     if training_args.mode == "medusa":
         data_module = make_medusa_supervised_data_module(tokenizer, data_args)
-    elif training_args.mode in ["eagle1", "eagle3"]:
+    elif training_args.mode == "eagle3":
         data_module = make_eagle_supervised_data_module(
             tokenizer, data_args, train_len=training_args.training_seq_len
         )

modelopt/torch/speculative/config.py

@@ -105,3 +105,8 @@ class EagleConfig(ModeloptBaseConfig):
         default="llama",
         description=("The class of eagle decoder to use. Available options: llama, kimik2"),
     )
+
+    draft_vocab_cache: str = ModeloptField(
+        default=None,
+        description=("Path to d2t.pt cache file."),
+    )

modelopt/torch/speculative/eagle/conversion.py

@@ -20,6 +20,7 @@
 from modelopt.torch.opt.conversion import ModelLikeModule
 from modelopt.torch.opt.dynamic import _DMRegistryCls
 from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict
+from modelopt.torch.speculative.config import eagle3_default_config, kimik2_eagle_default_config
 
 from ..config import EagleConfig
 
@@ -38,6 +39,14 @@ def convert_to_eagle_model(model: nn.Module, config: EagleConfig) -> ConvertRetu
                 EagleDMRegistry.register({original_cls: "base_model_class"})(EagleDMRegistry[cls])
                 break
 
+    # merge custom config with default config
+    default_arch_config = {
+        "llama": eagle3_default_config,
+        "kimik2": kimik2_eagle_default_config,
+    }[config.eagle_decoder_type]
+    custom_config = config.eagle_architecture_config
+    config.eagle_architecture_config = {**default_arch_config, **custom_config}
+
     eagle_model = EagleDMRegistry.convert(model)
     eagle_model.modify(
         eagle_offline=config.eagle_offline,

modelopt/torch/speculative/plugins/transformers.py

@@ -31,6 +31,7 @@
 
 import contextlib
 import copy
+import os
 from typing import Any
 
 import torch
@@ -49,6 +50,8 @@
 from transformers.utils import ModelOutput
 from transformers.utils.quantization_config import QuantizationMethod
 
+from modelopt.torch.utils import print_rank_0
+
 from ..eagle.conversion import EagleDMRegistry
 from ..eagle.eagle_model import EagleModel
 from ..eagle.utils import expand_mask, make_causal_mask
@@ -248,7 +251,16 @@ def __init__(self, config, decoder_layer_cls, bias=False):
             # Initialize the buffers to zero.
             # Their values depend on specific tokenzier and calibrate dataset, and should be set in training script.
             if config.draft_vocab_size < config.vocab_size:
-                self.register_buffer("d2t", torch.zeros(config.draft_vocab_size, dtype=torch.int64))
+                if config.draft_vocab_cache is not None and os.path.isfile(
+                    config.draft_vocab_cache
+                ):
+                    self.register_buffer("d2t", torch.load(config.draft_vocab_cache))
+                    print_rank_0(f"Loaded draft vocab cache from {config.draft_vocab_cache}.")
+                else:
+                    raise FileNotFoundError(
+                        f"Draft vocab cache file not found: {config.draft_vocab_cache}"
+                    )
+
             self.lm_head = nn.Linear(
                 config.hidden_size,
                 config.draft_vocab_size,
@@ -425,8 +437,11 @@ def _base_model_lm_head(self):
     @property
     def _base_llm_config(self):
         """Return the llm config for the base model, from LLM or VLM."""
-        # return self.config.llm_config if hasattr(self.config, "llm_config") else self.config
-        return self.config.text_config
+        return (
+            getattr(self.config, "text_config", None)
+            or getattr(self.config, "llm_config", None)
+            or self.config
+        )
 
     def _find_base_model_parts(self):
         """Find model parts from different models and set base_{part}_path attributes."""
@@ -574,13 +589,6 @@ def modify(
         ):
             self._set_default_aux_hidden_state_layers()
 
-        if self._base_llm_config.hidden_size != self.eagle_config.hidden_size:
-            raise ValueError(
-                "EAGLE module hidden size "
-                f"{self.eagle_config.hidden_size} must match base model hidden size "
-                f"{self._base_llm_config.hidden_size}!"
-            )
-
         # Freeze all parameters
         if self.eagle_freeze_base_model:
             for name, param in self.named_parameters():

modelopt/torch/utils/plugins/transformers_dataset.py

@@ -24,6 +24,8 @@
 from datasets import load_dataset
 from transformers.trainer_pt_utils import LabelSmoother
 
+from modelopt.torch.utils import print_rank_0
+
 REMOVE_THINK_CHAT_TEMPLATE = (
     "{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}"
 )
@@ -147,10 +149,15 @@ def __init__(
         else:
             self._post_process_chat_template()
 
+        self._post_process_tokenizer()
         if self.tokenizer.chat_template is None:
             raise ValueError("No valid chat template!")
 
     def _post_process_tokenizer(self):
+        if self.tokenizer.pad_token_id is None:
+            print_rank_0("The tokenizer has no pad_token_id, using eos_token_id instead.")
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
         if hasattr(self.tokenizer, "pad_token") and self.tokenizer.pad_token is None:
             if self.tokenizer.eos_token == "<|eot_id|>":  # nosec
                 self.tokenizer.pad_token = "<|end_of_text|>"  # nosec
@@ -264,7 +271,6 @@ def __call__(self, examples):
         for example in examples:
             messages = example.get("messages", None)
             if messages is None:
-                # print(example)
                 conversations = example.get("conversations", None)
                 if conversations is None:
                     raise ValueError(