Last time we discussed the structure of the Deepseek model. Today, I will share the inference acceleration tool in V3, the multi-token prediction (MTP) module, which is somewhat similar to speculative decoding, but the specific code has not been open-sourced. Speculative decoding was introduced in a series of articles on inference optimization, which essentially changes the autoregressive generation characteristics of large models. The idea in the paper by Meta published at ICML 2024 titled “Better & Faster Large Language Models via Multi-token Prediction” [1] is somewhat similar to this. I have included their pseudocode for reference.
This is very similar to Medusa, which adds several output heads to predict the next token. I haven’t introduced Medusa to everyone yet; I will write a separate article about it later, so we can ignore it for now. The model consists of a main part and several heads, each with its own RMSnorm, one Linear layer, and one Transformer Block. The predictions from each head are interdependent; after the Main Model computes, the results are passed to the subsequent heads, which is the MTP module. During training, the loss includes the prediction loss from the subsequent heads, enhancing the model’s understanding of long texts and complex reasoning, ensuring the accuracy of the model’s subsequent token predictions. During inference, the computational overhead of each head is much smaller than that of the Main Model, about 1/L, allowing for the quick retrieval of multiple consecutive tokens. It enables the generation of multiple tokens in a single inference. Of course, during inference, the Main Model can also be used alone for autoregressive generation results.
Talk is cheap, show me the code:
# Define the Transformer model for multi-token prediction
class MultiTokenTransformer(nn.Module):
def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, dim_feedforward,
max_seq_length, num_future_tokens, dropout=0.1):
super(MultiTokenTransformer, self).__init__()
self.d_model = d_model
self.num_future_tokens = num_future_tokens # Number of future tokens to predict
# Word embedding layer
self.embedding = nn.Embedding(vocab_size, d_model)
# Positional encoding
self.positional_encoding = PositionalEncoding(d_model, dropout, max_seq_length)
# Shared Transformer encoder (backbone)
encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead,
dim_feedforward=dim_feedforward,
dropout=dropout)
self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
# Define multiple independent output heads
self.output_heads = nn.ModuleList([
nn.Sequential(
nn.Linear(d_model, dim_feedforward),
nn.ReLU(),
nn.Linear(dim_feedforward, vocab_size)
) for _ in range(num_future_tokens)
])
# Initialize parameters
self._reset_parameters()
def _reset_parameters(self):
nn.init.xavier_uniform_(self.embedding.weight)
for head in self.output_heads:
for layer in head:
if isinstance(layer, nn.Linear):
nn.init.xavier_uniform_(layer.weight)
if layer.bias is not None:
nn.init.zeros_(layer.bias)
def encode(self, src):
"""
Encode the input sequence and return the shared representation (memory)
Input:
src: [seq_length, batch_size]
Output:
memory: [seq_length, batch_size, d_model]
"""
# Word embedding and add positional encoding
src_emb = self.embedding(src) * (self.d_model ** 0.5)
src_emb = self.positional_encoding(src_emb)
# Pass through the Transformer encoder (shared backbone)
memory = self.transformer_encoder(src_emb)
return memory
def forward_head(self, memory, head_index):
"""
Compute the output of the specified output head
Input:
memory: [seq_length, batch_size, d_model]
head_index: Index of the output head
Output:
output: [seq_length, batch_size, vocab_size]
"""
output = self.output_heads[head_index](memory)
return output
def forward(self, src):
"""
Complete forward propagation, returning results from all output heads
Input:
src: [seq_length, batch_size]
Output:
outputs: {'logits_head_0': ..., 'logits_head_1': ..., ...}
"""
memory = self.encode(src) # [seq_length, batch_size, d_model]
outputs = {}
for i in range(self.num_future_tokens):
logits = self.forward_head(memory, i) # [seq_length, batch_size, vocab_size]
outputs[f'logits_head_{i}'] = logits
return outputs
def generate_text_with_self_speculative_decoding(model, input_seq, max_length, tokenizer, prob_threshold=0.9):
"""
Generate text using self-speculative decoding
Parameters:
- model: Trained multi-token prediction model
- input_seq: Initial input sequence, shape [seq_length, batch_size]
- max_length: Maximum length of generated text
- tokenizer: Tokenizer for decoding generated sequences
- prob_threshold: Probability threshold; tokens exceeding this threshold are accepted
"""
model.eval()
generated = input_seq.clone() # Clone the input sequence to prevent modifying the original sequence
seq_length, batch_size = generated.size()
with torch.no_grad():
while generated.size(0) < max_length:
# Get model outputs
outputs = model(generated)
# Get speculative tokens
speculative_tokens = []
for i in range(model.num_future_tokens):
logits = outputs[f'logits_head_{i}'] # [seq_length, batch_size, vocab_size]
# Take the logits from the last time step
last_logits = logits[-1, :, :] # [batch_size, vocab_size]
probs = F.softmax(last_logits, dim=-1)
next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1]
speculative_tokens.append(next_token)
# Concatenate speculative tokens to the currently generated sequence
speculative_tokens_tensor = torch.cat(speculative_tokens, dim=1) # [batch_size, num_future_tokens]
extended_generated = torch.cat([generated.transpose(0,1), speculative_tokens_tensor], dim=1) # [batch_size, seq_length + num_future_tokens]
extended_generated = extended_generated.transpose(0,1) # Transpose back to [seq_length + num_future_tokens, batch_size]
# Use the first output head to validate speculative tokens
extended_memory = model.encode(extended_generated) # [extended_seq_length, batch_size, d_model]
validation_logits = model.forward_head(extended_memory, head_index=0) # [extended_seq_length, batch_size, vocab_size]
# Extract logits for the speculative token part
validation_logits = validation_logits[-model.num_future_tokens:, :, :] # [num_future_tokens, batch_size, vocab_size]
# Calculate acceptance probabilities
acceptance_probs = []
for i in range(model.num_future_tokens):
logits = validation_logits[i, :, :] # [batch_size, vocab_size]
probs = F.softmax(logits, dim=-1)
token = speculative_tokens[i] # [batch_size, 1]
prob = probs.gather(dim=1, index=token) # [batch_size, 1]
acceptance_probs.append(prob)
acceptance_probs_tensor = torch.cat(acceptance_probs, dim=1) # [batch_size, num_future_tokens]
# Decide which tokens to accept based on the probability threshold
accept_mask = acceptance_probs_tensor >= prob_threshold # [batch_size, num_future_tokens]
accept_lengths = accept_mask.sum(dim=1) # [batch_size]
# Process each sample (assuming batch_size = 1)
num_accept = accept_lengths.item()
if num_accept > 0:
# Accept the first num_accept tokens
accepted_tokens = speculative_tokens_tensor[0, :num_accept].unsqueeze(1) # [num_accept, 1]
generated = torch.cat([generated, accepted_tokens], dim=0)
else:
# If no tokens are accepted, generate one token
logits = outputs['logits_head_0'][-1, :, :] # [batch_size, vocab_size]
probs = F.softmax(logits, dim=-1)
next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1]
generated = torch.cat([generated, next_token], dim=0)
# Convert generated token IDs to text
generated_text_indices = generated.squeeze(1).tolist() # [total_seq_length]
generated_text = tokenizer.decode(generated_text_indices, skip_special_tokens=True)
return generated_text
# Create model instance
model = MultiTokenTransformer(
vocab_size=vocab_size,
d_model=d_model,
nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dim_feedforward=dim_feedforward,
max_seq_length=max_seq_length,
num_future_tokens=num_future_tokens,
dropout=dropout)
The logic of speculative decoding implemented here is as follows: during one forward pass, all heads generate results sequentially. The main module runs only once, and the prediction heads infer multiple times to generate multiple tokens. Then, all the next words generated by the heads are concatenated together, and the model runs the forward pass again, using only the main head’s results to verify how many of the speculative results are acceptable.
V3 pre-training has included the prediction loss from each head to ensure the prediction accuracy of each head. If it can accurately predict three consecutive tokens each time, achieving a threefold acceleration in inference is not surprising.
That’s all for this sharing! I have roughly introduced the structure of MTP in V3 and the process of achieving speculative decoding through it. The implementation in the subsequent pseudocode is very similar to V3’s MTP, but since V3’s MTP implementation has not been open-sourced, it is only for reference to aid understanding. Feel free to leave comments or contact me privately for discussion.
References:
[1]https://arxiv.org/abs/2404.19737
[2]https://zhuanlan.zhihu.com/p/15037286337
Currently, I have several clients looking for algorithm, large model, AI engineering, and AI infrastructure-related positions, with demands comparable to Alibaba’s P6-10 across various locations including Beijing, Shanghai, Guangzhou, Shenzhen, and Hangzhou.All experts are welcome to self-recommend or recommend others. Please send resumes to [email protected].Feel free to add our consultant on WeChat for more detailed job information: Linda_lucy2010. You can also discuss your job-seeking needs with our consultant or apply to join our related communities.