Use separate weight mapper for draft model

Signed-off-by: Anurag Mukkara <[email protected]>

Author: amukkara

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -554,11 +554,11 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                          hidden_size=model_config.pretrained_config.hidden_size,
                          vocab_size=model_config.pretrained_config.vocab_size)
         self.draft_model = None
+        self.draft_config = None
         spec_config = getattr(model_config, 'spec_config', None)
         if spec_config and spec_config.spec_dec_mode.use_one_engine():
-            draft_config = None
             if spec_config.spec_dec_mode.is_eagle3_one_model():
-                draft_config = ModelConfig.from_pretrained(
+                self.draft_config = ModelConfig.from_pretrained(
                     model_config.spec_config.speculative_model_dir,
                     trust_remote_code=True,
                     attn_backend=model_config.attn_backend,
@@ -567,17 +567,17 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                     spec_config=model_config.spec_config,
                     max_num_tokens=model_config.max_num_tokens,
                     moe_max_num_tokens=model_config.moe_max_num_tokens)
-                draft_config.quant_config.kv_cache_quant_algo = \
+                self.draft_config.quant_config.kv_cache_quant_algo = \
                 model_config.quant_config.kv_cache_quant_algo
 
-            self.draft_model = get_draft_model(model_config, draft_config,
+            self.draft_model = get_draft_model(model_config, self.draft_config,
                                                self.lm_head, self.model)
             self.spec_worker = get_spec_worker(model_config.spec_config,
                                                model_config,
                                                model_config.mapping)
 
-            if draft_config is not None:
-                for key, value in draft_config.extra_attrs.items():
+            if self.draft_config is not None:
+                for key, value in self.draft_config.extra_attrs.items():
                     assert key in ('attn_layers', 'mla_layers')
                     assert key in model_config.extra_attrs
                     model_config.extra_attrs[key].update(value)

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -278,8 +278,10 @@ def init_meta_tensor(t: torch.Tensor):
                 ):
                     weights = checkpoint_loader.load_weights(
                         self.spec_config.speculative_model_dir)
+                    draft_weight_mapper = checkpoint_loader.get_initialized_weight_mapper(
+                        model.draft_model, model.draft_config)
                     self._call_load_weights(model.load_draft_weights, weights,
-                                            self.weight_mapper)
+                                            draft_weight_mapper)
 
             elif load_format == LoadFormat.DUMMY:
                 self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper(