From aced3a1a5a11e42d699a54d9a7e15aa1fb725043 Mon Sep 17 00:00:00 2001 From: Bowen Fu <5812640+BowenFu@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:36:45 +0000 Subject: [PATCH 1/2] [TRTLLM-6934] Re-enable Nemotron NAS HF parity test Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_nas.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py index 01bff76114b7..ce83bc32490a 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 unittest from copy import deepcopy from dataclasses import dataclass @@ -360,7 +375,6 @@ def test_nemotron_nas_sanity(self): ], lambda testcase_func, param_num, param: f"{testcase_func.__name__}[{param.args[0]}]") @torch.no_grad() - @unittest.skip("https://nvbugspro.nvidia.com/bug/5439817") def test_nemotron_nas_allclose_to_hf(self, scenario: Scenario) -> None: """ Compare output to HF From 46c732addd3ebb48fd4eb208ce6d8a67f6f3b291 Mon Sep 17 00:00:00 2001 From: Bowen Fu <5812640+BowenFu@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:55:28 +0000 Subject: [PATCH 2/2] [https://nvbugs/5439817][fix] Make Nemotron NAS HF parity cache work with transformers>=5 The Nemotron NAS HF parity test builds its reference cache from the checkpoint's remote `VariableCache`, whose `__init__` assigns `self.max_batch_size` and `self.max_cache_len`. `transformers>=5` turned both into read-only properties on `Cache`, so those assignments now raise `AttributeError: property 'max_batch_size' of 'VariableCache' object has no setter` before any comparison runs. Instantiate the remote cache through a small helper that shadows the read-only property names with plain class attributes on a subclass, which restores writable instance attributes and is a no-op on `transformers` versions that never defined the properties. Reproduced and verified locally against the checkpoint's remote code with transformers 5.5.4: construction plus a context and a generation forward through the cache update path all succeed with the helper and fail without it. Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_nas.py | 27 ++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py index ce83bc32490a..5625ad26a542 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py @@ -38,6 +38,26 @@ # Setup NEED_SETUP_CACHE_CLASSES_MAPPING to an empty dict for modeling_nemotron_nas.py transformers.generation.utils.NEED_SETUP_CACHE_CLASSES_MAPPING = dict() + +def instantiate_variable_cache(cache_cls: type, **kwargs: Any) -> Any: + """Instantiate the checkpoint's remote ``VariableCache``. + + Its ``__init__`` assigns ``self.max_batch_size`` and ``self.max_cache_len``, + which ``transformers>=5`` exposes as read-only properties on ``Cache``, so the + assignments raise ``AttributeError``. Shadowing those names with plain class + attributes on a subclass makes them writable instance attributes again, and + is a no-op on ``transformers`` versions that never defined the properties. + """ + shadowed = { + name: None + for name in ("max_batch_size", "max_cache_len") + if isinstance(getattr(cache_cls, name, None), property) + } + if shadowed: + cache_cls = type(f"Compat{cache_cls.__name__}", (cache_cls, ), shadowed) + return cache_cls(**kwargs) + + NEMOTRON_NAS_MINI_CONFIG = { "architectures": ["DeciLMForCausalLM"], "attention_bias": @@ -479,9 +499,10 @@ def test_nemotron_nas_allclose_to_hf(self, scenario: Scenario) -> None: position_ids = [torch.arange(0, input_ids.size(-1))] position_ids = torch.cat(position_ids).unsqueeze(0).cuda() # And, lastly, this is the simplest way of creating a Cache that `hf_nemotron_nas` will accept - past_key_values = VariableCache(config=nemotron_nas_config, - dtype=dtype, - batch_size=1) + past_key_values = instantiate_variable_cache(VariableCache, + config=nemotron_nas_config, + dtype=dtype, + batch_size=1) with torch.inference_mode(): attn_metadata.prepare() logits = nemotron_nas.forward(input_ids=input_ids,