diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index e981dc233010..accb27ef310a 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -334,21 +334,22 @@ def split_kv_b_proj(kv_b_proj: torch.Tensor, # Check if weights supports mark_consumed (ConsumableWeightsDict) can_mark_consumed = hasattr(weights, 'mark_consumed') + # The pretrained_config's num_nextn_predict_layers may have been expanded + # by ModelLoader._load_and_validate_config to match user max_draft_len. + # The original checkpoint MTP layer count (used for mod-indexing) is + # preserved as `_ckpt_num_nextn_predict_layers`. + ckpt_num_nextn_predict_layers = ( + getattr(self.config, '_ckpt_num_nextn_predict_layers', None) + or self.config.num_nextn_predict_layers) + def detect_shared_mtp_weights() -> bool: - # Detect if MTP layers share checkpoint weights (model requests more - # MTP layers than the checkpoint provides). In this case, multiple - # model MTP layers map to the same checkpoint layer via modulo, and - # mark_consumed must be skipped to avoid deleting weights that later - # MTP layers still need. - ckpt_nextn = self.config.num_nextn_predict_layers or 0 - spec_config = self.model_config.spec_config - if spec_config is not None and hasattr( - spec_config, 'spec_dec_mode' - ) and spec_config.spec_dec_mode.is_mtp_one_model(): - model_nextn = spec_config.num_nextn_predict_layers or 0 - else: - model_nextn = 0 - return model_nextn > ckpt_nextn > 0 + # Detect if MTP layers share checkpoint weights (model has more MTP + # layer instances than the checkpoint provides). In this case, + # multiple model MTP layers map to the same checkpoint layer via + # modulo, and mark_consumed must be skipped to avoid deleting + # weights that later MTP layers still need. + model_nextn = self.config.num_nextn_predict_layers or 0 + return model_nextn > (ckpt_num_nextn_predict_layers or 0) > 0 has_shared_mtp_weights = detect_shared_mtp_weights() @@ -368,7 +369,7 @@ def detect_shared_mtp_weights() -> bool: mtp_layer_idx = int( names[2]) - self.config.num_hidden_layers names[2] = str(mtp_layer_idx % - self.config.num_nextn_predict_layers + + ckpt_num_nextn_predict_layers + self.config.num_hidden_layers) name = '.'.join(names) mark_consumed = can_mark_consumed and not is_shared_mtp_layer @@ -1857,7 +1858,11 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): if model_config.spec_config is not None and model_config.spec_config.spec_dec_mode.is_mtp_one_model( ): model_nextn = self.config.num_nextn_predict_layers - ckpt_nextn = self.config.num_nextn_predict_layers + # When MTP layers share checkpoint weights (vanilla MTP with + # max_draft_len > ckpt MTP count), the original checkpoint count is + # preserved on pretrained_config; otherwise it equals num_nextn. + ckpt_nextn = (getattr(self.config, '_ckpt_num_nextn_predict_layers', + None) or self.config.num_nextn_predict_layers) self.num_hidden_layers = self.config.num_hidden_layers assert ckpt_nextn > 0, "There is not MTP modules in the checkpoint." if ckpt_nextn == 1 and not model_config.spec_config.use_mtp_vanilla: diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 5e98e730b6b8..c2b1b069b2c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -864,6 +864,45 @@ def _load_and_validate_config( if hasattr(config.pretrained_config, sub_config): getattr(config.pretrained_config, sub_config).num_hidden_layers = num_layers_override + + # Shared-weights vanilla MTP: build extra MTP layer instances beyond + # what the checkpoint provides (one ckpt MTP layer, multiple draft + # tokens, one KV cache per draft position) by sharing the single + # ckpt MTP layer's weights via mod-indexing in + # DeepseekV3WeightLoader. We expand + # pretrained_config.num_nextn_predict_layers to max_draft_len before + # model construction and preserve the original ckpt count as + # `_ckpt_num_nextn_predict_layers` for downstream mod-indexing. + # + # NOTE: this is a very special MTP mode that has not been used in + # any real-world workload to date; only DeepSeek has indicated they + # want to keep the path alive for their model. We therefore only + # support it on DeepSeek model_types for now. Other MTP-capable + # model families don't need this mode -- when their users request + # vanilla with max_draft_len > ckpt count, the natural + # `min(max_draft_len, ckpt_nextn)` clamp inside MTPForCausalLM + # silently caps the draft length to ckpt_nextn, which is the + # expected behavior for them. + _DEEPSEEK_MTP_MODEL_TYPES = {"deepseek_v3", "deepseek_v32"} + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + spec_config = self.spec_config + if (isinstance(spec_config, MTPDecodingConfig) + and spec_config.use_mtp_vanilla + and spec_config.max_draft_len is not None + and getattr(config.pretrained_config, 'model_type', + None) in _DEEPSEEK_MTP_MODEL_TYPES + and getattr(config.pretrained_config, + 'num_nextn_predict_layers', None)): + ckpt_nextn = config.pretrained_config.num_nextn_predict_layers + if spec_config.max_draft_len > ckpt_nextn: + config.pretrained_config._ckpt_num_nextn_predict_layers = ckpt_nextn + config.pretrained_config.num_nextn_predict_layers = \ + spec_config.max_draft_len + logger.warning( + f"MTP vanilla: expanding num_nextn_predict_layers from " + f"{ckpt_nextn} to {spec_config.max_draft_len} to match " + f"max_draft_len. Extra MTP layer instances will share " + f"checkpoint weights via mod-indexing.") return config def _call_load_weights(self, diff --git a/tests/integration/defs/accuracy/accuracy_core.py b/tests/integration/defs/accuracy/accuracy_core.py index c464fff23c6d..7e5a21a29207 100644 --- a/tests/integration/defs/accuracy/accuracy_core.py +++ b/tests/integration/defs/accuracy/accuracy_core.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import gc import json import math import os @@ -21,6 +22,7 @@ import pytest import scipy +import torch import yaml import tensorrt_llm.evaluate @@ -975,6 +977,19 @@ def setup_class(cls): yield logger.set_level(original_level) + @pytest.fixture(autouse=True) + def _cleanup_cuda_between_tests(self): + # Force Python GC + CUDA cache release after each test method. + # The LLM context manager's __exit__ schedules destruction of CUDA + # resources (streams, graph captures, KV cache pools), but objects in + # reference cycles aren't reclaimed until the next GC cycle. Without + # this teardown, a leftover CUDA stream/graph from a previous test can + # land in the next test's allocations and corrupt them, producing + # cross-test IMA reports that look like the current test crashed. + yield + gc.collect() + torch.cuda.empty_cache() + def get_accuracy_task(dataset_name: str): try: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index da27c539bd09..25ebe82099f4 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -45,7 +45,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mt accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6198785) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] SKIP (https://nvbugs/6071081) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6185146) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales[mtp=vanilla-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6195110) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6162115) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6162115) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp2pp2-mtp_nextn=2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6112497)