diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 4c5c01f40948..a0964b63dece 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1936,11 +1936,22 @@ def get_draft_model(model_config, draft_config, lm_head, model): class SpecDecOneEngineForCausalLM(DecoderModelForCausalLM[TModel, TConfig], Generic[TModel, TConfig]): - def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): + def __init__(self, + model: TModel, + model_config: ModelConfig[TConfig], + hidden_size: int | None = None, + vocab_size: int | None = None) -> None: + # Composite configs (e.g. vision-language wrappers) may not expose + # hidden_size/vocab_size at the top level; callers can pass the + # text-config values explicitly. + if hidden_size is None: + hidden_size = model_config.pretrained_config.hidden_size + if vocab_size is None: + vocab_size = model_config.pretrained_config.vocab_size super().__init__(model, config=model_config, - hidden_size=model_config.pretrained_config.hidden_size, - vocab_size=model_config.pretrained_config.vocab_size) + hidden_size=hidden_size, + vocab_size=vocab_size) self.draft_model = None self.draft_config = None self.spec_worker = None diff --git a/tests/unittest/_torch/modeling/test_modeling_speculative.py b/tests/unittest/_torch/modeling/test_modeling_speculative.py index 59ddd3cee128..87ba48e8ac5b 100644 --- a/tests/unittest/_torch/modeling/test_modeling_speculative.py +++ b/tests/unittest/_torch/modeling/test_modeling_speculative.py @@ -13,13 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for Eagle3ForCausalLM.apply_eagle3_fc fc_norm branch.""" +"""Unit tests for speculative modeling classes.""" + +from unittest.mock import MagicMock, patch import pytest import torch from torch import nn +from transformers import PretrainedConfig -from tensorrt_llm._torch.models.modeling_speculative import Eagle3ForCausalLM +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + Eagle3ForCausalLM, + SpecDecOneEngineForCausalLM, +) from tensorrt_llm._torch.modules.rms_norm import RMSNorm @@ -144,3 +151,73 @@ def test_apply_eagle3_fc_with_fc_norm(num_capture_layers): "fc_norm should apply per-chunk normalization which differs from " "whole-tensor normalization" ) + + +# --------------------------------------------------------------------------- +# SpecDecOneEngineForCausalLM: optional hidden_size / vocab_size +# --------------------------------------------------------------------------- + +_BASE_CLS = "tensorrt_llm._torch.models.modeling_utils.DecoderModelForCausalLM" + + +def _init_specdec_with_mocked_base(model_config, **kwargs): + """Instantiate SpecDecOneEngineForCausalLM with the base class stubbed out. + + DecoderModelForCausalLM is built on the PostInitCaller metaclass, which + invokes __post_init__/__pp_init__ right after __init__ returns. Those + hooks must be stubbed too: the mocked __init__ never sets the attributes + (model_config, lm_head, ...) they rely on. + + Returns the kwargs captured by the mocked base __init__. + """ + with ( + patch(f"{_BASE_CLS}.__init__", return_value=None) as mock_init, + patch(f"{_BASE_CLS}.__post_init__"), + patch(f"{_BASE_CLS}.__pp_init__"), + ): + SpecDecOneEngineForCausalLM(MagicMock(), model_config, **kwargs) + _, captured_kwargs = mock_init.call_args + return captured_kwargs + + +def test_specdec_one_engine_reads_from_pretrained_config() -> None: + """Default path: hidden_size/vocab_size come from pretrained_config.""" + hidden_size = 4096 + vocab_size = 32000 + model_config = ModelConfig( + pretrained_config=PretrainedConfig(hidden_size=hidden_size, vocab_size=vocab_size) + ) + + kwargs = _init_specdec_with_mocked_base(model_config) + assert kwargs["hidden_size"] == hidden_size + assert kwargs["vocab_size"] == vocab_size + + +def test_specdec_one_engine_accepts_explicit_sizes() -> None: + """Composite configs (e.g. VL wrappers) can pass sizes explicitly.""" + hidden_size = 8192 + vocab_size = 128256 + # Bare PretrainedConfig lacks hidden_size/vocab_size; the caller + # supplies them instead. + model_config = ModelConfig(pretrained_config=PretrainedConfig()) + + kwargs = _init_specdec_with_mocked_base( + model_config, hidden_size=hidden_size, vocab_size=vocab_size + ) + assert kwargs["hidden_size"] == hidden_size + assert kwargs["vocab_size"] == vocab_size + + +def test_specdec_one_engine_explicit_overrides_pretrained_config() -> None: + """Explicit args take precedence over pretrained_config when both present.""" + hidden_size = 2048 + vocab_size = 64000 + model_config = ModelConfig( + pretrained_config=PretrainedConfig(hidden_size=4096, vocab_size=32000) + ) + + kwargs = _init_specdec_with_mocked_base( + model_config, hidden_size=hidden_size, vocab_size=vocab_size + ) + assert kwargs["hidden_size"] == hidden_size + assert kwargs["vocab_size"] == vocab_size