refactor: interface

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

Author: h-guo18

examples/speculative_decoding/eagle_utils.py

@@ -74,10 +74,10 @@ def make_speculative_data_module(
             chat_template = f.read()
         print_rank_0(f"Loaded chat template from {template_path}")
 
-    streaming_url = getattr(data_args, "streaming_server_url", None)
-    if streaming_url is not None:
-        # Streaming: trainer is a client of a running vllm serve
-        print_rank_0(f"Streaming hidden states from {streaming_url}")
+    mode = getattr(data_args, "mode", "online")
+    if mode == "streaming":
+        # Trainer is an HTTP client of a running vllm serve; samples stream in lazily.
+        print_rank_0(f"Streaming hidden states from {data_args.streaming_server_url}")
         from modelopt.torch.speculative.plugins.hf_streaming_dataset import (
             StreamingHiddenStatesDataset,
         )
@@ -89,15 +89,15 @@ def make_speculative_data_module(
         train_dataset = StreamingHiddenStatesDataset(
             entries=entries,
             tokenizer=tokenizer,
-            server_url=streaming_url,
+            server_url=data_args.streaming_server_url,
             model=data_args.streaming_model_name,
             max_seq_len=data_args.streaming_max_seq_len,
             answer_only_loss=answer_only_loss,
             prefetch=data_args.streaming_prefetch,
         )
         data_collator = EagleOfflineDataCollator(train_len=train_len)
 
-    elif data_args.offline_data_path is None:
+    elif mode == "online":
         train_dataset = ShardedDataset("json", data_files=data_args.data_path)
 
         if not data_args.vlm_processor:

examples/speculative_decoding/main.py

@@ -152,10 +152,8 @@ def train():
     training_args = HfTrainingArguments(**recipe.training.model_dump())
     init_distributed_env(training_args)
 
-    if not recipe.data.data_path and not recipe.data.offline_data_path:
-        raise ValueError(
-            "Either data.data_path or data.offline_data_path must be set in the config."
-        )
+    if recipe.data.mode in ("online", "streaming") and not recipe.data.data_path:
+        raise ValueError(f"data.mode={recipe.data.mode!r} requires data.data_path.")
     if training_args.cp_size > 1:
         patch_ring_attention_for_ttt()
         # Specific patch to accelerate 1.12.0. Removable after move to 1.13.0
@@ -174,10 +172,7 @@ def train():
 
     checkpoint = training_args.resume_from_checkpoint or last_checkpoint
 
-    use_offline_training = (
-        recipe.data.offline_data_path is not None
-        or recipe.data.streaming_server_url is not None
-    )
+    use_offline_training = recipe.data.mode != "online"
 
     if checkpoint:
         with patch_transformers5_params_loading():

modelopt/recipe/config.py

@@ -156,9 +156,7 @@ class ModelOptEagleRecipe(ModelOptSpeculativeRecipeBase):
 
     @model_validator(mode="after")
     def _derive_eagle_offline(self) -> ModelOptEagleRecipe:
-        self.eagle.eagle_offline = (
-            self.data.offline_data_path is not None or self.data.streaming_server_url is not None
-        )
+        self.eagle.eagle_offline = self.data.mode != "online"
         return self
 
     @model_validator(mode="after")

modelopt/torch/speculative/config.py

@@ -139,8 +139,9 @@ class EagleConfig(ModeloptBaseConfig):
     eagle_offline: bool = ModeloptField(
         default=False,
         description=(
-            "Whether to use detached Eagle. Derived by ModelOptEagleRecipe from "
-            "data.offline_data_path; not user-configurable."
+            "Whether the Eagle module consumes pre-computed hidden states (offline or streaming) "
+            "instead of running the base model in-process. Derived by ModelOptEagleRecipe from "
+            "``data.mode``; not user-configurable."
         ),
     )
 

modelopt/torch/speculative/plugins/hf_training_args.py

@@ -29,7 +29,9 @@
 
 from __future__ import annotations
 
-from pydantic import BaseModel, ConfigDict, field_validator
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 
 
 class ModelArguments(BaseModel):
@@ -47,6 +49,7 @@ class DataArguments(BaseModel):
 
     model_config = ConfigDict(extra="forbid")
 
+    mode: Literal["online", "offline", "streaming"] = "online"
     data_path: str | None = None
     offline_data_path: str | None = None
     lazy_preprocess: bool = True
@@ -67,6 +70,25 @@ def _check_sample_size(cls, v: int) -> int:
             raise ValueError("sample_size must be -1 (use all samples) or a positive integer")
         return v
 
+    @model_validator(mode="after")
+    def _check_mode_requirements(self) -> DataArguments:
+        # Backward-compat: if ``mode`` is left at its default ("online") but a mode-specific
+        # field is set, promote to the corresponding mode. Explicit ``mode`` always wins.
+        if self.mode == "online":
+            if self.offline_data_path is not None:
+                self.mode = "offline"
+            elif self.streaming_server_url is not None:
+                self.mode = "streaming"
+        if self.mode == "offline" and not self.offline_data_path:
+            raise ValueError("data.mode='offline' requires data.offline_data_path")
+        if self.mode == "streaming" and not (
+            self.streaming_server_url and self.streaming_model_name
+        ):
+            raise ValueError(
+                "data.mode='streaming' requires data.streaming_server_url and data.streaming_model_name"
+            )
+        return self
+
 
 class TrainingArguments(BaseModel):
     """Speculative-decoding extensions on top of ``transformers.TrainingArguments``.

modelopt_recipes/general/speculative_decoding/eagle3.yaml

@@ -13,6 +13,9 @@ model:
 
 # maps to DataArguments (main.py)
 data:
+  # online | offline | streaming. Set offline_data_path for offline; set
+  # streaming_server_url + streaming_model_name for streaming.
+  mode: online
   data_path: input_conversations/train.jsonl
   offline_data_path:
   draft_vocab_cache:
@@ -48,7 +51,7 @@ training:
 
 # maps to EagleConfig (modelopt/torch/speculative/config.py).
 eagle:
-  # eagle_offline is derived from data.offline_data_path; do not set here.
+  # eagle_offline is derived from data.mode; do not set here.
   eagle_decoder_type: llama
   eagle_ttt_steps: 3
   eagle_mix_hidden_states: false