diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index b0b195750..655775447 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -972,6 +972,8 @@ class NemotronHForCausalLM(nn.Module): continue is_expert_weight = True name_mapped = name.replace(weight_name, param_name) + if name_mapped not in params_dict: + continue param = params_dict[name_mapped] param.weight_loader( param, diff --git a/test/registered/unit/models/test_nemotron_h_weight_loading.py b/test/registered/unit/models/test_nemotron_h_weight_loading.py new file mode 100644 index 000000000..8a5e30a4e --- /dev/null +++ b/test/registered/unit/models/test_nemotron_h_weight_loading.py @@ -0,0 +1,87 @@ +""" +Unit tests for NemotronHForCausalLM.load_weights. + +Regression test for Nemotron-H expert scale checkpoint tensors that map to +parameters absent from the current runtime model. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=4, suite="stage-a-test-cpu") + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.models.nemotron_h import NemotronHForCausalLM + + +class _FakePPGroup: + is_first_rank = True + is_last_rank = True + + +class _FakeParam: + def __init__(self): + self.loaded = None + + def weight_loader( + self, param, loaded_weight, name, *, shard_id=None, expert_id=None + ): + self.loaded = (param, loaded_weight, name, shard_id, expert_id) + + +class TestNemotronHWeightLoading(unittest.TestCase): + def _make_minimal_model(self, named_parameters=()): + model = object.__new__(NemotronHForCausalLM) + model.config = SimpleNamespace(n_routed_experts=2) + model.model = SimpleNamespace() + model.pp_group = _FakePPGroup() + model.remap_prefix = {} + model.remap_substr = {} + model.stacked_params_mapping = [] + model.named_parameters = lambda: iter(named_parameters) + return model + + def test_expert_input_scale_without_target_parameter_is_skipped(self): + """Expert scale weights absent from params_dict should not raise KeyError.""" + model = self._make_minimal_model() + weights = [ + ( + "model.layers.1.mixer.experts.0.down_proj.input_scale", + torch.ones(1), + ) + ] + + model.load_weights(weights) + + def test_expert_weight_with_target_parameter_is_loaded(self): + param = _FakeParam() + model = self._make_minimal_model( + [("model.layers.1.mixer.experts.w2_weight", param)] + ) + loaded_weight = torch.ones(1) + weights = [ + ( + "model.layers.1.mixer.experts.0.down_proj.weight", + loaded_weight, + ) + ] + + model.load_weights(weights) + + self.assertEqual( + param.loaded, + ( + param, + loaded_weight, + "model.layers.1.mixer.experts.w2_weight", + "w2", + 0, + ), + ) + + +if __name__ == "__main__": + unittest.main()