| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | from unittest.mock import Mock |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | from vllm.models.deepseek_v32.nvidia import model as deepseek_v32_model |
| 11 | from vllm.models.deepseek_v32.nvidia import mtp as deepseek_v32_mtp |
| 12 | |
| 13 | |
| 14 | class _IdentityNorm(nn.Module): |
| 15 | def __init__(self, hidden_size: int = 2) -> None: |
| 16 | super().__init__() |
| 17 | self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False) |
| 18 | self.variance_epsilon = 1e-5 |
| 19 | |
| 20 | def forward( |
| 21 | self, |
| 22 | hidden_states: torch.Tensor, |
| 23 | residual: torch.Tensor | None = None, |
| 24 | ): |
| 25 | if residual is None: |
| 26 | return hidden_states |
| 27 | return hidden_states, residual |
| 28 | |
| 29 | |
| 30 | class _RecordingModule(nn.Module): |
| 31 | def __init__(self) -> None: |
| 32 | super().__init__() |
| 33 | self.num_tokens = 0 |
| 34 | |
| 35 | def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: |
| 36 | self.num_tokens = hidden_states.shape[0] |
| 37 | return hidden_states |
| 38 | |
| 39 | |
| 40 | class _RecordingProjection(nn.Module): |
| 41 | def __init__(self) -> None: |
| 42 | super().__init__() |
| 43 | self.weight = nn.Parameter(torch.ones(2, 4), requires_grad=False) |
| 44 | self.num_tokens = 0 |
| 45 | |
| 46 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 47 | self.num_tokens = hidden_states.shape[0] |
| 48 | return hidden_states[:, :2] |
| 49 | |
| 50 | |
| 51 | class _SequenceParallelMTPBlock: |
| 52 | use_sequence_parallel = True |
| 53 | |
| 54 | def __call__( |
| 55 | self, |
| 56 | *, |
| 57 | positions: torch.Tensor, |
| 58 | hidden_states: torch.Tensor, |
| 59 | residual: torch.Tensor | None, |
| 60 | ): |
| 61 | assert residual is None |
| 62 | return hidden_states * 2, hidden_states * 3 |
| 63 | |
| 64 | |
| 65 | def _mock_sequence_parallel_collectives(monkeypatch, module): |
| 66 | monkeypatch.setattr( |
| 67 | module, |
| 68 | "sp_reduce_scatter", |
| 69 | lambda tensor: tensor.chunk(2, dim=0)[0], |
| 70 | raising=False, |
| 71 | ) |
| 72 | monkeypatch.setattr( |
| 73 | module, |
| 74 | "sp_shard", |
| 75 | lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], |
| 76 | raising=False, |
| 77 | ) |
| 78 | monkeypatch.setattr( |
| 79 | module, |
| 80 | "sp_all_gather", |
| 81 | lambda tensor: torch.cat([tensor, tensor], dim=0), |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | def test_decoder_layer_keeps_dense_states_sequence_sharded(monkeypatch): |
| 86 | layer = object.__new__(deepseek_v32_model.DeepseekV32DecoderLayer) |
| 87 | nn.Module.__init__(layer) |
| 88 | layer.use_sequence_parallel = True |
| 89 | layer.input_layernorm = _IdentityNorm() |
| 90 | layer.post_attention_layernorm = _IdentityNorm() |
| 91 | layer.self_attn = _RecordingModule() |
| 92 | layer.mlp = _RecordingModule() |
| 93 | |
| 94 | _mock_sequence_parallel_collectives(monkeypatch, deepseek_v32_model) |
| 95 | |
| 96 | positions = torch.arange(3) |
| 97 | full_hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) |
| 98 | hidden_states = deepseek_v32_model.sp_shard(full_hidden_states) |
| 99 | hidden_states, residual = layer(positions, hidden_states, residual=None) |
| 100 | |
| 101 | assert hidden_states.shape == residual.shape == (2, 2) |
| 102 | assert layer.self_attn.num_tokens == 3 |
| 103 | assert layer.mlp.num_tokens == 2 |
| 104 | |
| 105 | hidden_states, residual = layer(positions, hidden_states, residual) |
| 106 | |
| 107 | assert hidden_states.shape == residual.shape == (2, 2) |
| 108 | assert layer.self_attn.num_tokens == 3 |
| 109 | assert layer.mlp.num_tokens == 2 |
| 110 | |
| 111 | |
| 112 | def test_mtp_projects_sequence_shard_and_restores_full_output(monkeypatch): |
| 113 | layer = object.__new__(deepseek_v32_mtp.DeepseekV32MultiTokenPredictorLayer) |
| 114 | nn.Module.__init__(layer) |
| 115 | layer.enorm = _IdentityNorm() |
| 116 | layer.hnorm = _IdentityNorm() |
| 117 | layer.eh_proj = _RecordingProjection() |
| 118 | layer._eh_plan = None |
| 119 | object.__setattr__(layer, "mtp_block", _SequenceParallelMTPBlock()) |
| 120 | norm = Mock( |
| 121 | side_effect=lambda hidden_states, residual: (hidden_states + residual, None) |
| 122 | ) |
| 123 | object.__setattr__(layer, "shared_head", SimpleNamespace(norm=norm)) |
| 124 | |
| 125 | monkeypatch.setattr( |
| 126 | deepseek_v32_mtp, |
| 127 | "fused_eh_norm", |
| 128 | lambda positions, inputs_embeds, previous_hidden_states, *args: torch.cat( |
| 129 | [inputs_embeds, previous_hidden_states], dim=-1 |
| 130 | ), |
| 131 | ) |
| 132 | monkeypatch.setattr(deepseek_v32_mtp, "run_glm52_plan", lambda *args: None) |
| 133 | _mock_sequence_parallel_collectives(monkeypatch, deepseek_v32_mtp) |
| 134 | |
| 135 | inputs_embeds = torch.arange(6, dtype=torch.float32).view(3, 2) |
| 136 | hidden_states, recycled_hidden_states = layer( |
| 137 | input_ids=torch.zeros(3, dtype=torch.long), |
| 138 | positions=torch.arange(3), |
| 139 | previous_hidden_states=torch.zeros_like(inputs_embeds), |
| 140 | inputs_embeds=inputs_embeds, |
| 141 | ) |
| 142 | |
| 143 | sharded_states = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, 1))[:2] |
| 144 | expected = torch.cat([sharded_states * 5, sharded_states * 5])[:3] |
| 145 | assert layer.eh_proj.num_tokens == 2 |
| 146 | torch.testing.assert_close(hidden_states, expected) |
| 147 | torch.testing.assert_close(recycled_hidden_states, expected) |
| 148 | norm.assert_called_once() |