Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / mtp.py: 100%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""DeepSeek-V3-style MTP layers and teacher-forced training orchestration. 

16 

17Fusion uses the released-weight convention [normalized embedding, normalized 

18hidden], as in Transformers MtpLayer. The training loop uses ground-truth future 

19tokens, not the candidate-token generation loop in Transformers MtpModel. 

20Model adapters own decoder construction, attention/position metadata, distributed 

21layouts and checkpoint key mapping. No model-specific routing policy lives here. 

22""" 

23 

24from __future__ import annotations 

25 

26import math 

27from collections.abc import Callable, Mapping, Sequence 

28from dataclasses import dataclass 

29from functools import partial 

30from typing import Any 

31 

32import torch 

33from torch import nn 

34 

35__all__ = ["MTPLayer", "MultiTokenPrediction", "MTPOutput", "build_mtp"] 

36 

37 

38class _MTPRMSNorm(nn.Module): 

39 """Match the HF DeepSeek normalization cast order, including FP32 weights.""" 

40 

41 def __init__(self, hidden_size: int, eps: float) -> None: 

42 """Create a direct scale with the requested normalization epsilon.""" 

43 super().__init__() 

44 self.weight = nn.Parameter(torch.ones(hidden_size)) 

45 self.eps = eps 

46 

47 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 

48 """Normalize in FP32 and preserve the reference multiplication order.""" 

49 input_dtype = hidden_states.dtype 

50 states = hidden_states.float() 

51 states = states * torch.rsqrt(states.square().mean(-1, keepdim=True) + self.eps) 

52 return self.weight * states.to(input_dtype) 

53 

54 

55class MTPLayer(nn.Module): 

56 """Fuse a future embedding with the previous depth's hidden state. 

57 

58 Args: 

59 hidden_size: Width of both inputs and the decoder output. 

60 decoder: Model-owned causal Transformer block returning a Tensor. 

61 rms_norm_eps: Epsilon for the default DeepSeek-style RMSNorms. 

62 norm_factory: Optional factory accepting hidden_size. It must return a 

63 fresh module for each of enorm, hnorm and post_norm. 

64 

65 Note: 

66 All three norms have independent parameters. The returned post-normalized 

67 state is used by both the shared head and the next MTP depth. A model with 

68 different fusion/normalization semantics should supply its own layer to 

69 MultiTokenPrediction, rather than changing this released-weight convention. 

70 """ 

71 

72 def __init__(self, hidden_size: int, decoder: nn.Module, *, rms_norm_eps: float = 1e-6, 

73 norm_factory: Callable[[int], nn.Module] | None = None) -> None: 

74 """Build independent fusion and output norms around the supplied decoder.""" 

75 super().__init__() 

76 if isinstance(hidden_size, bool) or not isinstance(hidden_size, int) or hidden_size <= 0: 

77 raise ValueError("hidden_size must be a positive integer") 

78 if not math.isfinite(rms_norm_eps) or rms_norm_eps <= 0: 

79 raise ValueError("rms_norm_eps must be positive and finite") 

80 factory = norm_factory or partial(_MTPRMSNorm, eps=rms_norm_eps) 

81 self.enorm = factory(hidden_size) 

82 self.hnorm = factory(hidden_size) 

83 self.eh_proj = nn.Linear(2 * hidden_size, hidden_size, bias=False) 

84 self.mtp_block = decoder 

85 self.post_norm = factory(hidden_size) 

86 

87 def forward(self, hidden_states: torch.Tensor, inputs_embeds: torch.Tensor, 

88 **decoder_kwargs: Any) -> torch.Tensor: 

89 """Return [B, S, H] states; kwargs must describe this depth's sequence.""" 

90 if hidden_states.shape != inputs_embeds.shape: 

91 raise ValueError("MTP hidden_states and inputs_embeds must have identical shapes") 

92 projection_input = torch.cat((self.enorm(inputs_embeds), self.hnorm(hidden_states)), dim=-1) 

93 return self.post_norm(self.mtp_block(self.eh_proj(projection_input), **decoder_kwargs)) 

94 

95 

96@dataclass 

97class MTPOutput: 

98 """Depth k (one-based) has S-k positions; logits are empty if head is None.""" 

99 

100 hidden_states: tuple[torch.Tensor, ...] 

101 logits: tuple[torch.Tensor, ...] 

102 

103 

