164302c7df
Co-authored-by: Christian Bahls <christian.bahls@planet-ai.de> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
from dataclasses import dataclass
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from transformers import PretrainedConfig
|
|
|
|
from sglang.srt.model_executor.model_runner import ForwardBatch
|
|
|
|
|
|
@dataclass
|
|
class SparseEmbeddingOutput:
|
|
embeddings: torch.Tensor # [batch_size, vocab_size]
|
|
|
|
|
|
class SparsePooler(nn.Module):
|
|
"""A layer that pools hidden states into sparse vocabulary-space embeddings.
|
|
|
|
This layer does the following:
|
|
1. Applies a linear transformation + ReLU to get token-level weights
|
|
2. Maps these weights to vocabulary positions using token IDs
|
|
3. Aggregates weights for repeated tokens using max pooling
|
|
4. Returns sparse embeddings in vocabulary space
|
|
|
|
Attributes:
|
|
config: Model configuration containing vocab_size and hidden_size
|
|
sparse_linear: Linear layer for computing token weights
|
|
vocab_size: Size of vocabulary for output embeddings
|
|
"""
|
|
|
|
def __init__(self, config: PretrainedConfig):
|
|
super().__init__()
|
|
|
|
# Validate required attributes
|
|
if not hasattr(config, "vocab_size"):
|
|
raise AttributeError(
|
|
f"Config {type(config)} missing required 'vocab_size' attribute"
|
|
)
|
|
if not hasattr(config, "hidden_size"):
|
|
raise AttributeError(
|
|
f"Config {type(config)} missing required 'hidden_size' attribute"
|
|
)
|
|
|
|
self.vocab_size = config.vocab_size
|
|
self.sparse_linear = nn.Linear(config.hidden_size, 1)
|
|
self._weights_loaded = False
|
|
|
|
def forward(
|
|
self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
|
|
) -> SparseEmbeddingOutput:
|
|
"""
|
|
Forward pass for sparse pooling.
|
|
|
|
Args:
|
|
hidden_states: Packed sequence hidden states [total_tokens, hidden_size]
|
|
forward_batch: Batch information with sequence lengths and input_ids
|
|
|
|
Returns:
|
|
SparseEmbeddingOutput with embeddings of shape [batch_size, vocab_size]
|
|
"""
|
|
if not self._weights_loaded:
|
|
raise ValueError(
|
|
"Sparse pooling weights not loaded. Call load_weights() first"
|
|
)
|
|
|
|
# Apply sparse linear + ReLU to get token weights
|
|
token_weights = F.relu(self.sparse_linear(hidden_states)).squeeze(
|
|
-1
|
|
) # [total_tokens]
|
|
|
|
# Create batch indices for packed sequences
|
|
batch_indices = torch.repeat_interleave(
|
|
torch.arange(
|
|
len(forward_batch.extend_seq_lens), device=hidden_states.device
|
|
),
|
|
forward_batch.extend_seq_lens,
|
|
)
|
|
|
|
# Initialize sparse embedding output
|
|
sparse_embedding = torch.zeros(
|
|
len(forward_batch.extend_seq_lens),
|
|
self.vocab_size,
|
|
dtype=token_weights.dtype,
|
|
device=token_weights.device,
|
|
)
|
|
|
|
# Map to vocabulary space using scatter_reduce with amax
|
|
flat_indices = batch_indices * self.vocab_size + forward_batch.input_ids
|
|
sparse_embedding.view(-1).scatter_reduce_(
|
|
0, flat_indices, token_weights, reduce="amax"
|
|
)
|
|
|
|
return SparseEmbeddingOutput(embeddings=sparse_embedding)
|
|
|
|
def load_weights(self, state_dict: dict):
|
|
"""Load weights from state dict (called by the model)."""
|
|
self.sparse_linear.load_state_dict(state_dict)
|
|
self._weights_loaded = True
|