@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 130 lines 3.91 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from types import SimpleNamespace
5
6 import pytest
7 import torch
8
9 from vllm.models.deepseek_v4.amd import dspark as dspark_module
10 from vllm.models.deepseek_v4.amd.dspark import DSparkDeepseekV4ForCausalLM
11
12
13 def _make_uninitialized_model(confidence_head):
14 model = DSparkDeepseekV4ForCausalLM.__new__(DSparkDeepseekV4ForCausalLM)
15 object.__setattr__(
16 model,
17 "model",
18 SimpleNamespace(confidence_head=confidence_head),
19 )
20 return model
21
22
23 def _prepare_loader_model(model, named_parameters):
24 object.__setattr__(
25 model,
26 "config",
27 SimpleNamespace(
28 n_routed_experts=1,
29 expert_dtype="fp4",
30 num_attention_heads=1,
31 ),
32 )
33 object.__setattr__(model, "named_parameters", lambda: named_parameters)
34
35
36 def _disable_distributed_loader_paths(monkeypatch: pytest.MonkeyPatch):
37 monkeypatch.setattr(
38 dspark_module,
39 "fused_moe_make_expert_params_mapping",
40 lambda *args, **kwargs: [],
41 )
42 monkeypatch.setattr(
43 dspark_module, "get_tensor_model_parallel_world_size", lambda: 1
44 )
45 monkeypatch.setattr(dspark_module, "get_tensor_model_parallel_rank", lambda: 0)
46
47
48 @pytest.mark.cpu_test
49 def test_deepseek_v4_rocm_dspark_maps_enabled_confidence_head():
50 model = _make_uninitialized_model(object())
51
52 assert (
53 model._remap_dspark_name("mtp.2.confidence_head.proj.weight")
54 == "model.confidence_head.proj.weight"
55 )
56
57
58 @pytest.mark.cpu_test
59 def test_deepseek_v4_rocm_dspark_skips_disabled_confidence_head():
60 model = _make_uninitialized_model(None)
61
62 assert model._remap_dspark_name("mtp.2.confidence_head.proj.weight") is None
63
64
65 @pytest.mark.cpu_test
66 def test_deepseek_v4_rocm_dspark_disables_unloaded_confidence_head(
67 monkeypatch: pytest.MonkeyPatch,
68 ):
69 model = _make_uninitialized_model(object())
70 _prepare_loader_model(model, [])
71 _disable_distributed_loader_paths(monkeypatch)
72
73 assert model.load_weights([]) == set()
74 assert model.model.confidence_head is None
75
76
77 @pytest.mark.cpu_test
78 def test_deepseek_v4_rocm_dspark_loads_complete_confidence_head(
79 monkeypatch: pytest.MonkeyPatch,
80 ):
81 class FakeParameter:
82 loaded_weight = None
83
84 def weight_loader(self, param, loaded_weight):
85 assert param is self
86 self.loaded_weight = loaded_weight
87
88 confidence_head = object()
89 parameter = FakeParameter()
90 model = _make_uninitialized_model(confidence_head)
91 _prepare_loader_model(
92 model,
93 [("model.confidence_head.proj.weight", parameter)],
94 )
95 _disable_distributed_loader_paths(monkeypatch)
96 loaded_weight = torch.tensor([[1.0, 2.0]])
97
98 loaded = model.load_weights([("mtp.2.confidence_head.proj.weight", loaded_weight)])
99
100 assert loaded == {"model.confidence_head.proj.weight"}
101 assert parameter.loaded_weight is loaded_weight
102 assert model.model.confidence_head is confidence_head
103
104
105 @pytest.mark.cpu_test
106 def test_deepseek_v4_rocm_dspark_confidence_is_probability():
107 class ConfidenceHead:
108 def __call__(self, head_hidden, markov_embed):
109 return (head_hidden[:, 0] + markov_embed[:, 0]).float()
110
111 model = _make_uninitialized_model(ConfidenceHead())
112 head_hidden = torch.tensor([[0.0], [1.0]], dtype=torch.bfloat16)
113 markov_embed = torch.tensor([[0.0], [-2.0]], dtype=torch.bfloat16)
114
115 confidence = model.compute_confidence(head_hidden, markov_embed)
116
117 torch.testing.assert_close(
118 confidence,
119 torch.sigmoid(torch.tensor([0.0, -1.0])),
120 )
121 assert torch.all((confidence >= 0) & (confidence <= 1))
122
123
124 @pytest.mark.cpu_test
125 def test_deepseek_v4_rocm_dspark_confidence_requires_a_head():
126 model = _make_uninitialized_model(None)
127 empty = torch.zeros((1, 1))
128
129 with pytest.raises(RuntimeError, match="confidence_head"):
130 model.compute_confidence(empty, empty)