104class MultiTokenPrediction(nn.Module): 

105 """Run independent MTP depths sequentially with ground-truth future tokens. 

106 

107 Args: 

108 layers: Independent layers with forward(hidden_states, inputs_embeds, 

109 **decoder_kwargs) -> Tensor. Normally MTPLayer instances; 

110 adapters may supply another layer with the same contract. 

111 

112 Note: 

113 This loop consumes complete, unsharded rows. It does not implement CP halo 

114 exchange, packed-document boundaries, generation/KV caching or automatic 

115 Trainer integration. Those belong to the model adapter. A distributed 

116 adapter may also call individual layers with already-aligned inputs. 

117 """ 

118 

119 def __init__(self, layers: Sequence[nn.Module]) -> None: 

120 """Register the prediction layers without copying external shared weights.""" 

121 super().__init__() 

122 if not layers: 

123 raise ValueError("MTP requires at least one prediction layer") 

124 self.layers = nn.ModuleList(layers) 

125 

126 def forward(self, hidden_states: torch.Tensor, input_ids: torch.Tensor, *, 

127 embedding: nn.Module, head: nn.Module | None = None, 

128 decoder_kwargs: Sequence[Mapping[str, Any]] | None = None) -> MTPOutput: 

129 """Predict future tokens without detaching the main-model computation. 

130 

131 Args: 

132 hidden_states: Main trunk output BEFORE its final norm, [B, S, H]. 

133 input_ids: Ground-truth tokens [B, S], not next-token labels. 

134 embedding: Main model's embedding, reused without registering a copy. 

135 head: Main model's output projection, or None for external chunk loss. 

136 decoder_kwargs: One mapping per depth. For depth k the caller must 

137 provide attention masks/positions for S-k query positions. None 

138 is suitable only for decoders that construct causal metadata. 

139 

140 Returns: 

141 Per-depth states and optional logits. At depth k, position i fuses 

142 the prior state at i with token i+k, and predicts token i+k+1. 

143 """ 

144 if hidden_states.ndim != 3 or input_ids.ndim != 2 or hidden_states.shape[:2] != input_ids.shape: 

145 raise ValueError("MTP expects aligned hidden_states [B, S, H] and input_ids [B, S]") 

146 if input_ids.shape[0] == 0 or input_ids.shape[1] <= len(self.layers): 

147 raise ValueError("MTP needs a nonempty batch and sequence length greater than its depth") 

148 if decoder_kwargs is not None and len(decoder_kwargs) != len(self.layers): 

149 raise ValueError("decoder_kwargs must contain one mapping per MTP depth") 

150 states, logits = [], [] 

151 for index, layer in enumerate(self.layers): 

152 kwargs = {} if decoder_kwargs is None else decoder_kwargs[index] 

153 future_embeddings = embedding(input_ids[:, index + 1:]) 

154 hidden_states = layer(hidden_states[:, :-1], future_embeddings, **kwargs) 

155 states.append(hidden_states) 

156 if head is not None: 

157 logits.append(head(hidden_states)) 

158 return MTPOutput(tuple(states), tuple(logits)) 

159 

160 

161def build_mtp( 

162 mtp_num_layers: int | None = None, *, layer_factory: Callable[[int], nn.Module], 

163) -> MultiTokenPrediction | None: 

164 """Build the configured MTP depths, or return None when disabled. 

165 

166 Args: 

167 mtp_num_layers: Megatron-style depth count. None or zero disables MTP; 

168 a positive integer creates that many independent prediction layers. 

169 layer_factory: Model-owned factory called with each zero-based depth. 

170 It must return a fresh MTP layer, including its causal decoder. 

171 

172 Returns: 

173 Registered MTP stack, or None without invoking the factory when disabled. 

174 

175 Note: 

176 Model builders pass their YAML mtp_num_layers value here, attach the 

177 returned module, and include its output in forward/loss when enabled. 

178 This factory does not attach MTP to arbitrary models automatically. 

179 """ 

180 if mtp_num_layers is None: 

181 return None 

182 if isinstance(mtp_num_layers, bool) or not isinstance(mtp_num_layers, int) or mtp_num_layers < 0: 

183 raise ValueError("mtp_num_layers must be None or a nonnegative integer") 

184 if mtp_num_layers == 0: 

185 return None 

186 return MultiTokenPrediction([layer_factory(depth) for depth in range(mtp_num_layers)])