A continuation of kimi source tree.txt. The complete source identity, every recorded occurrence, and linked equations are on its first page.
Source transcript · continued
Download complete source text ↓6,675 lines · 275.4 KB
bias: whether to use bias in linear layer.
"""
def __init__(self, dims: list[int], activation, bias=True):
super().__init__()
assert len(dims) == 3
self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
self.activation = activation
for m in [self.fc0, self.fc1]:
nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features))
if m.bias is not None:
nn.init.zeros_(m.bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.fc0(x)
x = self.activation(x)
return self.fc1(x)
class MoonViTEncoderLayer(nn.Module):
def __init__(
self,
num_heads: int,
hidden_dim: int,
mlp_dim: int,
*,
attn_implementation: str = 'flash_attention_2',
activation=F.gelu,
attn_bias: bool = False,
use_deterministic_attn: bool = False,
):
super().__init__()
self.num_heads = num_heads
self.hidden_dim = hidden_dim
self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
self.attn_implementation = attn_implementation
self.use_deterministic_attn = use_deterministic_attn
self.norm0 = nn.LayerNorm(hidden_dim)
self.norm1 = nn.LayerNorm(hidden_dim)
self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)
self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)
self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias)
def attention_qkvpacked(
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
max_seqlen: torch.Tensor,
rope_freqs_cis: torch.Tensor | None = None,
):
"""
Args:
x (torch.Tensor): (batch_size, seqlen, hidden_dim)
cu_seqlens (torch.Tensor):
"""
xqkv = self.wqkv(x)
qkv_shape = xqkv.size()[:-1] + (
3,
self.num_heads,
self.hidden_size_per_attention_head,
)
# xqkv: (batch_size, seqlen, 3, nheads, headdim)
xqkv = xqkv.view(*qkv_shape)
xq, xk, xv = torch.unbind(xqkv, dim=-3)
xq, xk = apply_rope(xq, xk, rope_freqs_cis)
attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation]
attn_out = attn_func(xq,
xk,
xv,
q_cu_seqlens=cu_seqlens,
k_cu_seqlens=cu_seqlens,
max_seqlen_k=max_seqlen,
max_seqlen_q=max_seqlen,
deterministic=self.use_deterministic_attn)
attn_out = self.wo(attn_out)
return attn_out
def forward(
self,
hidden_states: torch.Tensor,
cu_seqlens: torch.Tensor,
max_seqlen: int,
rope_freqs_cis: torch.Tensor | None = None,
):
residual = hidden_states
hidden_states = self.norm0(hidden_states)
hidden_states = self.attention_qkvpacked(hidden_states, cu_seqlens,
max_seqlen, rope_freqs_cis)
hidden_states = residual + hidden_states
residual = hidden_states
hidden_states = self.norm1(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
return hidden_states
class MoonViT3dEncoder(nn.Module):
def __init__(self,
hidden_dim: int,
num_layers: int,
block_cfg: dict,
video_attn_type: str = 'spatial_temporal') -> None:
super().__init__()
assert video_attn_type == 'spatial_temporal', f'video_attn_type must be "spatial_temporal", got {video_attn_type}'
self.video_attn_type = video_attn_type
self.rope_2d = Rope2DPosEmbRepeated(
block_cfg['hidden_dim'] // block_cfg['num_heads'], 512, 512)
self.blocks = nn.ModuleList([
MoonViTEncoderLayer(
**block_cfg,
use_deterministic_attn=self.use_deterministic_attn)
for _ in range(num_layers)
])
self.final_layernorm = nn.LayerNorm(hidden_dim)
def forward(
self,
hidden_states: torch.Tensor,
grid_thws: torch.Tensor,
) -> torch.Tensor:
rope_freqs_cis = self.rope_2d.get_freqs_cis(
grid_thws=grid_thws, device=hidden_states.device)
lengths = torch.cat((
torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device),
grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2],
))
max_seqlen = lengths.max()
cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0,
dtype=torch.int32)
for block in self.blocks:
hidden_states = block(hidden_states,
cu_seqlens,
max_seqlen,
rope_freqs_cis=rope_freqs_cis)
hidden_states = self.final_layernorm(hidden_states)
return hidden_states
def tpool_patch_merger(
x: torch.Tensor,
grid_thws: torch.Tensor,
merge_kernel_size: tuple[int, int] = (2, 2),
) -> list[torch.Tensor]:
d_model = x.size(-1)
outputs = []
pre_sum = 0
for t, h, w in grid_thws.tolist():
# Get the current sequence
seq = x[pre_sum:pre_sum + t * h * w]
# Reshape along self.merge_kernel_size and concat to the last dimension
kernel_height, kernel_width = merge_kernel_size
new_height, new_width = h // kernel_height, w // kernel_width
reshaped_seq = seq.view(t, new_height, kernel_height, new_width,
kernel_width, d_model)
reshaped_seq = reshaped_seq.permute(0, 1,
3, 2, 4, 5).contiguous().mean(
dim=0) # temporal pooling
padded_seq = reshaped_seq.view(new_height * new_width,
kernel_height * kernel_width, -1)
outputs.append(padded_seq)
pre_sum += t * h * w
return outputs
class MoonViT3dPretrainedModel(PreTrainedModel):
config_class = None
model_type = 'moonvit3d'
_no_split_modules = ['PackingTransformer']
_supports_flash_attn_2 = True
_supports_sdpa = True
def __init__(self, config, *inputs, **kwargs):
super().__init__(config, *inputs, **kwargs)
config = deepcopy(config)
self.merge_kernel_size = config.merge_kernel_size
self.patch_size = config.patch_size
self.merge_type = config.merge_type
self.patch_embed = MoonVision3dPatchEmbed(
out_dim=config.hidden_size,
patch_size=config.patch_size,
pos_emb_height=config.init_pos_emb_height,
pos_emb_width=config.init_pos_emb_width,
pos_emb_time=config.init_pos_emb_time,
pos_emb_type=config.pos_emb_type,
)
self.encoder = MoonViT3dEncoder(hidden_dim=config.hidden_size,
num_layers=config.num_hidden_layers,
block_cfg={
'num_heads':
config.num_attention_heads,
'hidden_dim':
config.hidden_size,
'mlp_dim':
config.intermediate_size,
'activation':
PytorchGELUTanh(),
'attn_bias':
True,
'attn_implementation':
config._attn_implementation,
},
video_attn_type=config.video_attn_type)
def forward(self, pixel_values: torch.Tensor,
grid_thws: torch.Tensor) -> torch.Tensor:
"""
Args:
pixel_values (torch.Tensor): The input pixel values.
grid_thws (torch.Tensor): Temporal, height and width.
Returns:
torch.Tensor: The output tokens.
"""
# grid_thws = grid_thws.to('cpu')
assert grid_thws.ndim == 2, f'grid_thws should be 2D, got {grid_thws.ndim}'
assert grid_thws.size(1) == 3, f'No support for thw: {grid_thws}'
hidden_states = self.patch_embed(pixel_values, grid_thws)
hidden_states = self.encoder(hidden_states, grid_thws)
if self.merge_type == 'sd2_tpool': # spatial downsampling 2x with temporal pooling all
hidden_states = tpool_patch_merger(
hidden_states,
grid_thws,
merge_kernel_size=self.merge_kernel_size)
else:
raise NotImplementedError(f'Not support {self.merge_type}')
return hidden_states
# ============================================================================
# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py)
# ============================================================================
class IdentityMap(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, *args, **kwargs):
return x
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
# TODO, use faster LayerNorm
self.pre_norm = nn.LayerNorm(config.mm_hidden_size)
self.proj = nn.Sequential(
nn.Linear(config.mm_hidden_size, config.hidden_size), nn.GELU(),
nn.Linear(config.hidden_size, config.hidden_size))
def forward(self, x, *args, **kwargs):
assert isinstance(x,
list | tuple), f'x is not a list or tuple: {type(x)}'
lengths = [item.shape[0] for item in x]
x = torch.cat(x, dim=0)
x = self.pre_norm(x)
x = self.proj(x)
x = torch.split(x, lengths, dim=0)
return x
class PatchMergerMLP(nn.Module):
def __init__(self, config):
super().__init__()
eps = config.projector_ln_eps
self.hidden_size = config.mm_hidden_size * (
config.merge_kernel_size[0] * config.merge_kernel_size[1])
self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps)
self.proj = nn.Sequential(
nn.Linear(self.hidden_size, self.hidden_size),
nn.GELU(),
nn.Linear(self.hidden_size, config.hidden_size),
)
def forward(self, x, *args, **kwargs):
if isinstance(x, list) or isinstance(x, tuple):
x = [
self.proj(self.pre_norm(item).view(item.shape[0], -1))
for item in x
]
else:
# B, N, N_k, C = x.shape
B = x.shape[0]
x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size))
return x
class KimiK25PreTrainedModel(PreTrainedModel):
config_class = KimiK25Config
base_model_prefix = "model"
_no_split_modules = [
"MoonViT3dPretrainedModel",
"MoonViTEncoderLayer",
"DeepseekDecoderLayer",
"PatchMergerMLP",
]
_skip_keys_device_placement = "past_key_values"
_supports_flash_attn_2 = True
_supports_sdpa = False
def _init_weights(self, module):
# important: this ported version of Llava isn't meant for training from scratch - only
# inference and fine-tuning - so the proper init weights code has been removed - the original codebase
# https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose
std = (self.config.initializer_range if hasattr(
self.config, "initializer_range") else
self.config.text_config.initializer_range)
if hasattr(module, "class_embedding"):
module.class_embedding.data.normal_(mean=0.0, std=std)
if isinstance(module, (nn.Linear, nn.Conv2d)):
module.weight.data.normal_(mean=0.0, std=std)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.Embedding):
module.weight.data.normal_(mean=0.0, std=std)
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
class VisionTowerConfig(PretrainedConfig):
model_type = 'moonvit3d'
def __init__(self, config: KimiK25Config, **kwargs):
super().__init__(**kwargs)
self.patch_size = config.patch_size
self.init_pos_emb_height = config.init_pos_emb_height
self.init_pos_emb_width = config.init_pos_emb_width
self.init_pos_emb_time = config.init_pos_emb_time
self.pos_emb_type = config.pos_emb_type
self.num_attention_heads = config.vt_num_attention_heads
self.num_hidden_layers = config.vt_num_hidden_layers
self.hidden_size = config.vt_hidden_size
self.intermediate_size = config.vt_intermediate_size
self.merge_kernel_size = config.merge_kernel_size
self.video_attn_type = config.video_attn_type
self.merge_type = config.merge_type
self._attn_implementation = config._attn_implementation
class ProjectorConfig:
def __init__(self, config: KimiK25Config):
self.mm_projector_type = config.mm_projector_type
self.mm_hidden_size = config.mm_hidden_size
self.hidden_size = config.text_hidden_size
self.merge_kernel_size = config.merge_kernel_size
self.projector_hidden_act = config.projector_hidden_act
self.projector_ln_eps = config.projector_ln_eps
# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240
class KimiK25ForConditionalGeneration(KimiK25PreTrainedModel):
def __init__(self, config: KimiK25Config):
super().__init__(config)
vt_config = VisionTowerConfig(config.vision_config)
self.vision_tower = MoonViT3dPretrainedModel(vt_config)
proj_config = ProjectorConfig(config.vision_config)
if proj_config.mm_projector_type == 'identity':
self.mm_projector = IdentityMap()
elif proj_config.mm_projector_type == 'mlp':
self.mm_projector = MLP(proj_config)
elif proj_config.mm_projector_type == 'patchmerger':
self.mm_projector = PatchMergerMLP(proj_config)
else:
raise ValueError(
f"Unsupported mm_projector_type: {proj_config.mm_projector_type}"
)
self.language_model = DeepseekV3ForCausalLM(config.text_config)
self.post_init()
if hasattr(self.language_model, 'dtype'):
target_dtype = self.language_model.dtype
self.vision_tower = self.vision_tower.to(dtype=target_dtype)
self.mm_projector = self.mm_projector.to(dtype=target_dtype)
def get_input_embeddings(self):
return self.language_model.get_input_embeddings()
def set_input_embeddings(self, value):
self.language_model.set_input_embeddings(value)
def get_output_embeddings(self):
return self.language_model.get_output_embeddings()
def set_output_embeddings(self, new_embeddings):
self.language_model.set_output_embeddings(new_embeddings)
def set_decoder(self, decoder):
self.language_model.set_decoder(decoder)
def get_decoder(self):
return self.language_model.get_decoder()
def tie_weights(self):
return self.language_model.tie_weights()
def resize_token_embeddings(self,
new_num_tokens: int | None = None,
pad_to_multiple_of=None) -> nn.Embedding:
model_embeds = self.language_model.resize_token_embeddings(
new_num_tokens, pad_to_multiple_of)
# update vocab size
self.config.text_config.vocab_size = model_embeds.num_embeddings
self.vocab_size = model_embeds.num_embeddings
return model_embeds
def _merge_input_ids_with_image_features(
self,
image_features: list[torch.Tensor],
inputs_embeds: torch.Tensor,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
labels: torch.Tensor | None = None,
):
"""
Args:
image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`):
The image features to merge with the input embeddings.
inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`):
The input embeddings.
input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`):
The input ids.
attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`):
The attention mask.
labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*):
The labels.
"""
_, embed_dim = image_features[0].shape
feature_lengths = [x.shape[0] for x in image_features]
image_features = torch.cat(image_features, dim=0)
image_token_index: int = self.config.media_placeholder_token_id
pad_token_id: int = self.config.pad_token_id
ignore_index: int = self.config.ignore_index
batch_size, sequence_length = input_ids.shape
left_padding = not torch.sum(
input_ids[:, -1] == torch.tensor(pad_token_id))
# 1. Create a mask to know where special image tokens are
_token_occupation_table = torch.ones_like(input_ids.flatten())
_token_occupation_table[input_ids.flatten() ==
image_token_index] = torch.tensor(
feature_lengths,
dtype=torch.long,
device=input_ids.device)
_token_occupation_table = _token_occupation_table.reshape(
input_ids.shape)
max_embed_dim = _token_occupation_table.sum(-1).max().item()
assert (
max_embed_dim >= sequence_length
), f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})"
batch_indices, non_image_indices = torch.where(
input_ids != image_token_index)
# 2. Compute the positions where text should be written
# Calculate new positions for text tokens in merged image-text sequence.
new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1
nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1]
if left_padding:
new_token_positions += nb_image_pad[:,
None] # offset for left padding
text_to_overwrite = new_token_positions[batch_indices,
non_image_indices]
# 3. Create the full embedding, already padded to the maximum position
final_embedding = torch.zeros(
batch_size,
max_embed_dim,
embed_dim,
dtype=inputs_embeds.dtype,
device=inputs_embeds.device,
)
final_attention_mask = torch.zeros(batch_size,
max_embed_dim,
dtype=attention_mask.dtype,
device=inputs_embeds.device)
if labels is not None:
final_labels = torch.full(
(batch_size, max_embed_dim),
ignore_index,
dtype=input_ids.dtype,
device=input_ids.device,
)
# In case the Vision model or the Language model has been offloaded to CPU, we need to manually
# set the corresponding tensors into their correct target device.
target_device = inputs_embeds.device
batch_indices, non_image_indices, text_to_overwrite = (
batch_indices.to(target_device),
non_image_indices.to(target_device),
text_to_overwrite.to(target_device),
)
attention_mask = attention_mask.to(target_device)
# 4. Fill the embeddings based on the mask.
final_embedding[batch_indices,
text_to_overwrite] = inputs_embeds[batch_indices,
non_image_indices]
final_attention_mask[batch_indices,
text_to_overwrite] = attention_mask[
batch_indices, non_image_indices]
if labels is not None:
final_labels[batch_indices,
text_to_overwrite] = labels[batch_indices,
non_image_indices]
# 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835)
image_to_overwrite = torch.full((batch_size, max_embed_dim),
True,
dtype=torch.bool,
device=inputs_embeds.device)
image_to_overwrite[batch_indices, text_to_overwrite] = False
image_to_overwrite &= image_to_overwrite.cumsum(
-1) - 1 >= nb_image_pad[:, None].to(target_device)
if image_to_overwrite.sum() != image_features.shape[:-1].numel():
raise ValueError(
f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while"
f" the number of image features given to the model is {image_features.shape[:-1].numel()}. "
"This prevents correct indexing and breaks batch generation.")
final_embedding[image_to_overwrite] = (
image_features.contiguous().reshape(-1,
embed_dim).to(target_device))
final_attention_mask |= image_to_overwrite
position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_(
(final_attention_mask == 0), 1)
# 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens.
batch_indices, pad_indices = torch.where(input_ids == pad_token_id)
indices_to_mask = new_token_positions[batch_indices, pad_indices]
final_embedding[batch_indices, indices_to_mask] = 0
if labels is None:
final_labels = None
return final_embedding, final_attention_mask, final_labels, position_ids
def _extract_image_features(self, pixel_values: torch.Tensor,
grid_thws: torch.Tensor) -> list[torch.Tensor]:
"""
Args:
pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`):
The pixel values of the images processed by image processor.
grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`):
The grid, height, width of the images.
Returns:
selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`):
The selected image features to use as input to the projector head.
"""
target_dtype = self.vision_tower.patch_embed.proj.weight.dtype
pixel_values = pixel_values.to(target_dtype)
image_features = self.vision_tower(pixel_values, grid_thws)
return image_features
def forward(
self,
input_ids: torch.LongTensor | None = None,
pixel_values: torch.FloatTensor | list[torch.FloatTensor]
| None = None,
grid_thws: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: list[torch.FloatTensor] | None = None,
inputs_embeds: torch.FloatTensor | None = None,
labels: torch.LongTensor | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
) -> tuple | LlavaCausalLMOutputWithPast:
r"""
Args:
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
```"""
assert self.vision_tower is not None, "vision_tower is not loaded"
output_attentions = (output_attentions if output_attentions is not None
else self.config.output_attentions)
output_hidden_states = (output_hidden_states
if output_hidden_states is not None else
self.config.output_hidden_states)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if inputs_embeds is None:
# 1. Extra the input embeddings
inputs_embeds = self.get_input_embeddings()(input_ids)
# 2. Merge text and images
if pixel_values is not None and len(
pixel_values) > 0 and input_ids.shape[1] != 1:
image_features = self._extract_image_features(
pixel_values, grid_thws)
if self.mm_projector:
image_features = self.mm_projector(image_features)
inputs_embeds = inputs_embeds.to(
image_features[0].dtype) # num_tokens, embed_dim
inputs_embeds, attention_mask, labels, position_ids = (
self._merge_input_ids_with_image_features(
image_features,
inputs_embeds,
input_ids,
attention_mask,
labels,
))
# In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of
# generation with cache
elif (past_key_values is not None and pixel_values is not None
and input_ids.shape[1] == 1):
# Retrieve the first layer to inspect the logits and mask out the hidden states
# that are set to 0
first_layer_past_key_value = past_key_values[0][0][:, :, :, 0]
# Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941
batch_index, non_attended_tokens = torch.where(
first_layer_past_key_value.float().sum(-2) == 0)
# Get the target length
target_length = input_ids.shape[1]
past_length = first_layer_past_key_value.shape[-1]
extended_attention_mask = torch.ones(
(attention_mask.shape[0], past_length),
dtype=attention_mask.dtype,
device=attention_mask.device,
)
# Filter out only the tokens that can be un-attended, this can happen
# if one uses Llava + Fused modules where the cache on the
# first iteration is already big enough, or if one passes custom cache
valid_indices = non_attended_tokens < extended_attention_mask.size(
-1)
new_batch_index = batch_index[valid_indices]
new_non_attended_tokens = non_attended_tokens[valid_indices]
# Zero-out the places where we don't need to attend
extended_attention_mask[new_batch_index,
new_non_attended_tokens] = 0
attention_mask = torch.cat(
(extended_attention_mask, attention_mask[:,
-target_length:]),
dim=1)
position_ids = torch.sum(attention_mask,
dim=1).unsqueeze(-1) - 1
outputs = self.language_model(
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
logits = outputs[0]
loss = None
if labels is not None:
# Shift so that tokens < n predict n
if attention_mask is not None:
shift_attention_mask = attention_mask[..., 1:]
shift_logits = logits[..., :-1, :][shift_attention_mask.to(
logits.device) != 0].contiguous()
shift_labels = labels[..., 1:][shift_attention_mask.to(
labels.device) != 0].contiguous()
else:
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1).to(shift_logits.device),
)
if not return_dict:
output = (logits, ) + outputs[1:]
return (loss, ) + output if loss is not None else output
return LlavaCausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
def prepare_inputs_for_generation(
self,
input_ids,
past_key_values=None,
inputs_embeds=None,
pixel_values=None,
grid_thws=None,
attention_mask=None,
**kwargs,
):
if past_key_values is not None:
if isinstance(past_key_values, Cache):
cache_length = past_key_values.get_seq_length()
past_length = getattr(past_key_values, 'seen_tokens',
cache_length)
else:
cache_length = past_length = past_key_values[0][0].shape[2]
# Keep only the unprocessed tokens:
# 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
# some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
# input)
if attention_mask is not None and attention_mask.shape[
1] > input_ids.shape[1]:
input_ids = input_ids[:, -(attention_mask.shape[1] -
past_length):]
# 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
# input_ids based on the past_length.
elif past_length < input_ids.shape[1]:
input_ids = input_ids[:, past_length:]
# 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
elif self.config.media_placeholder_token_id in input_ids:
input_ids = input_ids[:, input_ids.shape[1] - 1:]
# If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the
# older attention values, as their corresponding values are not part of the input.
if cache_length < past_length and attention_mask is not None:
attention_mask = attention_mask[:, -(cache_length +
input_ids.shape[1]):]
position_ids = kwargs.get("position_ids", None)
if attention_mask is not None and position_ids is None:
# create position_ids on the fly for batch generation
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 1)
if past_key_values:
position_ids = position_ids[:, -input_ids.shape[1]:]
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
if inputs_embeds is not None and past_key_values is None:
model_inputs = {"inputs_embeds": inputs_embeds}
else:
model_inputs = {"input_ids": input_ids}
model_inputs.update({
"position_ids": position_ids,
"past_key_values": past_key_values,
"use_cache": kwargs.get("use_cache"),
"attention_mask": attention_mask,
"pixel_values": pixel_values,
"grid_thws": grid_thws,
})
return model_inputs
def _reorder_cache(self, *args, **kwargs):
return self.language_model._reorder_cache(*args, **kwargs)
>>>>> FILE: modeling_deepseek.py (80663 bytes) <<<<<
# coding=utf-8
# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" PyTorch DeepSeek model."""
import math
import warnings
from typing import List, Optional, Tuple, Union
import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_attn_mask_utils import \
_prepare_4d_causal_attention_mask
from transformers.modeling_outputs import (BaseModelOutputWithPast,
CausalLMOutputWithPast,
SequenceClassifierOutputWithPast)
from transformers.modeling_utils import PreTrainedModel
from transformers.pytorch_utils import (ALL_LAYERNORM_LAYERS,
is_torch_greater_or_equal_than_1_13)
from transformers.utils import (add_start_docstrings,
add_start_docstrings_to_model_forward,
is_flash_attn_2_available,
is_flash_attn_greater_or_equal_2_10, logging,
replace_return_docstrings)
from transformers.utils.import_utils import is_torch_fx_available
from .configuration_deepseek import DeepseekV3Config
if is_flash_attn_2_available():
from flash_attn import flash_attn_func, flash_attn_varlen_func
from flash_attn.bert_padding import pad_input # noqa
from flash_attn.bert_padding import index_first_axis, unpad_input
# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
# It means that the function will not be traced through and simply appear as a node in the graph.
if is_torch_fx_available():
if not is_torch_greater_or_equal_than_1_13:
import torch.fx
_prepare_4d_causal_attention_mask = torch.fx.wrap(
_prepare_4d_causal_attention_mask)
logger = logging.get_logger(__name__)
_CONFIG_FOR_DOC = "DeepseekV3Config"
def _get_unpad_data(attention_mask):
seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
max_seqlen_in_batch = seqlens_in_batch.max().item()
cu_seqlens = F.pad(
torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
return (
indices,
cu_seqlens,
max_seqlen_in_batch,
)
# code modified from transformers 4.48.3 to amend breaks in newer transformers versions
def get_usable_length(past_key_value,
new_seq_length: int,
layer_idx: Optional[int] = 0) -> int:
max_length = past_key_value.get_max_cache_shape()
previous_seq_length = past_key_value.get_seq_length(layer_idx)
if max_length is not None and max_length > 0 and previous_seq_length + new_seq_length > max_length:
return max_length - new_seq_length
return previous_seq_length
class DeepseekV3RMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
"""
DeepseekV3RMSNorm is equivalent to T5LayerNorm
"""
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance +
self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
ALL_LAYERNORM_LAYERS.append(DeepseekV3RMSNorm)
class DeepseekV3RotaryEmbedding(nn.Module):
def __init__(self,
dim,
max_position_embeddings=2048,
base=10000,
device=None):
super().__init__()
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
inv_freq = 1.0 / (self.base**(
torch.arange(0, self.dim, 2).float().to(device) / self.dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
# Build here to make `torch.jit.trace` work.
self._set_cos_sin_cache(
seq_len=max_position_embeddings,
device=self.inv_freq.device,
dtype=torch.get_default_dtype(),
)
self.max_seq_len_cached = None
def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached,
device=device,
dtype=self.inv_freq.dtype)
freqs = torch.outer(t, self.inv_freq.to(t.device))
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim=-1)
self.register_buffer("cos_cached",
emb.cos().to(dtype),
persistent=False)
self.register_buffer("sin_cached",
emb.sin().to(dtype),
persistent=False)
def forward(self, x, seq_len=None):
# x: [bs, num_attention_heads, seq_len, head_size]
if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached:
self._set_cos_sin_cache(seq_len=seq_len,
device=x.device,
dtype=x.dtype)
return (
self.cos_cached[:seq_len].to(dtype=x.dtype),
self.sin_cached[:seq_len].to(dtype=x.dtype),
)
# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->DeepseekV3
class DeepseekV3LinearScalingRotaryEmbedding(DeepseekV3RotaryEmbedding):
"""DeepseekV3RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
def __init__(
self,
dim,
max_position_embeddings=2048,
base=10000,
device=None,
scaling_factor=1.0,
):
self.scaling_factor = scaling_factor
super().__init__(dim, max_position_embeddings, base, device)
def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached,
device=device,
dtype=self.inv_freq.dtype)
t = t / self.scaling_factor
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim=-1)
self.register_buffer("cos_cached",
emb.cos().to(dtype),
persistent=False)
self.register_buffer("sin_cached",
emb.sin().to(dtype),
persistent=False)
# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->DeepseekV3
class DeepseekV3DynamicNTKScalingRotaryEmbedding(DeepseekV3RotaryEmbedding):
"""DeepseekV3RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
def __init__(
self,
dim,
max_position_embeddings=2048,
base=10000,
device=None,
scaling_factor=1.0,
):
self.scaling_factor = scaling_factor
super().__init__(dim, max_position_embeddings, base, device)
def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len
if seq_len > self.max_position_embeddings:
base = self.base * ((self.scaling_factor * seq_len /
self.max_position_embeddings) -
(self.scaling_factor - 1))**(self.dim /
(self.dim - 2))
inv_freq = 1.0 / (base**(
torch.arange(0, self.dim, 2).float().to(device) / self.dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
t = torch.arange(self.max_seq_len_cached,
device=device,
dtype=self.inv_freq.dtype)
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim=-1)
self.register_buffer("cos_cached",
emb.cos().to(dtype),
persistent=False)
self.register_buffer("sin_cached",
emb.sin().to(dtype),
persistent=False)
# Inverse dim formula to find dim based on number of rotations
def yarn_find_correction_dim(num_rotations,
dim,
base=10000,
max_position_embeddings=2048):
return (dim * math.log(max_position_embeddings /
(num_rotations * 2 * math.pi))) / (2 *
math.log(base))
# Find dim range bounds based on rotations
def yarn_find_correction_range(low_rot,
high_rot,
dim,
base=10000,
max_position_embeddings=2048):
low = math.floor(
yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings))
high = math.ceil(
yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings))
return max(low, 0), min(high, dim - 1) # Clamp values just in case
def yarn_get_mscale(scale=1, mscale=1):
if scale <= 1:
return 1.0
return 0.1 * mscale * math.log(scale) + 1.0
def yarn_linear_ramp_mask(min, max, dim):
if min == max:
max += 0.001 # Prevent singularity
linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
ramp_func = torch.clamp(linear_func, 0, 1)
return ramp_func
class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding):
def __init__(
self,
dim,
max_position_embeddings=2048,
base=10000,
device=None,
scaling_factor=1.0,
original_max_position_embeddings=4096,
beta_fast=32,
beta_slow=1,
mscale=1,
mscale_all_dim=0,
):
self.scaling_factor = scaling_factor
self.original_max_position_embeddings = original_max_position_embeddings
self.beta_fast = beta_fast
self.beta_slow = beta_slow
self.mscale = mscale
self.mscale_all_dim = mscale_all_dim
super().__init__(dim, max_position_embeddings, base, device)
def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len
dim = self.dim
freq_extra = 1.0 / (self.base**(
torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
freq_inter = 1.0 / (self.scaling_factor * self.base**(
torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
low, high = yarn_find_correction_range(
self.beta_fast,
self.beta_slow,
dim,
self.base,
self.original_max_position_embeddings,
)
inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to(
device=device, dtype=torch.float32)
inv_freq = freq_inter * (1 -
inv_freq_mask) + freq_extra * inv_freq_mask
self.register_buffer("inv_freq", inv_freq, persistent=False)
t = torch.arange(seq_len, device=device, dtype=torch.float32)
freqs = torch.outer(t, inv_freq)
_mscale = float(
yarn_get_mscale(self.scaling_factor, self.mscale) /
yarn_get_mscale(self.scaling_factor, self.mscale_all_dim))
emb = torch.cat((freqs, freqs), dim=-1)
self.register_buffer("cos_cached", (emb.cos() * _mscale).to(dtype),
persistent=False)
self.register_buffer("sin_cached", (emb.sin() * _mscale).to(dtype),
persistent=False)
# Copied from transformers.models.llama.modeling_llama.rotate_half
def rotate_half(x):
"""Rotates half the hidden dims of the input."""
x1 = x[..., :x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2:]
return torch.cat((-x2, x1), dim=-1)
# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
"""Applies Rotary Position Embedding to the query and key tensors.
Args:
q (`torch.Tensor`): The query tensor.
k (`torch.Tensor`): The key tensor.
cos (`torch.Tensor`): The cosine part of the rotary embedding.
sin (`torch.Tensor`): The sine part of the rotary embedding.
position_ids (`torch.Tensor`):
The position indices of the tokens corresponding to the query and key tensors. For example, this can be
used to pass offsetted position ids when working with a KV-cache.
unsqueeze_dim (`int`, *optional*, defaults to 1):
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
Returns:
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
"""
cos = cos[position_ids].unsqueeze(unsqueeze_dim)
sin = sin[position_ids].unsqueeze(unsqueeze_dim)
b, h, s, d = q.shape
q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
b, h, s, d = k.shape
k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed, k_embed
class DeepseekV3MLP(nn.Module):
def __init__(self, config, hidden_size=None, intermediate_size=None, is_routed_expert=False):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
self.intermediate_size = (config.intermediate_size if intermediate_size
is None else intermediate_size)
# ROUTED EXPERTS (the int4 intelligence bulk): when exact generators are on, the cells ARE
# PackedInt4Linear whose buffers (weight_packed, weight_scale) match the shipped checkpoint keys — so the
# native int4 loads straight in with zero conversion, never expands to bf16, and computes the MAC exactly.
if is_routed_expert and getattr(config, "use_exact_generators", False):
from exact_generators import PackedInt4Linear
ab = getattr(config, "exact_act_bits", 16)
self.gate_proj = PackedInt4Linear(self.hidden_size, self.intermediate_size, act_bits=ab)
self.up_proj = PackedInt4Linear(self.hidden_size, self.intermediate_size, act_bits=ab)
self.down_proj = PackedInt4Linear(self.intermediate_size, self.hidden_size, act_bits=ab)
else: # shared expert + dense layer stay bf16 nn.Linear
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, x):
down_proj = self.down_proj(
self.act_fn(self.gate_proj(x)) * self.up_proj(x))
return down_proj
def to_exact(self, act_bits=16):
"""Rewrite this MLP's three projections into EXACT-integer generators, in place. Experts (int4 native)
become ~950x more precise than the shipped bf16-dequant path AND smaller; the dense layer uses int8.
This is the bottom-up rewrite: every MAC becomes an exact int64 accumulation, deterministic."""
from exact_generators import ExactLinear
moe_i = getattr(self.config, "moe_intermediate_size", None)
num_bits = 4 if (moe_i is not None and self.intermediate_size == moe_i) else 8 # experts=int4, dense=int8
self.gate_proj = ExactLinear.from_weight(self.gate_proj.weight.data, num_bits, 32, act_bits)
self.up_proj = ExactLinear.from_weight(self.up_proj.weight.data, num_bits, 32, act_bits)
self.down_proj = ExactLinear.from_weight(self.down_proj.weight.data, num_bits, 32, act_bits)
return self
def convert_kimi_to_exact(model, act_bits=16):
"""Rewrite the whole model into exact-integer generators: every expert + dense MLP. Attention/router/lm_head
stay at bf16 (quantizing them down loses precision). Result: the intelligence bulk is exact + smaller."""
n = 0
for m in model.modules():
if isinstance(m, DeepseekV3MLP):
m.to_exact(act_bits); n += 1
return n
class MoEGate(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.top_k = config.num_experts_per_tok
self.n_routed_experts = config.n_routed_experts
self.routed_scaling_factor = config.routed_scaling_factor
self.scoring_func = config.scoring_func
self.seq_aux = config.seq_aux
self.topk_method = config.topk_method
self.n_group = config.n_group
self.topk_group = config.topk_group
# topk selection algorithm
self.norm_topk_prob = config.norm_topk_prob
self.gating_dim = config.hidden_size
self.weight = nn.Parameter(
torch.empty((self.n_routed_experts, self.gating_dim)))
if self.topk_method == "noaux_tc":
self.e_score_correction_bias = nn.Parameter(
torch.empty((self.n_routed_experts)))
self.reset_parameters()
def reset_parameters(self) -> None:
import torch.nn.init as init
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
def forward(self, hidden_states):
bsz, seq_len, h = hidden_states.shape
### compute gating score
hidden_states = hidden_states.view(-1, h)
logits = F.linear(hidden_states.type(torch.float32),
self.weight.type(torch.float32), None)
if self.scoring_func == "sigmoid":
scores = logits.sigmoid()
else:
raise NotImplementedError(
f"insupportable scoring function for MoE gating: {self.scoring_func}"
)
### select top-k experts
if self.topk_method == "noaux_tc":
assert not self.training
scores_for_choice = scores.view(
bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)
group_scores = (scores_for_choice.view(
bsz * seq_len, self.n_group,
-1).topk(2, dim=-1)[0].sum(dim=-1)) # [n, n_group]
group_idx = torch.topk(group_scores,
k=self.topk_group,
dim=-1,
sorted=False)[1] # [n, top_k_group]
group_mask = torch.zeros_like(group_scores) # [n, n_group]
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
score_mask = (group_mask.unsqueeze(-1).expand(
bsz * seq_len, self.n_group,
self.n_routed_experts // self.n_group).reshape(
bsz * seq_len, -1)) # [n, e]
tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(),
0.0) # [n, e]
_, topk_idx = torch.topk(tmp_scores,
k=self.top_k,
dim=-1,
sorted=False)
topk_weight = scores.gather(1, topk_idx)
else:
raise NotImplementedError(
f"insupportable TopK function for MoE gating: {self.topk_method}"
)
### norm gate to sum 1
if self.top_k > 1 and self.norm_topk_prob:
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
topk_weight = topk_weight / denominator
topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor
return topk_idx, topk_weight
class DeepseekV3MoE(nn.Module):
"""
A mixed expert module containing shared experts.
"""
def __init__(self, config):
super().__init__()
self.config = config
self.num_experts_per_tok = config.num_experts_per_tok
if hasattr(config, "ep_size") and config.ep_size > 1:
assert config.ep_size == dist.get_world_size()
self.ep_size = config.ep_size
self.experts_per_rank = config.n_routed_experts // config.ep_size
self.ep_rank = dist.get_rank()
self.experts = nn.ModuleList([
(DeepseekV3MLP(config,
intermediate_size=config.moe_intermediate_size)
if i >= self.ep_rank * self.experts_per_rank
and i < (self.ep_rank + 1) * self.experts_per_rank else None)
for i in range(config.n_routed_experts)
])
else:
self.ep_size = 1
self.experts_per_rank = config.n_routed_experts
self.ep_rank = 0
self.experts = nn.ModuleList([
DeepseekV3MLP(config,
intermediate_size=config.moe_intermediate_size,
is_routed_expert=True)
for i in range(config.n_routed_experts)
])
self.gate = MoEGate(config)
if config.n_shared_experts is not None:
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
self.shared_experts = DeepseekV3MLP(
config=config, intermediate_size=intermediate_size)
def forward(self, hidden_states):
identity = hidden_states
orig_shape = hidden_states.shape
topk_idx, topk_weight = self.gate(hidden_states)
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
flat_topk_idx = topk_idx.view(-1)
if not self.training:
y = self.moe_infer(hidden_states, topk_idx,
topk_weight).view(*orig_shape)
if self.config.n_shared_experts is not None:
y = y + self.shared_experts(identity)
return y
@torch.no_grad()
def moe_infer(self, x, topk_ids, topk_weight):
cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
cnts.scatter_(1, topk_ids, 1)
tokens_per_expert = cnts.sum(dim=0)
idxs = topk_ids.view(-1).argsort()
sorted_tokens = x[idxs // topk_ids.shape[1]]
sorted_tokens_shape = sorted_tokens.shape
if self.ep_size > 1:
tokens_per_ep_rank = tokens_per_expert.view(self.ep_size,
-1).sum(dim=1)
tokens_per_expert_group = tokens_per_expert.new_empty(
tokens_per_expert.shape[0])
dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert)
output_splits = (tokens_per_expert_group.view(
self.ep_size, -1).sum(1).cpu().numpy().tolist())
gathered_tokens = sorted_tokens.new_empty(
tokens_per_expert_group.sum(dim=0).cpu().item(),
sorted_tokens.shape[1])
input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()
dist.all_to_all(
list(gathered_tokens.split(output_splits)),
list(sorted_tokens.split(input_split_sizes)),
)
tokens_per_expert_post_gather = tokens_per_expert_group.view(
self.ep_size, self.experts_per_rank).sum(dim=0)
gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0], ),
dtype=np.int32)
s = 0
for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
gatherd_idxs[s:s + k] = i % self.experts_per_rank
s += k
gatherd_idxs = gatherd_idxs.argsort()
sorted_tokens = gathered_tokens[gatherd_idxs]
tokens_per_expert = tokens_per_expert_post_gather
tokens_per_expert = tokens_per_expert.cpu().numpy()
outputs = []
start_idx = 0
for i, num_tokens in enumerate(tokens_per_expert):
end_idx = start_idx + num_tokens
if num_tokens == 0:
continue
expert = self.experts[i + self.ep_rank * self.experts_per_rank]
tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
expert_out = expert(tokens_for_this_expert)
outputs.append(expert_out)
start_idx = end_idx
outs = torch.cat(outputs,
dim=0) if len(outputs) else sorted_tokens.new_empty(0)
if self.ep_size > 1:
new_x = torch.empty_like(outs)
new_x[gatherd_idxs] = outs
gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
dist.all_to_all(
list(gathered_tokens.split(input_split_sizes)),
list(new_x.split(output_splits)),
)
outs = gathered_tokens
new_x = torch.empty_like(outs)
new_x[idxs] = outs
final_out = (new_x.view(
*topk_ids.shape, -1).type(topk_weight.dtype).mul_(
topk_weight.unsqueeze(dim=-1)).sum(dim=1).type(new_x.dtype))
return final_out
# Copied from transformers.models.llama.modeling_llama.repeat_kv
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
"""
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
"""
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
if n_rep == 1:
return hidden_states
hidden_states = hidden_states[:, :,
None, :, :].expand(batch,
num_key_value_heads,
n_rep, slen, head_dim)
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen,
head_dim)
# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV3
class DeepseekV3Attention(nn.Module):
"""Multi-headed attention from 'Attention Is All You Need' paper"""
def __init__(self,
config: DeepseekV3Config,
layer_idx: Optional[int] = None):
super().__init__()
self.config = config
self.layer_idx = layer_idx
if layer_idx is None:
logger.warning_once(
f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
"to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
"when creating this class.")
self.attention_dropout = config.attention_dropout
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.max_position_embeddings = config.max_position_embeddings
self.rope_theta = config.rope_theta
self.q_lora_rank = config.q_lora_rank
self.qk_rope_head_dim = config.qk_rope_head_dim
self.kv_lora_rank = config.kv_lora_rank
self.v_head_dim = config.v_head_dim
self.qk_nope_head_dim = config.qk_nope_head_dim
self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
self.is_causal = True
if self.q_lora_rank is None:
self.q_proj = nn.Linear(self.hidden_size,
self.num_heads * self.q_head_dim,
bias=False)
else:
self.q_a_proj = nn.Linear(self.hidden_size,
config.q_lora_rank,
bias=config.attention_bias)
self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)
self.q_b_proj = nn.Linear(config.q_lora_rank,
self.num_heads * self.q_head_dim,
bias=False)
self.kv_a_proj_with_mqa = nn.Linear(
self.hidden_size,