Skip to content

Commit ebc77ce

Browse files
committed
fix tests
Signed-off-by: Izzy Putterman <[email protected]>
1 parent 1c12aa9 commit ebc77ce

File tree

2 files changed

+4
-0
lines changed

2 files changed

+4
-0
lines changed

tensorrt_llm/_torch/pyexecutor/model_engine.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1027,6 +1027,9 @@ def init_meta_tensor(t: torch.Tensor):
10271027

10281028
elif load_format == LoadFormat.DUMMY:
10291029
initialize_dummy_weights(model)
1030+
if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
1031+
):
1032+
model.draft_model.load_weights_from_target_model(model)
10301033

10311034
elif load_format == LoadFormat.VISION_ONLY:
10321035
# Vision weights are already loaded within the model.

tests/unittest/_torch/speculative/test_eagle3.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,7 @@ def test_multi_eagle3(use_one_model: bool):
305305
max_batch_size=max_batch_size,
306306
kv_cache_config=kv_cache_config,
307307
enable_chunked_prefill=enable_chunked_prefill,
308+
load_format="dummy",
308309
)
309310

310311
spec_config = EagleDecodingConfig(

0 commit comments

Comments
 (0)