From b926c3dad52f405c580a4a2bbe9e32ca9bd8e8f6 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 9 Sep 2024 08:06:45 +0200 Subject: [PATCH 01/86] init commit --- src/transformers/modeling_utils.py | 29 +++++++++++++++++++++++++++++ tests/utils/test_modeling_utils.py | 7 ++++++- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 2faa60210ed4..5caaffe49729 100755 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -26,6 +26,7 @@ import shutil import tempfile import warnings +from collections import defaultdict from contextlib import contextmanager from dataclasses import dataclass from functools import partial, wraps @@ -990,6 +991,31 @@ def _add_variant(weights_name: str, variant: Optional[str] = None) -> str: return weights_name +def update_key_name(keys): + """ + Updates a dictionary of keys to pack layers together as layer.{0, 1, 4} instead of layers.0, layers.1, layers.4. + """ + key_dict = defaultdict(set) + for key in keys: + modified_key = r"" + for char in key: + if char.isdigit(): + if modified_key != "": + modified_key += r"(\d+)" + modified_key = modified_key.replace(".", r"\.") + key_dict[modified_key].add(int(char)) + modified_key = r"" + else: + modified_key += char + final_keys = set() + for key in keys: + text = key + for pattern, values in key_dict.items(): + text = re.sub(pattern, lambda match: match.group(0)[: -len(match.group(1))] + str(values), text) + final_keys.add(text) + return final_keys + + class ModuleUtilsMixin: """ A few utilities for `torch.nn.Modules`, to be used as a mixin. @@ -4506,6 +4532,9 @@ def _find_mismatched_keys( ) raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}") + missing_keys = update_key_name(missing_keys) + unexpected_keys = update_key_name(unexpected_keys) + if len(unexpected_keys) > 0: archs = [] if model.config.architectures is None else model.config.architectures warner = logger.warning if model.__class__.__name__ in archs else logger.info diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index f78285fdb90d..15edefc314b9 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -72,7 +72,7 @@ is_torch_sdpa_available, is_torchdynamo_available, ) - +from transformers.modeling_utils import update_key_name sys.path.append(str(Path(__file__).parent.parent.parent / "utils")) @@ -1715,6 +1715,11 @@ def test_isin_mps_friendly(self): torch.equal(torch.isin(random_ids, random_test_tensor), isin_mps_friendly(random_ids, random_test_tensor)) ) + def test_update_key_name(self): + original_keys = ['model.language_model.0.self_attn.0.mlp0', 'model.language_model.0.self_attn.1.mlp.conv1.weight', 'model.language_model.1.self_attn.0.mlp.conv1.weight', 'model.language_model.1.self_attn.1.mlp.conv1.weight', 'model.language_model.2.self_attn.0.mlp.conv1.weight', 'model.language_model.2.self_attn.1.conv1.weight', 'model.language_model.3.self_attn.0.mlp.conv1.weight', 'model.language_model.3.self_attn.1.mlp.conv1.weight', 'model.language_model.0.self_attn.0.mlp.layer.conv1.weight', 'model.language_model.1.self_attn.0.mlp.layer.conv1.weight', 'model.language_model.2.self_attn.0.mlp.layer.conv1.weight', 'model.language_model.3.self_attn.0.mlp.layer.conv1.weight', 'model.language_model.0.self_attn.0.mlp.layer.1.weight', 'model.language_model.1.self_attn.0.mlp.layer.1.weight', 'model.language_model.2.self_attn.0.mlp.layer.1.weight', 'model.language_model.3.self_attn.0.mlp.layer.1.weight'] + new_keys = update_key_name(original_keys) + expected_new_keys = ['model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.{1}.weight', 'model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.conv{1}.weight', 'model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp{0}', 'model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.conv{1}.weight', 'model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.conv{1}.weight'] + self.assertListEqual(list(new_keys), expected_new_keys) @slow @require_torch From 5b850236256af630f4f7fe68e5c9d3ece0e725ac Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 9 Sep 2024 08:08:43 +0200 Subject: [PATCH 02/86] style --- tests/utils/test_modeling_utils.py | 31 +++++++++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 15edefc314b9..6ef9917b9b4e 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -41,6 +41,7 @@ is_torch_available, logging, ) +from transformers.modeling_utils import update_key_name from transformers.testing_utils import ( TOKEN, USER, @@ -72,7 +73,7 @@ is_torch_sdpa_available, is_torchdynamo_available, ) -from transformers.modeling_utils import update_key_name + sys.path.append(str(Path(__file__).parent.parent.parent / "utils")) @@ -1716,11 +1717,35 @@ def test_isin_mps_friendly(self): ) def test_update_key_name(self): - original_keys = ['model.language_model.0.self_attn.0.mlp0', 'model.language_model.0.self_attn.1.mlp.conv1.weight', 'model.language_model.1.self_attn.0.mlp.conv1.weight', 'model.language_model.1.self_attn.1.mlp.conv1.weight', 'model.language_model.2.self_attn.0.mlp.conv1.weight', 'model.language_model.2.self_attn.1.conv1.weight', 'model.language_model.3.self_attn.0.mlp.conv1.weight', 'model.language_model.3.self_attn.1.mlp.conv1.weight', 'model.language_model.0.self_attn.0.mlp.layer.conv1.weight', 'model.language_model.1.self_attn.0.mlp.layer.conv1.weight', 'model.language_model.2.self_attn.0.mlp.layer.conv1.weight', 'model.language_model.3.self_attn.0.mlp.layer.conv1.weight', 'model.language_model.0.self_attn.0.mlp.layer.1.weight', 'model.language_model.1.self_attn.0.mlp.layer.1.weight', 'model.language_model.2.self_attn.0.mlp.layer.1.weight', 'model.language_model.3.self_attn.0.mlp.layer.1.weight'] + original_keys = [ + "model.language_model.0.self_attn.0.mlp0", + "model.language_model.0.self_attn.1.mlp.conv1.weight", + "model.language_model.1.self_attn.0.mlp.conv1.weight", + "model.language_model.1.self_attn.1.mlp.conv1.weight", + "model.language_model.2.self_attn.0.mlp.conv1.weight", + "model.language_model.2.self_attn.1.conv1.weight", + "model.language_model.3.self_attn.0.mlp.conv1.weight", + "model.language_model.3.self_attn.1.mlp.conv1.weight", + "model.language_model.0.self_attn.0.mlp.layer.conv1.weight", + "model.language_model.1.self_attn.0.mlp.layer.conv1.weight", + "model.language_model.2.self_attn.0.mlp.layer.conv1.weight", + "model.language_model.3.self_attn.0.mlp.layer.conv1.weight", + "model.language_model.0.self_attn.0.mlp.layer.1.weight", + "model.language_model.1.self_attn.0.mlp.layer.1.weight", + "model.language_model.2.self_attn.0.mlp.layer.1.weight", + "model.language_model.3.self_attn.0.mlp.layer.1.weight", + ] new_keys = update_key_name(original_keys) - expected_new_keys = ['model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.{1}.weight', 'model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.conv{1}.weight', 'model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp{0}', 'model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.conv{1}.weight', 'model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.conv{1}.weight'] + expected_new_keys = [ + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.{1}.weight", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.conv{1}.weight", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp{0}", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.conv{1}.weight", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.conv{1}.weight", + ] self.assertListEqual(list(new_keys), expected_new_keys) + @slow @require_torch class ModelOnTheFlyConversionTester(unittest.TestCase): From 3b76bdae692863de40ad945ba90318eb3c7bf6b0 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 9 Sep 2024 10:50:03 +0200 Subject: [PATCH 03/86] take comments into account --- src/transformers/modeling_utils.py | 5 +++-- tests/utils/test_modeling_utils.py | 10 +++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 5caaffe49729..b6614517a1bb 100755 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -999,7 +999,7 @@ def update_key_name(keys): for key in keys: modified_key = r"" for char in key: - if char.isdigit(): + if char.isdigit() and modified_key[-1] == ".": if modified_key != "": modified_key += r"(\d+)" modified_key = modified_key.replace(".", r"\.") @@ -1011,7 +1011,8 @@ def update_key_name(keys): for key in keys: text = key for pattern, values in key_dict.items(): - text = re.sub(pattern, lambda match: match.group(0)[: -len(match.group(1))] + str(values), text) + if len(values) > 1: + text = re.sub(pattern, lambda match: match.group(0)[: -len(match.group(1))] + str(values), text) final_keys.add(text) return final_keys diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 6ef9917b9b4e..f0147f7dada6 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -1737,11 +1737,11 @@ def test_update_key_name(self): ] new_keys = update_key_name(original_keys) expected_new_keys = [ - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.{1}.weight", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.conv{1}.weight", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp{0}", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.conv{1}.weight", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.conv{1}.weight", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.conv1.weight", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.1.weight", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp0", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.conv1.weight", + "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.conv1.weight", ] self.assertListEqual(list(new_keys), expected_new_keys) From 704767e05c8245e6bfbd0a2c1f4f5dddc86a9702 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 14:42:30 +0900 Subject: [PATCH 04/86] add deepseekv3 modeling --- .../deepseekv3/configuration_deepseekv3.py | 204 +++ .../models/deepseekv3/modeling_deepseekv3.py | 1222 +++++++++++++++++ .../models/deepseekv3/modular_deepseekv3.py | 521 +++++++ 3 files changed, 1947 insertions(+) create mode 100644 src/transformers/models/deepseekv3/configuration_deepseekv3.py create mode 100644 src/transformers/models/deepseekv3/modeling_deepseekv3.py create mode 100644 src/transformers/models/deepseekv3/modular_deepseekv3.py diff --git a/src/transformers/models/deepseekv3/configuration_deepseekv3.py b/src/transformers/models/deepseekv3/configuration_deepseekv3.py new file mode 100644 index 000000000000..c2a32399d6d5 --- /dev/null +++ b/src/transformers/models/deepseekv3/configuration_deepseekv3.py @@ -0,0 +1,204 @@ +from ...configuration_utils import PretrainedConfig +from ...utils import logging + +logger = logging.get_logger(__name__) + +DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} +class DeepseekV3Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the DeepSeek-V3. + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + Args: + vocab_size (`int`, *optional*, defaults to 129280): + Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`DeepseekV3Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations. + moe_intermediate_size (`int`, *optional*, defaults to 1407): + Dimension of the MoE representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_nextn_predict_layers (`int`, *optional*, defaults to 1): + Number of nextn predict layers in the DeepSeekV3 Model. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + n_shared_experts (`int`, *optional*, defaults to None): + Number of shared experts, None means dense model. + n_routed_experts (`int`, *optional*, defaults to None): + Number of routed experts, None means dense model. + routed_scaling_factor (`float`, *optional*, defaults to 1.0): + Scaling factor or routed experts. + topk_method (`str`, *optional*, defaults to `gready`): + Topk method used in routed gate. + n_group (`int`, *optional*, defaults to None): + Number of groups for routed experts. + topk_group (`int`, *optional*, defaults to None): + Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). + num_experts_per_tok (`int`, *optional*, defaults to None): + Number of selected experts, None means dense model. + moe_layer_freq (`int`, *optional*, defaults to 1): + The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. + first_k_dense_replace (`int`, *optional*, defaults to 0): + Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). + \--k dense layers--/ + norm_topk_prob (`bool`, *optional*, defaults to False): + Whether to normalize the weights of the routed experts. + scoring_func (`str`, *optional*, defaults to 'softmax'): + Method of computing expert weights. + aux_loss_alpha (`float`, *optional*, defaults to 0.001): + Auxiliary loss weight coefficient. + seq_aux = (`bool`, *optional*, defaults to True): + Whether to compute the auxiliary loss for each individual sample. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + pad_token_id (`int`, *optional*): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + Beginning of stream token id. + eos_token_id (`int`, *optional*, defaults to 2): + End of stream token id. + pretraining_tp (`int`, *optional*, defaults to 1): + Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this + document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is + necessary to ensure exact reproducibility of the pretraining results. Please refer to [this + issue](https://github.com/pytorch/pytorch/issues/76232). + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + rope_scaling (`Dict`, *optional*): + Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling + strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is + `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update + `max_position_embeddings` to the expected new maximum. + attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + Whether to use a bias in the query, key, value and output projection layers during self-attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + ```python + >>> from transformers import DeepseekV3Model, DeepseekV3Config + >>> # Initializing a Deepseek-V3 style configuration + >>> configuration = DeepseekV3Config() + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "deepseek_v3" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=129280, + hidden_size=7168, + intermediate_size=18432, + moe_intermediate_size = 2048, + num_hidden_layers=61, + num_nextn_predict_layers=1, + num_attention_heads=128, + num_key_value_heads=128, + n_shared_experts = 1, + n_routed_experts = 256, + ep_size = 1, + routed_scaling_factor = 2.5, + kv_lora_rank = 512, + q_lora_rank = 1536, + qk_rope_head_dim = 64, + v_head_dim = 128, + qk_nope_head_dim = 128, + topk_method = 'noaux_tc', + n_group = 8, + topk_group = 4, + num_experts_per_tok = 8, + moe_layer_freq = 1, + first_k_dense_replace = 3, + norm_topk_prob = True, + scoring_func = 'sigmoid', + aux_loss_alpha = 0.001, + seq_aux = True, + hidden_act="silu", + max_position_embeddings=4096, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=0, + eos_token_id=1, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_nextn_predict_layers = num_nextn_predict_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.ep_size = ep_size + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.topk_method = topk_method + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + self.aux_loss_alpha = aux_loss_alpha + self.seq_aux = seq_aux + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/src/transformers/models/deepseekv3/modeling_deepseekv3.py b/src/transformers/models/deepseekv3/modeling_deepseekv3.py new file mode 100644 index 000000000000..9c00887e30f5 --- /dev/null +++ b/src/transformers/models/deepseekv3/modeling_deepseekv3.py @@ -0,0 +1,1222 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/deepseekv3/modular_deepseekv3.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_deepseekv3.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +import math +from typing import Callable, List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache, StaticCache +from ...generation import GenerationMixin +from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import ( + LossKwargs, + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, + replace_return_docstrings, +) +from ...utils.deprecation import deprecate_kwarg +from .configuration_deepseekv3 import DeepseekV3Config + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "DeepseekV3Config" + + +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) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class DeepseekV3RotaryEmbedding(nn.Module): + def __init__(self, config: DeepseekV3Config, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + def _dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len) + self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @torch.no_grad() + def forward(self, x, position_ids): + if "dynamic" in self.rope_type: + self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class DeepseekV3MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + 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 + + +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) + 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]) + 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 + + +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) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +# 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 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, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank) + self.kv_b_proj = nn.Linear( + config.kv_lora_rank, + self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=config.attention_bias, + ) + + self.rotary_emb = DeepseekV3RotaryEmbedding( + config=self.config, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor], + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + + k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + + if self.q_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class DeepseekV3DecoderLayer(nn.Module): + def __init__(self, config: DeepseekV3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = DeepseekV3Attention(config=config, layer_idx=layer_idx) + + self.mlp = ( + DeepseekV3MoE(config) + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + else DeepseekV3MLP(config) + ) + self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +DEEPSEEK_V3_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`DeepseekV3Config`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", + DEEPSEEK_V3_START_DOCSTRING, +) +class DeepseekV3PreTrainedModel(PreTrainedModel): + config_class = DeepseekV3Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["DeepseekV3DecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + _supports_attention_backend = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + 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_() + + +DEEPSEEK_V3_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance, see our + [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache); + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", + DEEPSEEK_V3_START_DOCSTRING, +) +class DeepseekV3Model(DeepseekV3PreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeepseekV3DecoderLayer`] + + Args: + config: DeepseekV3Config + """ + + def __init__(self, config: DeepseekV3Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [DeepseekV3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = DeepseekV3RotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ) -> Union[Tuple, BaseModelOutputWithPast]: + 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 + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + output = BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + return output if return_dict else output.to_tuple() + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and (attention_mask == 0.0).any(): + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_cache_shape() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + device=device, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + ) + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + min_dtype = torch.finfo(dtype).min + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + @staticmethod + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + device: torch.device, + cache_position: torch.Tensor, + batch_size: int, + **kwargs, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape + `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, + to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + device (`torch.device`): + The device to plcae the 4D attention mask on. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + + return causal_mask + + +class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... + + +class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + + def __init__(self, config): + super().__init__(config) + self.model = DeepseekV3Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **kwargs: Unpack[KwargsForCausalLM], + ) -> Union[Tuple, CausalLMOutputWithPast]: + 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]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM + + >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + 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 + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + 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, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + The DeepseekV3 Model transformer with a sequence classification head on top (linear layer). + + [`DeepseekV3ForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + DEEPSEEK_V3_START_DOCSTRING, +) +class DeepseekV3ForSequenceClassification(DeepseekV3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = DeepseekV3Model(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.model( + input_ids, + 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, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) diff --git a/src/transformers/models/deepseekv3/modular_deepseekv3.py b/src/transformers/models/deepseekv3/modular_deepseekv3.py new file mode 100644 index 000000000000..e68053765524 --- /dev/null +++ b/src/transformers/models/deepseekv3/modular_deepseekv3.py @@ -0,0 +1,521 @@ +import math +import warnings +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn + +from ...cache_utils import Cache +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import Unpack +from ...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 ..llama.modeling_llama import ( + LlamaDecoderLayer, + LlamaForCausalLM, + LlamaForSequenceClassification, + LlamaMLP, + LlamaModel, + LlamaPreTrainedModel, + LlamaRMSNorm, + LlamaRotaryEmbedding, + eager_attention_forward, +) +from .configuration_deepseekv3 import DeepseekV3Config +import torch.distributed as dist +import numpy as np + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + +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, + ) + + +class DeepseekV3RMSNorm(LlamaRMSNorm): + pass + + +class DeepseekV3RotaryEmbedding(LlamaRotaryEmbedding): + pass + + +# 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(LlamaMLP): + pass + + +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 + ) + 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]) + 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 + + +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, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank) + self.kv_b_proj = nn.Linear( + config.kv_lora_rank, + self.num_heads + * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=config.attention_bias, + ) + + self.rotary_emb = DeepseekV3RotaryEmbedding( + config=self.config, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor], + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + + k_nope, value_states = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + + if self.q_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class DeepseekV3DecoderLayer(LlamaDecoderLayer): + def __init__(self, config: DeepseekV3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = DeepseekV3Attention(config=config, layer_idx=layer_idx) + + self.mlp = ( + DeepseekV3MoE(config) + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + else DeepseekV3MLP(config) + ) + self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): + pass + + +class DeepseekV3Model(LlamaModel): + pass + + +class DeepseekV3ForCausalLM(LlamaForCausalLM): + pass + + +class DeepseekV3ForSequenceClassification(LlamaForSequenceClassification): + pass From 244e793c1ead36fb96fb1f6cd03780569a11abbd Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 14:51:04 +0900 Subject: [PATCH 05/86] remove redundant code --- .../models/deepseekv3/modular_deepseekv3.py | 32 ++----------------- 1 file changed, 2 insertions(+), 30 deletions(-) diff --git a/src/transformers/models/deepseekv3/modular_deepseekv3.py b/src/transformers/models/deepseekv3/modular_deepseekv3.py index e68053765524..f52d9ed560a3 100644 --- a/src/transformers/models/deepseekv3/modular_deepseekv3.py +++ b/src/transformers/models/deepseekv3/modular_deepseekv3.py @@ -1,6 +1,5 @@ import math -import warnings -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable, Optional, Tuple import torch import torch.nn.functional as F @@ -11,14 +10,7 @@ from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack -from ...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 ...utils import logging from ..llama.modeling_llama import ( LlamaDecoderLayer, LlamaForCausalLM, @@ -34,29 +26,9 @@ import torch.distributed as dist import numpy as np -if is_flash_attn_2_available(): - from flash_attn import flash_attn_func, flash_attn_varlen_func - from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa - 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, - ) - class DeepseekV3RMSNorm(LlamaRMSNorm): pass From 4fb2a80b3c9aeb8776c19ea5389e71085c5ddcb0 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 14:56:13 +0900 Subject: [PATCH 06/86] apply make style --- .../models/deepseekv3/configuration_deepseekv3.py | 1 + src/transformers/models/deepseekv3/modular_deepseekv3.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/deepseekv3/configuration_deepseekv3.py b/src/transformers/models/deepseekv3/configuration_deepseekv3.py index c2a32399d6d5..6098ac1755e7 100644 --- a/src/transformers/models/deepseekv3/configuration_deepseekv3.py +++ b/src/transformers/models/deepseekv3/configuration_deepseekv3.py @@ -1,6 +1,7 @@ from ...configuration_utils import PretrainedConfig from ...utils import logging + logger = logging.get_logger(__name__) DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} diff --git a/src/transformers/models/deepseekv3/modular_deepseekv3.py b/src/transformers/models/deepseekv3/modular_deepseekv3.py index f52d9ed560a3..cba30e878527 100644 --- a/src/transformers/models/deepseekv3/modular_deepseekv3.py +++ b/src/transformers/models/deepseekv3/modular_deepseekv3.py @@ -1,7 +1,9 @@ import math from typing import Callable, Optional, Tuple +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 @@ -23,8 +25,6 @@ eager_attention_forward, ) from .configuration_deepseekv3 import DeepseekV3Config -import torch.distributed as dist -import numpy as np logger = logging.get_logger(__name__) From 6b002e5e4386f8f51ab8ffca628d02145aa80732 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 15:02:17 +0900 Subject: [PATCH 07/86] apply fix-copies --- .../models/deepseekv3/modeling_deepseekv3.py | 19 ++++++------------- .../models/deepseekv3/modular_deepseekv3.py | 19 ++++++------------- 2 files changed, 12 insertions(+), 26 deletions(-) diff --git a/src/transformers/models/deepseekv3/modeling_deepseekv3.py b/src/transformers/models/deepseekv3/modeling_deepseekv3.py index 9c00887e30f5..b9c0bb928bc3 100644 --- a/src/transformers/models/deepseekv3/modeling_deepseekv3.py +++ b/src/transformers/models/deepseekv3/modeling_deepseekv3.py @@ -368,16 +368,16 @@ def rotate_half(x): # 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): +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. 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 @@ -388,15 +388,8 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): 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) - + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/transformers/models/deepseekv3/modular_deepseekv3.py b/src/transformers/models/deepseekv3/modular_deepseekv3.py index cba30e878527..de401b6e49b3 100644 --- a/src/transformers/models/deepseekv3/modular_deepseekv3.py +++ b/src/transformers/models/deepseekv3/modular_deepseekv3.py @@ -47,16 +47,16 @@ def rotate_half(x): # 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): +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. 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 @@ -67,15 +67,8 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): 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) - + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed From 4ec1e8871054ffa3b0b76caf6952fe820f199511 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 15:08:38 +0900 Subject: [PATCH 08/86] make format --- src/transformers/models/deepseekv3/modeling_deepseekv3.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/transformers/models/deepseekv3/modeling_deepseekv3.py b/src/transformers/models/deepseekv3/modeling_deepseekv3.py index b9c0bb928bc3..a6cb6c493d90 100644 --- a/src/transformers/models/deepseekv3/modeling_deepseekv3.py +++ b/src/transformers/models/deepseekv3/modeling_deepseekv3.py @@ -34,7 +34,6 @@ logger = logging.get_logger(__name__) - _CONFIG_FOR_DOC = "DeepseekV3Config" From 114ab84c127f932782627bc09657d17787aa5c63 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 15:28:11 +0900 Subject: [PATCH 09/86] add init files --- src/transformers/models/__init__.py | 1 + .../models/auto/configuration_auto.py | 2 ++ src/transformers/models/auto/modeling_auto.py | 2 ++ .../models/auto/tokenization_auto.py | 7 +++++ .../models/deepseekv3/__init__.py | 27 ++++++++++++++++++ .../models/deepseekv3/modeling_deepseekv3.py | 20 ++++++------- src/transformers/utils/dummy_pt_objects.py | 28 +++++++++++++++++++ 7 files changed, 77 insertions(+), 10 deletions(-) create mode 100644 src/transformers/models/deepseekv3/__init__.py diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index 43cf2fe42951..322d254aecf1 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -69,6 +69,7 @@ deberta, deberta_v2, decision_transformer, + deepseekv3, deformable_detr, deit, deprecated, diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 95c5bc4d008d..d1ad6291b39f 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -87,6 +87,7 @@ ("deberta", "DebertaConfig"), ("deberta-v2", "DebertaV2Config"), ("decision_transformer", "DecisionTransformerConfig"), + ("deepseekv3", "DeepseekV3Config"), ("deformable_detr", "DeformableDetrConfig"), ("deit", "DeiTConfig"), ("depth_anything", "DepthAnythingConfig"), @@ -406,6 +407,7 @@ ("deberta", "DeBERTa"), ("deberta-v2", "DeBERTa-v2"), ("decision_transformer", "Decision Transformer"), + ("deepseekv3", "DeepseekV3"), ("deformable_detr", "Deformable DETR"), ("deit", "DeiT"), ("deplot", "DePlot"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 8030f5dbbdaa..4368419a3a01 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -86,6 +86,7 @@ ("deberta", "DebertaModel"), ("deberta-v2", "DebertaV2Model"), ("decision_transformer", "DecisionTransformerModel"), + ("deepseekv3", "DeepSeekV3Model"), ("deformable_detr", "DeformableDetrModel"), ("deit", "DeiTModel"), ("deta", "DetaModel"), @@ -501,6 +502,7 @@ ("ctrl", "CTRLLMHeadModel"), ("data2vec-text", "Data2VecTextForCausalLM"), ("dbrx", "DbrxForCausalLM"), + ("deepseekv3", "DeepseekV3ForCausalLM"), ("diffllama", "DiffLlamaForCausalLM"), ("electra", "ElectraForCausalLM"), ("emu3", "Emu3ForCausalLM"), diff --git a/src/transformers/models/auto/tokenization_auto.py b/src/transformers/models/auto/tokenization_auto.py index 5ee4f612285f..b55148560853 100644 --- a/src/transformers/models/auto/tokenization_auto.py +++ b/src/transformers/models/auto/tokenization_auto.py @@ -170,6 +170,13 @@ "DebertaV2TokenizerFast" if is_tokenizers_available() else None, ), ), + ( + "deepseekv3", + ( + "LlamaTokenizer" if is_sentencepiece_available() else None, + "LlamaTokenizerFast" if is_tokenizers_available() else None, + ), + ), ( "diffllama", ( diff --git a/src/transformers/models/deepseekv3/__init__.py b/src/transformers/models/deepseekv3/__init__.py new file mode 100644 index 000000000000..649d0b88bc30 --- /dev/null +++ b/src/transformers/models/deepseekv3/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# 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. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_deepseekv3 import * + from .modeling_deepseekv3 import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/deepseekv3/modeling_deepseekv3.py b/src/transformers/models/deepseekv3/modeling_deepseekv3.py index a6cb6c493d90..47a1a00730f0 100644 --- a/src/transformers/models/deepseekv3/modeling_deepseekv3.py +++ b/src/transformers/models/deepseekv3/modeling_deepseekv3.py @@ -597,7 +597,7 @@ def forward( return outputs -DEEPSEEK_V3_START_DOCSTRING = r""" +DEEPSEEKV3_START_DOCSTRING = r""" This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) @@ -616,7 +616,7 @@ def forward( @add_start_docstrings( "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", - DEEPSEEK_V3_START_DOCSTRING, + DEEPSEEKV3_START_DOCSTRING, ) class DeepseekV3PreTrainedModel(PreTrainedModel): config_class = DeepseekV3Config @@ -644,7 +644,7 @@ def _init_weights(self, module): module.weight.data[module.padding_idx].zero_() -DEEPSEEK_V3_INPUTS_DOCSTRING = r""" +DEEPSEEKV3_INPUTS_DOCSTRING = r""" Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide @@ -721,7 +721,7 @@ def _init_weights(self, module): @add_start_docstrings( "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", - DEEPSEEK_V3_START_DOCSTRING, + DEEPSEEKV3_START_DOCSTRING, ) class DeepseekV3Model(DeepseekV3PreTrainedModel): """ @@ -753,7 +753,7 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.embed_tokens = value - @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) + @add_start_docstrings_to_model_forward(DEEPSEEKV3_INPUTS_DOCSTRING) def forward( self, input_ids: torch.LongTensor = None, @@ -1017,7 +1017,7 @@ def get_decoder(self): return self.model @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") - @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) + @add_start_docstrings_to_model_forward(DEEPSEEKV3_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) def forward( self, @@ -1056,8 +1056,8 @@ def forward( ```python >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM - >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") - >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") + >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseekv3/DeepseekV3-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseekv3/DeepseekV3-2-7b-hf") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") @@ -1123,7 +1123,7 @@ def forward( padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in each row of the batch). """, - DEEPSEEK_V3_START_DOCSTRING, + DEEPSEEKV3_START_DOCSTRING, ) class DeepseekV3ForSequenceClassification(DeepseekV3PreTrainedModel): def __init__(self, config): @@ -1141,7 +1141,7 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.embed_tokens = value - @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) + @add_start_docstrings_to_model_forward(DEEPSEEKV3_INPUTS_DOCSTRING) def forward( self, input_ids: Optional[torch.LongTensor] = None, diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index f379119289c6..6d08cf903d3e 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -3516,6 +3516,34 @@ def __init__(self, *args, **kwargs): requires_backends(self, ["torch"]) +class DeepseekV3ForCasualLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + class DepthAnythingForDepthEstimation(metaclass=DummyObject): _backends = ["torch"] From 779f8d2a823ac21423726545c33cfc9f9fc59dd4 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 23:21:36 +0900 Subject: [PATCH 10/86] rename deepseekv3 into deepseek_v3 based on its model_type --- docs/source/en/index.md | 1 + src/transformers/__init__.py | 18 ++++++ src/transformers/models/__init__.py | 2 +- .../{deepseekv3 => deepseek_v3}/__init__.py | 4 +- .../configuration_deepseek_v3.py} | 21 ++++++- .../modeling_deepseek_v3.py} | 28 +++++----- .../modular_deepseek_v3.py} | 4 +- src/transformers/utils/dummy_pt_objects.py | 56 +++++++++---------- 8 files changed, 82 insertions(+), 52 deletions(-) rename src/transformers/models/{deepseekv3 => deepseek_v3}/__init__.py (91%) rename src/transformers/models/{deepseekv3/configuration_deepseekv3.py => deepseek_v3/configuration_deepseek_v3.py} (92%) rename src/transformers/models/{deepseekv3/modeling_deepseekv3.py => deepseek_v3/modeling_deepseek_v3.py} (98%) rename src/transformers/models/{deepseekv3/modular_deepseekv3.py => deepseek_v3/modular_deepseek_v3.py} (98%) diff --git a/docs/source/en/index.md b/docs/source/en/index.md index 2233630128ae..ce9e41783d49 100644 --- a/docs/source/en/index.md +++ b/docs/source/en/index.md @@ -118,6 +118,7 @@ Flax), PyTorch, and/or TensorFlow. | [DeBERTa](model_doc/deberta) | βœ… | βœ… | ❌ | | [DeBERTa-v2](model_doc/deberta-v2) | βœ… | βœ… | ❌ | | [Decision Transformer](model_doc/decision_transformer) | βœ… | ❌ | ❌ | +| [Deepseek-V3](model_doc/deepseek_v3) | βœ… | ❌ | ❌ | | [Deformable DETR](model_doc/deformable_detr) | βœ… | ❌ | ❌ | | [DeiT](model_doc/deit) | βœ… | βœ… | ❌ | | [DePlot](model_doc/deplot) | βœ… | ❌ | ❌ | diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 02a55c28ac66..a66401355651 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -341,6 +341,7 @@ ], "models.deberta_v2": ["DebertaV2Config"], "models.decision_transformer": ["DecisionTransformerConfig"], + "models.deepseek_v3": ["DeepseekV3Config"], "models.deformable_detr": ["DeformableDetrConfig"], "models.deit": ["DeiTConfig"], "models.deprecated": [], @@ -1956,6 +1957,14 @@ "DecisionTransformerPreTrainedModel", ] ) + _import_structure["models.deepseek_v3"].extend( + [ + "DeepseekV3ForCausalLM", + "DeepseekV3ForSequenceClassification", + "DeepseekV3Model", + "DeepseekV3PreTrainedModel", + ] + ) _import_structure["models.deformable_detr"].extend( [ "DeformableDetrForObjectDetection", @@ -5395,6 +5404,9 @@ from .models.decision_transformer import ( DecisionTransformerConfig, ) + from .models.deepseek_v3 import ( + DeepseekV3Config, + ) from .models.deformable_detr import ( DeformableDetrConfig, ) @@ -6966,6 +6978,12 @@ DecisionTransformerModel, DecisionTransformerPreTrainedModel, ) + from .models.deepseek_v3 import ( + DeepseekV3ForCausalLM, + DeepseekV3ForSequenceClassification, + DeepseekV3Model, + DeepseekV3PreTrainedModel, + ) from .models.deformable_detr import ( DeformableDetrForObjectDetection, DeformableDetrModel, diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index 322d254aecf1..35df76833cb8 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -69,7 +69,7 @@ deberta, deberta_v2, decision_transformer, - deepseekv3, + deepseek_v3, deformable_detr, deit, deprecated, diff --git a/src/transformers/models/deepseekv3/__init__.py b/src/transformers/models/deepseek_v3/__init__.py similarity index 91% rename from src/transformers/models/deepseekv3/__init__.py rename to src/transformers/models/deepseek_v3/__init__.py index 649d0b88bc30..298f4c968375 100644 --- a/src/transformers/models/deepseekv3/__init__.py +++ b/src/transformers/models/deepseek_v3/__init__.py @@ -18,8 +18,8 @@ if TYPE_CHECKING: - from .configuration_deepseekv3 import * - from .modeling_deepseekv3 import * + from .configuration_deepseek_v3 import * + from .modeling_deepseek_v3 import * else: import sys diff --git a/src/transformers/models/deepseekv3/configuration_deepseekv3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py similarity index 92% rename from src/transformers/models/deepseekv3/configuration_deepseekv3.py rename to src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 6098ac1755e7..6841442a2229 100644 --- a/src/transformers/models/deepseekv3/configuration_deepseekv3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -1,8 +1,23 @@ -from ...configuration_utils import PretrainedConfig -from ...utils import logging +# coding=utf-8 +# Copyright 2025 bzantium and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on the DeepSeekV3 implementations from the DeepSeek AI team. (https://huggingface.co/deepseek-ai/DeepSeek-V3) + +# 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. +""" DeepSeekV3 model configuration """ +from ...configuration_utils import PretrainedConfig -logger = logging.get_logger(__name__) DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} class DeepseekV3Config(PretrainedConfig): diff --git a/src/transformers/models/deepseekv3/modeling_deepseekv3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py similarity index 98% rename from src/transformers/models/deepseekv3/modeling_deepseekv3.py rename to src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 47a1a00730f0..e49975b36c99 100644 --- a/src/transformers/models/deepseekv3/modeling_deepseekv3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -1,8 +1,8 @@ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 -# This file was automatically generated from src/transformers/models/deepseekv3/modular_deepseekv3.py. +# This file was automatically generated from src/transformers/models/deepseek_v3/modular_deepseek_v3.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the -# modular_deepseekv3.py file directly. One of our CI enforces this. +# modular_deepseek_v3.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 import math from typing import Callable, List, Optional, Tuple, Union @@ -30,7 +30,7 @@ replace_return_docstrings, ) from ...utils.deprecation import deprecate_kwarg -from .configuration_deepseekv3 import DeepseekV3Config +from .configuration_deepseek_v3 import DeepseekV3Config logger = logging.get_logger(__name__) @@ -358,7 +358,6 @@ def eager_attention_forward( return attn_output, attn_weights -# 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] @@ -366,7 +365,6 @@ def rotate_half(x): 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=None, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. @@ -597,7 +595,7 @@ def forward( return outputs -DEEPSEEKV3_START_DOCSTRING = r""" +DEEPSEEK_V3_START_DOCSTRING = r""" This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) @@ -616,7 +614,7 @@ def forward( @add_start_docstrings( "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", - DEEPSEEKV3_START_DOCSTRING, + DEEPSEEK_V3_START_DOCSTRING, ) class DeepseekV3PreTrainedModel(PreTrainedModel): config_class = DeepseekV3Config @@ -644,7 +642,7 @@ def _init_weights(self, module): module.weight.data[module.padding_idx].zero_() -DEEPSEEKV3_INPUTS_DOCSTRING = r""" +DEEPSEEK_V3_INPUTS_DOCSTRING = r""" Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide @@ -721,7 +719,7 @@ def _init_weights(self, module): @add_start_docstrings( "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", - DEEPSEEKV3_START_DOCSTRING, + DEEPSEEK_V3_START_DOCSTRING, ) class DeepseekV3Model(DeepseekV3PreTrainedModel): """ @@ -753,7 +751,7 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.embed_tokens = value - @add_start_docstrings_to_model_forward(DEEPSEEKV3_INPUTS_DOCSTRING) + @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) def forward( self, input_ids: torch.LongTensor = None, @@ -1017,7 +1015,7 @@ def get_decoder(self): return self.model @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") - @add_start_docstrings_to_model_forward(DEEPSEEKV3_INPUTS_DOCSTRING) + @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) def forward( self, @@ -1056,8 +1054,8 @@ def forward( ```python >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM - >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseekv3/DeepseekV3-2-7b-hf") - >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseekv3/DeepseekV3-2-7b-hf") + >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") @@ -1123,7 +1121,7 @@ def forward( padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in each row of the batch). """, - DEEPSEEKV3_START_DOCSTRING, + DEEPSEEK_V3_START_DOCSTRING, ) class DeepseekV3ForSequenceClassification(DeepseekV3PreTrainedModel): def __init__(self, config): @@ -1141,7 +1139,7 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.embed_tokens = value - @add_start_docstrings_to_model_forward(DEEPSEEKV3_INPUTS_DOCSTRING) + @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) def forward( self, input_ids: Optional[torch.LongTensor] = None, diff --git a/src/transformers/models/deepseekv3/modular_deepseekv3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py similarity index 98% rename from src/transformers/models/deepseekv3/modular_deepseekv3.py rename to src/transformers/models/deepseek_v3/modular_deepseek_v3.py index de401b6e49b3..913d35730dc7 100644 --- a/src/transformers/models/deepseekv3/modular_deepseekv3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -24,7 +24,7 @@ LlamaRotaryEmbedding, eager_attention_forward, ) -from .configuration_deepseekv3 import DeepseekV3Config +from .configuration_deepseek_v3 import DeepseekV3Config logger = logging.get_logger(__name__) @@ -38,7 +38,6 @@ class DeepseekV3RotaryEmbedding(LlamaRotaryEmbedding): pass -# 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] @@ -46,7 +45,6 @@ def rotate_half(x): 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=None, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index 6d08cf903d3e..6f1496bd1e16 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -2762,6 +2762,34 @@ def __init__(self, *args, **kwargs): requires_backends(self, ["torch"]) +class DeepseekV3ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + class DeformableDetrForObjectDetection(metaclass=DummyObject): _backends = ["torch"] @@ -3516,34 +3544,6 @@ def __init__(self, *args, **kwargs): requires_backends(self, ["torch"]) -class DeepseekV3ForCasualLM(metaclass=DummyObject): - _backends = ["torch"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch"]) - - -class DeepseekV3ForSequenceClassification(metaclass=DummyObject): - _backends = ["torch"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch"]) - - -class DeepseekV3Model(metaclass=DummyObject): - _backends = ["torch"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch"]) - - -class DeepseekV3PreTrainedModel(metaclass=DummyObject): - _backends = ["torch"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch"]) - - class DepthAnythingForDepthEstimation(metaclass=DummyObject): _backends = ["torch"] From 22623a39d6f5fbe8329fed581a21016b059a93db Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 23:23:15 +0900 Subject: [PATCH 11/86] rename deepseekv3 into deepseek_v3 based on its model_type --- src/transformers/models/auto/configuration_auto.py | 4 ++-- src/transformers/models/auto/modeling_auto.py | 4 ++-- src/transformers/models/auto/tokenization_auto.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index d1ad6291b39f..8b8145b14de2 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -87,7 +87,7 @@ ("deberta", "DebertaConfig"), ("deberta-v2", "DebertaV2Config"), ("decision_transformer", "DecisionTransformerConfig"), - ("deepseekv3", "DeepseekV3Config"), + ("deepseek_v3", "DeepseekV3Config"), ("deformable_detr", "DeformableDetrConfig"), ("deit", "DeiTConfig"), ("depth_anything", "DepthAnythingConfig"), @@ -407,7 +407,7 @@ ("deberta", "DeBERTa"), ("deberta-v2", "DeBERTa-v2"), ("decision_transformer", "Decision Transformer"), - ("deepseekv3", "DeepseekV3"), + ("deepseek_v3", "deepseek_v3"), ("deformable_detr", "Deformable DETR"), ("deit", "DeiT"), ("deplot", "DePlot"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 4368419a3a01..592dff25db49 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -86,7 +86,7 @@ ("deberta", "DebertaModel"), ("deberta-v2", "DebertaV2Model"), ("decision_transformer", "DecisionTransformerModel"), - ("deepseekv3", "DeepSeekV3Model"), + ("deepseek_v3", "DeepseekV3Model"), ("deformable_detr", "DeformableDetrModel"), ("deit", "DeiTModel"), ("deta", "DetaModel"), @@ -502,7 +502,7 @@ ("ctrl", "CTRLLMHeadModel"), ("data2vec-text", "Data2VecTextForCausalLM"), ("dbrx", "DbrxForCausalLM"), - ("deepseekv3", "DeepseekV3ForCausalLM"), + ("deepseek_v3", "DeepseekV3ForCausalLM"), ("diffllama", "DiffLlamaForCausalLM"), ("electra", "ElectraForCausalLM"), ("emu3", "Emu3ForCausalLM"), diff --git a/src/transformers/models/auto/tokenization_auto.py b/src/transformers/models/auto/tokenization_auto.py index b55148560853..c01de1986a53 100644 --- a/src/transformers/models/auto/tokenization_auto.py +++ b/src/transformers/models/auto/tokenization_auto.py @@ -171,7 +171,7 @@ ), ), ( - "deepseekv3", + "deepseek_v3", ( "LlamaTokenizer" if is_sentencepiece_available() else None, "LlamaTokenizerFast" if is_tokenizers_available() else None, From 78b19b058896789edcaa3cc723d0b2f336a20d4f Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 23:36:08 +0900 Subject: [PATCH 12/86] deepseek-v3 not deepseek_v3 --- docs/source/en/index.md | 2 +- src/transformers/models/auto/configuration_auto.py | 4 ++-- src/transformers/models/auto/modeling_auto.py | 4 ++-- src/transformers/models/auto/tokenization_auto.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/source/en/index.md b/docs/source/en/index.md index ce9e41783d49..9d34c2aaada7 100644 --- a/docs/source/en/index.md +++ b/docs/source/en/index.md @@ -118,7 +118,7 @@ Flax), PyTorch, and/or TensorFlow. | [DeBERTa](model_doc/deberta) | βœ… | βœ… | ❌ | | [DeBERTa-v2](model_doc/deberta-v2) | βœ… | βœ… | ❌ | | [Decision Transformer](model_doc/decision_transformer) | βœ… | ❌ | ❌ | -| [Deepseek-V3](model_doc/deepseek_v3) | βœ… | ❌ | ❌ | +| [Deepseek-V3](model_doc/deepseek-v3) | βœ… | ❌ | ❌ | | [Deformable DETR](model_doc/deformable_detr) | βœ… | ❌ | ❌ | | [DeiT](model_doc/deit) | βœ… | βœ… | ❌ | | [DePlot](model_doc/deplot) | βœ… | ❌ | ❌ | diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 8b8145b14de2..2402360e4893 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -87,7 +87,7 @@ ("deberta", "DebertaConfig"), ("deberta-v2", "DebertaV2Config"), ("decision_transformer", "DecisionTransformerConfig"), - ("deepseek_v3", "DeepseekV3Config"), + ("deepseek-v3", "DeepseekV3Config"), ("deformable_detr", "DeformableDetrConfig"), ("deit", "DeiTConfig"), ("depth_anything", "DepthAnythingConfig"), @@ -407,7 +407,7 @@ ("deberta", "DeBERTa"), ("deberta-v2", "DeBERTa-v2"), ("decision_transformer", "Decision Transformer"), - ("deepseek_v3", "deepseek_v3"), + ("deepseek-v3", "Deepseek-V3"), ("deformable_detr", "Deformable DETR"), ("deit", "DeiT"), ("deplot", "DePlot"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 592dff25db49..00e4892c0b6a 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -86,7 +86,7 @@ ("deberta", "DebertaModel"), ("deberta-v2", "DebertaV2Model"), ("decision_transformer", "DecisionTransformerModel"), - ("deepseek_v3", "DeepseekV3Model"), + ("deepseek-v3", "DeepseekV3Model"), ("deformable_detr", "DeformableDetrModel"), ("deit", "DeiTModel"), ("deta", "DetaModel"), @@ -502,7 +502,7 @@ ("ctrl", "CTRLLMHeadModel"), ("data2vec-text", "Data2VecTextForCausalLM"), ("dbrx", "DbrxForCausalLM"), - ("deepseek_v3", "DeepseekV3ForCausalLM"), + ("deepseek-v3", "DeepseekV3ForCausalLM"), ("diffllama", "DiffLlamaForCausalLM"), ("electra", "ElectraForCausalLM"), ("emu3", "Emu3ForCausalLM"), diff --git a/src/transformers/models/auto/tokenization_auto.py b/src/transformers/models/auto/tokenization_auto.py index c01de1986a53..8b862069c7d4 100644 --- a/src/transformers/models/auto/tokenization_auto.py +++ b/src/transformers/models/auto/tokenization_auto.py @@ -171,7 +171,7 @@ ), ), ( - "deepseek_v3", + "deepseek-v3", ( "LlamaTokenizer" if is_sentencepiece_available() else None, "LlamaTokenizerFast" if is_tokenizers_available() else None, From eb0e3a4e923ff70ac1735219171d98de2723fb31 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 28 Jan 2025 23:58:52 +0900 Subject: [PATCH 13/86] set model_type as deepseek_v3 --- .../models/auto/configuration_auto.py | 4 +- src/transformers/models/auto/modeling_auto.py | 4 +- .../models/auto/tokenization_auto.py | 2 +- .../deepseek_v3/configuration_deepseek_v3.py | 63 +++++++++++-------- .../deepseek_v3/modeling_deepseek_v3.py | 4 +- 5 files changed, 43 insertions(+), 34 deletions(-) diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 2402360e4893..7d07efcfbbdd 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -87,7 +87,7 @@ ("deberta", "DebertaConfig"), ("deberta-v2", "DebertaV2Config"), ("decision_transformer", "DecisionTransformerConfig"), - ("deepseek-v3", "DeepseekV3Config"), + ("deepseek_v3", "DeepseekV3Config"), ("deformable_detr", "DeformableDetrConfig"), ("deit", "DeiTConfig"), ("depth_anything", "DepthAnythingConfig"), @@ -407,7 +407,7 @@ ("deberta", "DeBERTa"), ("deberta-v2", "DeBERTa-v2"), ("decision_transformer", "Decision Transformer"), - ("deepseek-v3", "Deepseek-V3"), + ("deepseek_v3", "Deepseek-V3"), ("deformable_detr", "Deformable DETR"), ("deit", "DeiT"), ("deplot", "DePlot"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 00e4892c0b6a..592dff25db49 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -86,7 +86,7 @@ ("deberta", "DebertaModel"), ("deberta-v2", "DebertaV2Model"), ("decision_transformer", "DecisionTransformerModel"), - ("deepseek-v3", "DeepseekV3Model"), + ("deepseek_v3", "DeepseekV3Model"), ("deformable_detr", "DeformableDetrModel"), ("deit", "DeiTModel"), ("deta", "DetaModel"), @@ -502,7 +502,7 @@ ("ctrl", "CTRLLMHeadModel"), ("data2vec-text", "Data2VecTextForCausalLM"), ("dbrx", "DbrxForCausalLM"), - ("deepseek-v3", "DeepseekV3ForCausalLM"), + ("deepseek_v3", "DeepseekV3ForCausalLM"), ("diffllama", "DiffLlamaForCausalLM"), ("electra", "ElectraForCausalLM"), ("emu3", "Emu3ForCausalLM"), diff --git a/src/transformers/models/auto/tokenization_auto.py b/src/transformers/models/auto/tokenization_auto.py index 8b862069c7d4..c01de1986a53 100644 --- a/src/transformers/models/auto/tokenization_auto.py +++ b/src/transformers/models/auto/tokenization_auto.py @@ -171,7 +171,7 @@ ), ), ( - "deepseek-v3", + "deepseek_v3", ( "LlamaTokenizer" if is_sentencepiece_available() else None, "LlamaTokenizerFast" if is_tokenizers_available() else None, diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 6841442a2229..eaaf170c28d8 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -31,56 +31,62 @@ class DeepseekV3Config(PretrainedConfig): vocab_size (`int`, *optional*, defaults to 129280): Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`DeepseekV3Model`] - hidden_size (`int`, *optional*, defaults to 4096): + hidden_size (`int`, *optional*, defaults to 7168): Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 11008): + intermediate_size (`int`, *optional*, defaults to 18432): Dimension of the MLP representations. - moe_intermediate_size (`int`, *optional*, defaults to 1407): + moe_intermediate_size (`int`, *optional*, defaults to 2048): Dimension of the MoE representations. - num_hidden_layers (`int`, *optional*, defaults to 32): + num_hidden_layers (`int`, *optional*, defaults to 61): Number of hidden layers in the Transformer decoder. num_nextn_predict_layers (`int`, *optional*, defaults to 1): Number of nextn predict layers in the DeepSeekV3 Model. - num_attention_heads (`int`, *optional*, defaults to 32): + num_attention_heads (`int`, *optional*, defaults to 128): Number of attention heads for each attention layer in the Transformer decoder. - n_shared_experts (`int`, *optional*, defaults to None): + num_key_value_heads (`int`, *optional*, defaults to 128): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + n_shared_experts (`int`, *optional*, defaults to 1): Number of shared experts, None means dense model. - n_routed_experts (`int`, *optional*, defaults to None): + n_routed_experts (`int`, *optional*, defaults to 256): Number of routed experts, None means dense model. - routed_scaling_factor (`float`, *optional*, defaults to 1.0): + ep_size (``, *optional*, defaults to 1): + routed_scaling_factor (`float`, *optional*, defaults to 2.5): Scaling factor or routed experts. - topk_method (`str`, *optional*, defaults to `gready`): + kv_lora_rank (``, *optional*, defaults to 512): + q_lora_rank (``, *optional*, defaults to 1536): + qk_rope_head_dim (``, *optional*, defaults to 64): + v_head_dim (``, *optional*, defaults to 128): + qk_nope_head_dim (``, *optional*, defaults to 128): + topk_method (`str`, *optional*, defaults to `"noaux_tc"`): Topk method used in routed gate. - n_group (`int`, *optional*, defaults to None): + n_group (`int`, *optional*, defaults to 8): Number of groups for routed experts. - topk_group (`int`, *optional*, defaults to None): + topk_group (`int`, *optional*, defaults to 4): Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). - num_experts_per_tok (`int`, *optional*, defaults to None): + num_experts_per_tok (`int`, *optional*, defaults to 8): Number of selected experts, None means dense model. moe_layer_freq (`int`, *optional*, defaults to 1): The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. - first_k_dense_replace (`int`, *optional*, defaults to 0): + first_k_dense_replace (`int`, *optional*, defaults to 3): Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). \--k dense layers--/ - norm_topk_prob (`bool`, *optional*, defaults to False): + norm_topk_prob (`bool`, *optional*, defaults to `True`): Whether to normalize the weights of the routed experts. - scoring_func (`str`, *optional*, defaults to 'softmax'): + scoring_func (`str`, *optional*, defaults to `"sigmoid"`): Method of computing expert weights. aux_loss_alpha (`float`, *optional*, defaults to 0.001): Auxiliary loss weight coefficient. - seq_aux = (`bool`, *optional*, defaults to True): Whether to compute the auxiliary loss for each individual sample. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details checkout [this - paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to - `num_attention_heads`. + seq_aux (``, *optional*, defaults to `True`): hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. - max_position_embeddings (`int`, *optional*, defaults to 2048): + max_position_embeddings (`int`, *optional*, defaults to 4096): The maximum sequence length that this model might ever be used with. initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. @@ -91,9 +97,9 @@ class DeepseekV3Config(PretrainedConfig): relevant if `config.is_decoder=True`. pad_token_id (`int`, *optional*): Padding token id. - bos_token_id (`int`, *optional*, defaults to 1): + bos_token_id (`int`, *optional*, defaults to 0): Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 2): + eos_token_id (`int`, *optional*, defaults to 1): End of stream token id. pretraining_tp (`int`, *optional*, defaults to 1): Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this @@ -218,3 +224,6 @@ def __init__( tie_word_embeddings=tie_word_embeddings, **kwargs, ) + + +__all__ = ["DeepseekV3Config"] diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index e49975b36c99..6ceb44e15832 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -1054,8 +1054,8 @@ def forward( ```python >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM - >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") - >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") + >>> model = DeepseekV3ForCausalLM.from_pretrained("deepseek-ai/DeepSeek-V3") + >>> tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-V3") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") From 57088cc53e83f17117cc021336ffee0d05cce661 Mon Sep 17 00:00:00 2001 From: ryan u Date: Wed, 29 Jan 2025 00:18:58 +0900 Subject: [PATCH 14/86] use default docs --- src/transformers/models/deepseek_v3/modeling_deepseek_v3.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 6ceb44e15832..e49975b36c99 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -1054,8 +1054,8 @@ def forward( ```python >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM - >>> model = DeepseekV3ForCausalLM.from_pretrained("deepseek-ai/DeepSeek-V3") - >>> tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-V3") + >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") From 0ef561b9910356ec818ccf0ec23bafade61006a8 Mon Sep 17 00:00:00 2001 From: ryan u Date: Wed, 29 Jan 2025 01:33:52 +0900 Subject: [PATCH 15/86] apply make --- docs/source/en/index.md | 2 +- docs/source/en/model_doc/deepseek_v3.md | 46 ++++++++ .../deepseek_v3/configuration_deepseek_v3.py | 50 ++++---- .../deepseek_v3/modeling_deepseek_v3.py | 8 ++ .../models/deepseek_v3/modular_deepseek_v3.py | 107 ++++++------------ 5 files changed, 119 insertions(+), 94 deletions(-) create mode 100644 docs/source/en/model_doc/deepseek_v3.md diff --git a/docs/source/en/index.md b/docs/source/en/index.md index 9d34c2aaada7..ce9e41783d49 100644 --- a/docs/source/en/index.md +++ b/docs/source/en/index.md @@ -118,7 +118,7 @@ Flax), PyTorch, and/or TensorFlow. | [DeBERTa](model_doc/deberta) | βœ… | βœ… | ❌ | | [DeBERTa-v2](model_doc/deberta-v2) | βœ… | βœ… | ❌ | | [Decision Transformer](model_doc/decision_transformer) | βœ… | ❌ | ❌ | -| [Deepseek-V3](model_doc/deepseek-v3) | βœ… | ❌ | ❌ | +| [Deepseek-V3](model_doc/deepseek_v3) | βœ… | ❌ | ❌ | | [Deformable DETR](model_doc/deformable_detr) | βœ… | ❌ | ❌ | | [DeiT](model_doc/deit) | βœ… | βœ… | ❌ | | [DePlot](model_doc/deplot) | βœ… | ❌ | ❌ | diff --git a/docs/source/en/model_doc/deepseek_v3.md b/docs/source/en/model_doc/deepseek_v3.md new file mode 100644 index 000000000000..596805cdb70f --- /dev/null +++ b/docs/source/en/model_doc/deepseek_v3.md @@ -0,0 +1,46 @@ + + +# DeepSeek-V3 + +## Overview + +The DeepSeek-V3 model was proposed in [DeepSeek-V3 Technical Report](https://arxiv.org/abs/2412.19437) by DeepSeek-AI Team. + +The abstract from the paper is the following: +We present DeepSeek-V3, a strong Mixture-of-Experts (MoE) language model with 671B total parameters with 37B activated for each token. To achieve efficient inference and cost-effective training, DeepSeek-V3 adopts Multi-head Latent Attention (MLA) and DeepSeekMoE architectures, which were thoroughly validated in DeepSeek-V2. Furthermore, DeepSeek-V3 pioneers an auxiliary-loss-free strategy for load balancing and sets a multi-token prediction training objective for stronger performance. We pre-train DeepSeek-V3 on 14.8 trillion diverse and high-quality tokens, followed by Supervised Fine-Tuning and Reinforcement Learning stages to fully harness its capabilities. Comprehensive evaluations reveal that DeepSeek-V3 outperforms other open-source models and achieves performance comparable to leading closed-source models. Despite its excellent performance, DeepSeek-V3 requires only 2.788M H800 GPU hours for its full training. In addition, its training process is remarkably stable. Throughout the entire training process, we did not experience any irrecoverable loss spikes or perform any rollbacks. The model checkpoints are available at https://github.com/deepseek-ai/DeepSeek-V3. + +### Usage tips +The model uses Multi-head Latent Attention (MLA) and DeepSeekMoE architectures for efficient inference and cost-effective training. It employs an auxiliary-loss-free strategy for load balancing and multi-token prediction training objective. The model can be used for various language tasks after being pre-trained on 14.8 trillion tokens and going through Supervised Fine-Tuning and Reinforcement Learning stages. + +## DeepseekV3Config + +[[autodoc]] DeepseekV3Config + +## DeepseekV3Model + +[[autodoc]] DeepseekV3Model + - forward + +## DeepseekV3ForCausalLM + +[[autodoc]] DeepseekV3ForCausalLM + - forward + +## DeepseekV3ForSequenceClassification + +[[autodoc]] DeepseekV3ForSequenceClassification + - forward diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index eaaf170c28d8..1cfd57b4c73f 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -14,19 +14,24 @@ # 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. -""" DeepSeekV3 model configuration """ +"""DeepSeekV3 model configuration""" from ...configuration_utils import PretrainedConfig DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} + + class DeepseekV3Config(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the DeepSeek-V3. + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. + + Args: vocab_size (`int`, *optional*, defaults to 129280): Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the @@ -119,10 +124,13 @@ class DeepseekV3Config(PretrainedConfig): Whether to use a bias in the query, key, value and output projection layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. + ```python >>> from transformers import DeepseekV3Model, DeepseekV3Config + >>> # Initializing a Deepseek-V3 style configuration >>> configuration = DeepseekV3Config() + >>> # Accessing the model configuration >>> configuration = model.config ```""" @@ -135,30 +143,30 @@ def __init__( vocab_size=129280, hidden_size=7168, intermediate_size=18432, - moe_intermediate_size = 2048, + moe_intermediate_size=2048, num_hidden_layers=61, num_nextn_predict_layers=1, num_attention_heads=128, num_key_value_heads=128, - n_shared_experts = 1, - n_routed_experts = 256, - ep_size = 1, - routed_scaling_factor = 2.5, - kv_lora_rank = 512, - q_lora_rank = 1536, - qk_rope_head_dim = 64, - v_head_dim = 128, - qk_nope_head_dim = 128, - topk_method = 'noaux_tc', - n_group = 8, - topk_group = 4, - num_experts_per_tok = 8, - moe_layer_freq = 1, - first_k_dense_replace = 3, - norm_topk_prob = True, - scoring_func = 'sigmoid', - aux_loss_alpha = 0.001, - seq_aux = True, + n_shared_experts=1, + n_routed_experts=256, + ep_size=1, + routed_scaling_factor=2.5, + kv_lora_rank=512, + q_lora_rank=1536, + qk_rope_head_dim=64, + v_head_dim=128, + qk_nope_head_dim=128, + topk_method="noaux_tc", + n_group=8, + topk_group=4, + num_experts_per_tok=8, + moe_layer_freq=1, + first_k_dense_replace=3, + norm_topk_prob=True, + scoring_func="sigmoid", + aux_loss_alpha=0.001, + seq_aux=True, hidden_act="silu", max_position_embeddings=4096, initializer_range=0.02, diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index e49975b36c99..7aa67b70c7af 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -1210,3 +1210,11 @@ def forward( hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions, ) + + +__all__ = [ + "DeepseekV3PreTrainedModel", + "DeepseekV3Model", + "DeepseekV3ForCausalLM", + "DeepseekV3ForSequenceClassification", +] diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 913d35730dc7..11d16eac1cee 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -92,13 +92,9 @@ def __init__(self, config): # 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)) - ) + 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.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) self.reset_parameters() def reset_parameters(self) -> None: @@ -110,55 +106,42 @@ 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 - ) + 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}" - ) + 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) + 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_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 - ) + .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_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}" - ) + 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 + 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. @@ -177,11 +160,8 @@ def __init__(self, config): 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 + 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) @@ -193,18 +173,14 @@ def __init__(self, config): self.ep_rank = 0 self.experts = nn.ModuleList( [ - DeepseekV3MLP( - config, intermediate_size=config.moe_intermediate_size - ) + DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size) 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 - ) + self.shared_experts = DeepseekV3MLP(config=config, intermediate_size=intermediate_size) def forward(self, hidden_states): identity = hidden_states @@ -227,17 +203,9 @@ def moe_infer(self, x, topk_ids, topk_weight): 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] - ) + 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() - ) + 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] ) @@ -246,9 +214,9 @@ def moe_infer(self, x, topk_ids, topk_weight): 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) + 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()): @@ -324,17 +292,11 @@ def __init__(self, config: DeepseekV3Config, layer_idx: Optional[int] = None): 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 - ) + 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_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.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, @@ -344,8 +306,7 @@ def __init__(self, config: DeepseekV3Config, layer_idx: Optional[int] = None): self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank) self.kv_b_proj = nn.Linear( config.kv_lora_rank, - self.num_heads - * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), bias=False, ) @@ -375,14 +336,10 @@ def forward( else: q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) - q_nope, q_pe = torch.split( - q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 - ) + q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - compressed_kv, k_pe = torch.split( - compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) + compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) kv = ( self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) @@ -390,9 +347,7 @@ def forward( .transpose(1, 2) ) - k_nope, value_states = torch.split( - kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 - ) + k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) kv_seq_len = value_states.shape[-2] if past_key_value is not None: if self.layer_idx is None: @@ -482,3 +437,11 @@ class DeepseekV3ForCausalLM(LlamaForCausalLM): class DeepseekV3ForSequenceClassification(LlamaForSequenceClassification): pass + + +__all__ = [ + "DeepseekV3PreTrainedModel", + "DeepseekV3Model", + "DeepseekV3ForCausalLM", + "DeepseekV3ForSequenceClassification", +] From 9a75a56ad2163f060d98d797c7bfdfb69e5dbd11 Mon Sep 17 00:00:00 2001 From: ryan u Date: Wed, 29 Jan 2025 01:41:38 +0900 Subject: [PATCH 16/86] fill type and docstring --- .../deepseek_v3/configuration_deepseek_v3.py | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 1cfd57b4c73f..1f61ebe6b174 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -60,14 +60,20 @@ class DeepseekV3Config(PretrainedConfig): Number of shared experts, None means dense model. n_routed_experts (`int`, *optional*, defaults to 256): Number of routed experts, None means dense model. - ep_size (``, *optional*, defaults to 1): + ep_size (`int`, *optional*, defaults to 1): + Expert parallelism size for distributed training. routed_scaling_factor (`float`, *optional*, defaults to 2.5): Scaling factor or routed experts. - kv_lora_rank (``, *optional*, defaults to 512): - q_lora_rank (``, *optional*, defaults to 1536): - qk_rope_head_dim (``, *optional*, defaults to 64): - v_head_dim (``, *optional*, defaults to 128): - qk_nope_head_dim (``, *optional*, defaults to 128): + kv_lora_rank (`int`, *optional*, defaults to 512): + Rank of the LoRA matrices for key and value projections. + q_lora_rank (`int`, *optional*, defaults to 1536): + Rank of the LoRA matrices for query projections. + qk_rope_head_dim (`int`, *optional*, defaults to 64): + Dimension of the query/key heads that use rotary position embeddings. + v_head_dim (`int`, *optional*, defaults to 128): + Dimension of the value heads. + qk_nope_head_dim (`int`, *optional*, defaults to 128): + Dimension of the query/key heads that don't use rotary position embeddings. topk_method (`str`, *optional*, defaults to `"noaux_tc"`): Topk method used in routed gate. n_group (`int`, *optional*, defaults to 8): @@ -88,7 +94,8 @@ class DeepseekV3Config(PretrainedConfig): aux_loss_alpha (`float`, *optional*, defaults to 0.001): Auxiliary loss weight coefficient. Whether to compute the auxiliary loss for each individual sample. - seq_aux (``, *optional*, defaults to `True`): + seq_aux (`bool`, *optional*, defaults to `True`): + Whether to compute auxiliary loss at sequence level. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. max_position_embeddings (`int`, *optional*, defaults to 4096): From cdf83e4591ab3654ceb5de595808699db0144bac Mon Sep 17 00:00:00 2001 From: ryan u Date: Thu, 30 Jan 2025 00:28:39 +0900 Subject: [PATCH 17/86] add rope_config_validation --- .../models/deepseek_v3/configuration_deepseek_v3.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 1f61ebe6b174..7011fdeeb61a 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -17,6 +17,7 @@ """DeepSeekV3 model configuration""" from ...configuration_utils import PretrainedConfig +from ...modeling_rope_utils import rope_config_validation DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} @@ -217,6 +218,7 @@ def __init__( self.scoring_func = scoring_func self.aux_loss_alpha = aux_loss_alpha self.seq_aux = seq_aux + # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads @@ -231,6 +233,11 @@ def __init__( self.rope_scaling = rope_scaling self.attention_bias = attention_bias self.attention_dropout = attention_dropout + # Validate the correctness of rotary position embeddings parameters + # BC: if there is a 'type' field, copy it it to 'rope_type'. + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + rope_config_validation(self) super().__init__( pad_token_id=pad_token_id, From 51990b943686fbc289fd9c41b45e2f65f8061b6d Mon Sep 17 00:00:00 2001 From: ryan u Date: Thu, 30 Jan 2025 00:29:16 +0900 Subject: [PATCH 18/86] use custom DeepseekV3MLP --- .../deepseek_v3/modeling_deepseek_v3.py | 13 +++++++------ .../models/deepseek_v3/modular_deepseek_v3.py | 19 ++++++++++++++++--- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 7aa67b70c7af..c823a24af6db 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -119,14 +119,15 @@ def forward(self, x, position_ids): class DeepseekV3MLP(nn.Module): - def __init__(self, config): + def __init__(self, config, hidden_size=None, intermediate_size=None): super().__init__() self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + 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 + + 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): diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 11d16eac1cee..ee648fae517e 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -8,6 +8,7 @@ import torch.utils.checkpoint from torch import nn +from ...activations import ACT2FN from ...cache_utils import Cache from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_utils import ALL_ATTENTION_FUNCTIONS @@ -17,7 +18,6 @@ LlamaDecoderLayer, LlamaForCausalLM, LlamaForSequenceClassification, - LlamaMLP, LlamaModel, LlamaPreTrainedModel, LlamaRMSNorm, @@ -72,8 +72,21 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): return q_embed, k_embed -class DeepseekV3MLP(LlamaMLP): - pass +class DeepseekV3MLP(nn.Module): + def __init__(self, config, hidden_size=None, intermediate_size=None): + 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 + + 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 class MoEGate(nn.Module): From f4f0ebd81cef3ecf140bab89fcd931b57b89b8a1 Mon Sep 17 00:00:00 2001 From: ryan u Date: Thu, 30 Jan 2025 20:00:39 +0900 Subject: [PATCH 19/86] hold code only for checkpoints congifuration; remove redundant --- .../deepseek_v3/configuration_deepseek_v3.py | 36 +- .../deepseek_v3/modeling_deepseek_v3.py | 372 +++++++--------- .../models/deepseek_v3/modular_deepseek_v3.py | 411 ++++++++---------- 3 files changed, 343 insertions(+), 476 deletions(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 7011fdeeb61a..8d1c56343638 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -45,8 +45,6 @@ class DeepseekV3Config(PretrainedConfig): Dimension of the MoE representations. num_hidden_layers (`int`, *optional*, defaults to 61): Number of hidden layers in the Transformer decoder. - num_nextn_predict_layers (`int`, *optional*, defaults to 1): - Number of nextn predict layers in the DeepSeekV3 Model. num_attention_heads (`int`, *optional*, defaults to 128): Number of attention heads for each attention layer in the Transformer decoder. num_key_value_heads (`int`, *optional*, defaults to 128): @@ -58,11 +56,9 @@ class DeepseekV3Config(PretrainedConfig): paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `num_attention_heads`. n_shared_experts (`int`, *optional*, defaults to 1): - Number of shared experts, None means dense model. + Number of shared experts. n_routed_experts (`int`, *optional*, defaults to 256): - Number of routed experts, None means dense model. - ep_size (`int`, *optional*, defaults to 1): - Expert parallelism size for distributed training. + Number of routed experts. routed_scaling_factor (`float`, *optional*, defaults to 2.5): Scaling factor or routed experts. kv_lora_rank (`int`, *optional*, defaults to 512): @@ -75,28 +71,20 @@ class DeepseekV3Config(PretrainedConfig): Dimension of the value heads. qk_nope_head_dim (`int`, *optional*, defaults to 128): Dimension of the query/key heads that don't use rotary position embeddings. - topk_method (`str`, *optional*, defaults to `"noaux_tc"`): - Topk method used in routed gate. n_group (`int`, *optional*, defaults to 8): Number of groups for routed experts. topk_group (`int`, *optional*, defaults to 4): Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). num_experts_per_tok (`int`, *optional*, defaults to 8): Number of selected experts, None means dense model. - moe_layer_freq (`int`, *optional*, defaults to 1): - The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. first_k_dense_replace (`int`, *optional*, defaults to 3): Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). \--k dense layers--/ norm_topk_prob (`bool`, *optional*, defaults to `True`): Whether to normalize the weights of the routed experts. - scoring_func (`str`, *optional*, defaults to `"sigmoid"`): - Method of computing expert weights. aux_loss_alpha (`float`, *optional*, defaults to 0.001): Auxiliary loss weight coefficient. Whether to compute the auxiliary loss for each individual sample. - seq_aux (`bool`, *optional*, defaults to `True`): - Whether to compute auxiliary loss at sequence level. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. max_position_embeddings (`int`, *optional*, defaults to 4096): @@ -145,6 +133,12 @@ class DeepseekV3Config(PretrainedConfig): model_type = "deepseek_v3" keys_to_ignore_at_inference = ["past_key_values"] + # Default tensor parallel plan for base model `DeepseekV3Model` + base_model_tp_plan = { + "layers.*.gate_proj": "colwise", + "layers.*.up_proj": "colwise", + "layers.*.down_proj": "rowwise", + } def __init__( self, @@ -153,28 +147,22 @@ def __init__( intermediate_size=18432, moe_intermediate_size=2048, num_hidden_layers=61, - num_nextn_predict_layers=1, num_attention_heads=128, num_key_value_heads=128, n_shared_experts=1, n_routed_experts=256, - ep_size=1, routed_scaling_factor=2.5, kv_lora_rank=512, q_lora_rank=1536, qk_rope_head_dim=64, v_head_dim=128, qk_nope_head_dim=128, - topk_method="noaux_tc", n_group=8, topk_group=4, num_experts_per_tok=8, - moe_layer_freq=1, first_k_dense_replace=3, norm_topk_prob=True, - scoring_func="sigmoid", aux_loss_alpha=0.001, - seq_aux=True, hidden_act="silu", max_position_embeddings=4096, initializer_range=0.02, @@ -197,27 +185,23 @@ def __init__( self.intermediate_size = intermediate_size self.moe_intermediate_size = moe_intermediate_size self.num_hidden_layers = num_hidden_layers - self.num_nextn_predict_layers = num_nextn_predict_layers self.num_attention_heads = num_attention_heads self.n_shared_experts = n_shared_experts self.n_routed_experts = n_routed_experts - self.ep_size = ep_size self.routed_scaling_factor = routed_scaling_factor self.kv_lora_rank = kv_lora_rank self.q_lora_rank = q_lora_rank self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim - self.topk_method = topk_method + self.q_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.head_dim = qk_rope_head_dim self.n_group = n_group self.topk_group = topk_group self.num_experts_per_tok = num_experts_per_tok - self.moe_layer_freq = moe_layer_freq self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob - self.scoring_func = scoring_func self.aux_loss_alpha = aux_loss_alpha - self.seq_aux = seq_aux # for backward compatibility if num_key_value_heads is None: diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index c823a24af6db..bb5aad4077df 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -7,9 +7,7 @@ import math from typing import Callable, List, Optional, Tuple, Union -import numpy as np import torch -import torch.distributed as dist import torch.nn.functional as F from torch import nn @@ -142,63 +140,41 @@ def __init__(self, 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)) + self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) + self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) 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 + batch_size, seq_length = hidden_states.shape[:-1] + hidden_states = hidden_states.view(-1, self.config.hidden_size) + logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + + scores = logits.sigmoid() + scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) + group_scores = ( + scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) + .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(batch_size * seq_length, self.n_group, self.n_routed_experts // self.n_group) + .reshape(-1, self.n_routed_experts) + ) # [n, e] + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] + _, topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) + topk_weights = scores.gather(1, topk_indices) + if self.norm_topk_prob: + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator + topk_weights = topk_weights * self.routed_scaling_factor # must multiply the scaling factor + return topk_indices, topk_weights class DeepseekV3MoE(nn.Module): @@ -209,116 +185,84 @@ class DeepseekV3MoE(nn.Module): 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) - for i in range(config.n_routed_experts) - ] - ) + self.experts = nn.ModuleList( + [ + DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size) + for _ 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) + 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) + topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-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) + y = self.moe_infer(hidden_states, topk_indices, topk_weights).view(*orig_shape) + 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 + def moe_infer(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + """ + Perform inference using a Mixture of Experts (MoE) model. + + Args: + hidden_states (torch.Tensor): Input hidden states. + topk_indices (torch.Tensor): Indices of the top-k experts for each token. + topk_weights (torch.Tensor): Weights associated with the top-k experts. + + Returns: + torch.Tensor: Output of the MoE model. + """ + num_experts = len(self.experts) + batch_size, num_topk = topk_indices.shape + + # Count the number of tokens assigned to each expert + expert_counts = topk_indices.new_zeros((batch_size, num_experts)) + expert_counts.scatter_(1, topk_indices, 1) + tokens_per_expert = expert_counts.sum(dim=0) + + # Sort tokens by their assigned expert + sorted_indices = topk_indices.view(-1).argsort() + sorted_tokens = hidden_states[sorted_indices // num_topk] 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 + # Process tokens through their assigned experts + expert_outputs = [] + current_pos = 0 + + for expert_idx, num_tokens in enumerate(tokens_per_expert): 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 + + next_pos = current_pos + num_tokens + expert = self.experts[expert_idx] + expert_tokens = sorted_tokens[current_pos:next_pos] + expert_outputs.append(expert(expert_tokens)) + current_pos = next_pos + + # Combine the outputs from all experts + expert_outputs = torch.cat(expert_outputs, dim=0) if expert_outputs else sorted_tokens.new_empty(0) + + # Reorder the outputs to match the original token sequence + reordered_outputs = torch.empty_like(expert_outputs) + reordered_outputs[sorted_indices] = expert_outputs + + # Reshape and apply the expert weights + reordered_outputs = reordered_outputs.view(batch_size, num_topk, -1).type(topk_weights.dtype) + moe_output = torch.matmul(topk_weights.unsqueeze(1), reordered_outputs) + moe_output = moe_output.sum(dim=1).type(hidden_states.dtype) + return moe_output + + +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) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: @@ -359,16 +303,14 @@ def eager_attention_forward( return attn_output, attn_weights -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) +def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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. @@ -379,15 +321,22 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): 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 + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_length, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_length, 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. + the shape [batch_size, seq_length, 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.unsqueeze(unsqueeze_dim) sin = sin.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 @@ -396,111 +345,87 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): class DeepseekV3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" - def __init__(self, config: DeepseekV3Config, layer_idx: Optional[int] = None): + def __init__(self, config: DeepseekV3Config, layer_idx: int): 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.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads 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.q_head_dim = config.q_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.q_a_proj = nn.Linear(config.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, - config.kv_lora_rank + config.qk_rope_head_dim, + config.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, bias=config.attention_bias, ) - self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank) + self.kv_a_layernorm = DeepseekV3RMSNorm(self.kv_lora_rank) self.kv_b_proj = nn.Linear( - config.kv_lora_rank, + self.kv_lora_rank, self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), bias=False, ) self.o_proj = nn.Linear( self.num_heads * self.v_head_dim, - self.hidden_size, + config.hidden_size, bias=config.attention_bias, ) - self.rotary_emb = DeepseekV3RotaryEmbedding( - config=self.config, - ) + self.scaling = self.q_head_dim ** (-0.5) + if self.config.rope_scaling is not None: + mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = self.config.rope_scaling["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.scaling = self.scaling * mscale * mscale def forward( self, hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], - position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - bsz, q_len, _ = hidden_states.size() + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, self.num_heads, -1) + batch_size, seq_length = input_shape - if self.q_lora_rank is None: - q = self.q_proj(hidden_states) - else: - q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) - q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) - kv = ( - self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) - .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) - .transpose(1, 2) - ) + k_pe = k_pe.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(hidden_shape).transpose(1, 2) k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - kv_seq_len = value_states.shape[-2] - if past_key_value is not None: - if self.layer_idx is None: - raise ValueError( - f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " - "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " - "with a layer index." - ) - kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) - cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) - q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + cos, sin = position_embeddings + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin) - query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states = k_pe.new_empty(batch_size, self.num_heads, seq_length, self.q_head_dim) query_states[:, :, :, : self.qk_nope_head_dim] = q_nope query_states[:, :, :, self.qk_nope_head_dim :] = q_pe - key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states = k_pe.new_empty(batch_size, self.num_heads, seq_length, self.q_head_dim) key_states[:, :, :, : self.qk_nope_head_dim] = k_nope key_states[:, :, :, self.qk_nope_head_dim :] = k_pe - if self.q_head_dim != self.v_head_dim: + if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) if past_key_value is not None: @@ -528,9 +453,12 @@ def forward( scaling=self.scaling, **kwargs, ) - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) - attn_output = self.o_proj(attn_output) + if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) return attn_output, attn_weights @@ -541,15 +469,11 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): self.self_attn = DeepseekV3Attention(config=config, layer_idx=layer_idx) - self.mlp = ( - DeepseekV3MoE(config) - if ( - config.n_routed_experts is not None - and layer_idx >= config.first_k_dense_replace - and layer_idx % config.moe_layer_freq == 0 - ) - else DeepseekV3MLP(config) - ) + if layer_idx >= config.first_k_dense_replace: + self.mlp = DeepseekV3MoE(config) + else: + self.mlp = DeepseekV3MLP(config) + self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index ee648fae517e..8b20f068d489 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -1,9 +1,7 @@ import math from typing import Callable, Optional, Tuple -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 @@ -15,7 +13,6 @@ from ...processing_utils import Unpack from ...utils import logging from ..llama.modeling_llama import ( - LlamaDecoderLayer, LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, @@ -23,6 +20,7 @@ LlamaRMSNorm, LlamaRotaryEmbedding, eager_attention_forward, + rotate_half, ) from .configuration_deepseek_v3 import DeepseekV3Config @@ -38,16 +36,14 @@ class DeepseekV3RotaryEmbedding(LlamaRotaryEmbedding): pass -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) +def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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. @@ -58,15 +54,22 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): 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 + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_length, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_length, 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. + the shape [batch_size, seq_length, 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.unsqueeze(unsqueeze_dim) sin = sin.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 @@ -96,63 +99,41 @@ def __init__(self, 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)) + self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) + self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) 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 + batch_size, seq_length = hidden_states.shape[:-1] + hidden_states = hidden_states.view(-1, self.config.hidden_size) + logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + + scores = logits.sigmoid() + scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) + group_scores = ( + scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) + .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(batch_size * seq_length, self.n_group, self.n_routed_experts // self.n_group) + .reshape(-1, self.n_routed_experts) + ) # [n, e] + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] + _, topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) + topk_weights = scores.gather(1, topk_indices) + if self.norm_topk_prob: + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator + topk_weights = topk_weights * self.routed_scaling_factor # must multiply the scaling factor + return topk_indices, topk_weights class DeepseekV3MoE(nn.Module): @@ -163,226 +144,163 @@ class DeepseekV3MoE(nn.Module): 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) - for i in range(config.n_routed_experts) - ] - ) + self.experts = nn.ModuleList( + [ + DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size) + for _ 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) + 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) + topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-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) + y = self.moe_infer(hidden_states, topk_indices, topk_weights).view(*orig_shape) + 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 + def moe_infer(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + """ + Perform inference using a Mixture of Experts (MoE) model. + + Args: + hidden_states (torch.Tensor): Input hidden states. + topk_indices (torch.Tensor): Indices of the top-k experts for each token. + topk_weights (torch.Tensor): Weights associated with the top-k experts. + + Returns: + torch.Tensor: Output of the MoE model. + """ + num_experts = len(self.experts) + batch_size, num_topk = topk_indices.shape + + # Count the number of tokens assigned to each expert + expert_counts = topk_indices.new_zeros((batch_size, num_experts)) + expert_counts.scatter_(1, topk_indices, 1) + tokens_per_expert = expert_counts.sum(dim=0) + + # Sort tokens by their assigned expert + sorted_indices = topk_indices.view(-1).argsort() + sorted_tokens = hidden_states[sorted_indices // num_topk] 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 + # Process tokens through their assigned experts + expert_outputs = [] + current_pos = 0 + + for expert_idx, num_tokens in enumerate(tokens_per_expert): 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 + + next_pos = current_pos + num_tokens + expert = self.experts[expert_idx] + expert_tokens = sorted_tokens[current_pos:next_pos] + expert_outputs.append(expert(expert_tokens)) + current_pos = next_pos + + # Combine the outputs from all experts + expert_outputs = torch.cat(expert_outputs, dim=0) if expert_outputs else sorted_tokens.new_empty(0) + + # Reorder the outputs to match the original token sequence + reordered_outputs = torch.empty_like(expert_outputs) + reordered_outputs[sorted_indices] = expert_outputs + + # Reshape and apply the expert weights + reordered_outputs = reordered_outputs.view(batch_size, num_topk, -1).type(topk_weights.dtype) + moe_output = torch.matmul(topk_weights.unsqueeze(1), reordered_outputs) + moe_output = moe_output.sum(dim=1).type(hidden_states.dtype) + return moe_output class DeepseekV3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" - def __init__(self, config: DeepseekV3Config, layer_idx: Optional[int] = None): + def __init__(self, config: DeepseekV3Config, layer_idx: int): 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.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads 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.q_head_dim = config.q_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.q_a_proj = nn.Linear(config.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, - config.kv_lora_rank + config.qk_rope_head_dim, + config.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, bias=config.attention_bias, ) - self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank) + self.kv_a_layernorm = DeepseekV3RMSNorm(self.kv_lora_rank) self.kv_b_proj = nn.Linear( - config.kv_lora_rank, + self.kv_lora_rank, self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), bias=False, ) self.o_proj = nn.Linear( self.num_heads * self.v_head_dim, - self.hidden_size, + config.hidden_size, bias=config.attention_bias, ) - self.rotary_emb = DeepseekV3RotaryEmbedding( - config=self.config, - ) + self.scaling = self.q_head_dim ** (-0.5) + if self.config.rope_scaling is not None: + mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = self.config.rope_scaling["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.scaling = self.scaling * mscale * mscale def forward( self, hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], - position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - bsz, q_len, _ = hidden_states.size() + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, self.num_heads, -1) + batch_size, seq_length = input_shape - if self.q_lora_rank is None: - q = self.q_proj(hidden_states) - else: - q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) - q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) - kv = ( - self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) - .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) - .transpose(1, 2) - ) + k_pe = k_pe.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(hidden_shape).transpose(1, 2) k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - kv_seq_len = value_states.shape[-2] - if past_key_value is not None: - if self.layer_idx is None: - raise ValueError( - f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " - "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " - "with a layer index." - ) - kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) - cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) - q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + cos, sin = position_embeddings + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin) - query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states = k_pe.new_empty(batch_size, self.num_heads, seq_length, self.q_head_dim) query_states[:, :, :, : self.qk_nope_head_dim] = q_nope query_states[:, :, :, self.qk_nope_head_dim :] = q_pe - key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states = k_pe.new_empty(batch_size, self.num_heads, seq_length, self.q_head_dim) key_states[:, :, :, : self.qk_nope_head_dim] = k_nope key_states[:, :, :, self.qk_nope_head_dim :] = k_pe - if self.q_head_dim != self.v_head_dim: + if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) if past_key_value is not None: @@ -410,31 +328,72 @@ def forward( scaling=self.scaling, **kwargs, ) - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) - attn_output = self.o_proj(attn_output) + if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) return attn_output, attn_weights -class DeepseekV3DecoderLayer(LlamaDecoderLayer): +class DeepseekV3DecoderLayer(nn.Module): def __init__(self, config: DeepseekV3Config, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = DeepseekV3Attention(config=config, layer_idx=layer_idx) - self.mlp = ( - DeepseekV3MoE(config) - if ( - config.n_routed_experts is not None - and layer_idx >= config.first_k_dense_replace - and layer_idx % config.moe_layer_freq == 0 - ) - else DeepseekV3MLP(config) - ) + if layer_idx >= config.first_k_dense_replace: + self.mlp = DeepseekV3MoE(config) + else: + self.mlp = DeepseekV3MLP(config) + self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): pass From 4b72b30bb26b948cff5f1b7ea54deafad5a64415 Mon Sep 17 00:00:00 2001 From: ryan u Date: Thu, 30 Jan 2025 20:02:08 +0900 Subject: [PATCH 20/86] revise rope yarn for DeepSeek variation --- src/transformers/modeling_rope_utils.py | 36 ++++++++++++++++++++----- 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/src/transformers/modeling_rope_utils.py b/src/transformers/modeling_rope_utils.py index b2d343e0237f..beb0e9e2f793 100644 --- a/src/transformers/modeling_rope_utils.py +++ b/src/transformers/modeling_rope_utils.py @@ -189,13 +189,31 @@ def _compute_yarn_parameters( partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) dim = int(head_dim * partial_rotary_factor) - max_position_embeddings = config.max_position_embeddings factor = config.rope_scaling["factor"] + attention_factor = config.rope_scaling.get("attention_factor") + mscale = config.rope_scaling.get("mscale") + mscale_all_dim = config.rope_scaling.get("mscale_all_dim") + + # NOTE: DeekSeek-V3 (and potentially other models) modify `max_position_embeddings` and have a + # `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two + # values to compute the default attention scaling factor, instead of using `factor`. + if "original_max_position_embeddings" in config.rope_scaling: + original_max_position_embeddings = config.rope_scaling["original_max_position_embeddings"] + factor = config.max_position_embeddings / original_max_position_embeddings + else: + original_max_position_embeddings = config.max_position_embeddings + + def get_mscale(scale, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 # Sets the attention factor as suggested in the paper - attention_factor = config.rope_scaling.get("attention_factor") if attention_factor is None: - attention_factor = 0.1 * math.log(factor) + 1.0 + if mscale and mscale_all_dim: + attention_factor = float(get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim)) + else: + attention_factor = get_mscale(factor) # Optional config options # beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly) @@ -227,7 +245,7 @@ def linear_ramp_factor(min, max, dim): inv_freq_extrapolation = 1.0 / pos_freqs inv_freq_interpolation = 1.0 / (factor * pos_freqs) - low, high = find_correction_range(beta_fast, beta_slow, dim, base, max_position_embeddings) + low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_max_position_embeddings) # Get n-dimensional rotational scaling corrected for extrapolation inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).float().to(device) @@ -235,7 +253,6 @@ def linear_ramp_factor(min, max, dim): inv_freq_interpolation * (1 - inv_freq_extrapolation_factor) + inv_freq_extrapolation * inv_freq_extrapolation_factor ) - return inv_freq, attention_factor @@ -425,7 +442,14 @@ def _validate_yarn_parameters(config: PretrainedConfig, ignore_keys: Optional[se rope_scaling = config.rope_scaling rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type" required_keys = {"rope_type", "factor"} - optional_keys = {"attention_factor", "beta_fast", "beta_slow"} + optional_keys = { + "attention_factor", + "beta_fast", + "beta_slow", + "original_max_position_embeddings", + "mscale", + "mscale_all_dim", + } received_keys = set(rope_scaling.keys()) _check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys) From 6792cb52ce0a40097709fb4f3eed471ab267352c Mon Sep 17 00:00:00 2001 From: ryan u Date: Thu, 30 Jan 2025 20:29:10 +0900 Subject: [PATCH 21/86] rename DeepSeek-V3 --- docs/source/en/index.md | 2 +- src/transformers/models/auto/configuration_auto.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/en/index.md b/docs/source/en/index.md index ce9e41783d49..af086f68163c 100644 --- a/docs/source/en/index.md +++ b/docs/source/en/index.md @@ -118,7 +118,7 @@ Flax), PyTorch, and/or TensorFlow. | [DeBERTa](model_doc/deberta) | βœ… | βœ… | ❌ | | [DeBERTa-v2](model_doc/deberta-v2) | βœ… | βœ… | ❌ | | [Decision Transformer](model_doc/decision_transformer) | βœ… | ❌ | ❌ | -| [Deepseek-V3](model_doc/deepseek_v3) | βœ… | ❌ | ❌ | +| [DeepSeek-V3](model_doc/deepseek_v3) | βœ… | ❌ | ❌ | | [Deformable DETR](model_doc/deformable_detr) | βœ… | ❌ | ❌ | | [DeiT](model_doc/deit) | βœ… | βœ… | ❌ | | [DePlot](model_doc/deplot) | βœ… | ❌ | ❌ | diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 7d07efcfbbdd..69c1939afc76 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -407,7 +407,7 @@ ("deberta", "DeBERTa"), ("deberta-v2", "DeBERTa-v2"), ("decision_transformer", "Decision Transformer"), - ("deepseek_v3", "Deepseek-V3"), + ("deepseek_v3", "DeepSeek-V3"), ("deformable_detr", "Deformable DETR"), ("deit", "DeiT"), ("deplot", "DePlot"), From 3bf3b323c35d513026318fc8cdc8053abc1d655e Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 31 Jan 2025 11:50:29 +0100 Subject: [PATCH 22/86] some refactoring --- .../deepseek_v3/modeling_deepseek_v3.py | 395 +++++++++++++----- .../models/deepseek_v3/modular_deepseek_v3.py | 159 ++++--- 2 files changed, 360 insertions(+), 194 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index bb5aad4077df..f0ee6bf50f83 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -12,11 +12,17 @@ from torch import nn from ...activations import ACT2FN -from ...cache_utils import Cache, DynamicCache, StaticCache +from ...cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache from ...generation import GenerationMixin from ...modeling_attn_mask_utils import AttentionMaskConverter from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast +from ...modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, + SequenceClassifierOutputWithPast, +) from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -133,7 +139,7 @@ def forward(self, x): return down_proj -class MoEGate(nn.Module): +class DeepseekV3TopkRouter(nn.Module): def __init__(self, config): super().__init__() self.config = config @@ -142,7 +148,6 @@ def __init__(self, config): self.routed_scaling_factor = config.routed_scaling_factor self.n_group = config.n_group self.topk_group = config.topk_group - self.norm_topk_prob = config.norm_topk_prob self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) @@ -150,9 +155,9 @@ def __init__(self, config): def forward(self, hidden_states): batch_size, seq_length = hidden_states.shape[:-1] hidden_states = hidden_states.view(-1, self.config.hidden_size) - logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) - scores = logits.sigmoid() + scores = router_logits.sigmoid() scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) @@ -170,11 +175,10 @@ def forward(self, hidden_states): scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] _, topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) topk_weights = scores.gather(1, topk_indices) - if self.norm_topk_prob: - denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 - topk_weights /= denominator + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator topk_weights = topk_weights * self.routed_scaling_factor # must multiply the scaling factor - return topk_indices, topk_weights + return topk_indices, topk_weights, router_logits class DeepseekV3MoE(nn.Module): @@ -191,24 +195,21 @@ def __init__(self, config): for _ in range(config.n_routed_experts) ] ) - self.gate = MoEGate(config) - intermediate_size = config.moe_intermediate_size * config.n_shared_experts - self.shared_experts = DeepseekV3MLP(config=config, intermediate_size=intermediate_size) + self.gate = DeepseekV3TopkRouter(config) + self.shared_experts = DeepseekV3MLP(config=config, intermediate_size=config.moe_intermediate_size) def forward(self, hidden_states): - identity = hidden_states + residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) + topk_indices, topk_weights, router_logits = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - y = self.moe_infer(hidden_states, topk_indices, topk_weights).view(*orig_shape) - y = y + self.shared_experts(identity) - return y + hidden_states = self.moe_infer(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = hidden_states + self.shared_experts(residuals) + return hidden_states, router_logits - @torch.no_grad() def moe_infer(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): """ Perform inference using a Mixture of Experts (MoE) model. - Args: hidden_states (torch.Tensor): Input hidden states. topk_indices (torch.Tensor): Indices of the top-k experts for each token. @@ -219,20 +220,20 @@ def moe_infer(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, top """ num_experts = len(self.experts) batch_size, num_topk = topk_indices.shape + with torch.no_grad(): + # Count the number of tokens assigned to each expert + expert_counts = topk_indices.new_zeros((batch_size, num_experts)) + expert_counts.scatter_(1, topk_indices, 1) + tokens_per_expert = expert_counts.sum(dim=0) - # Count the number of tokens assigned to each expert - expert_counts = topk_indices.new_zeros((batch_size, num_experts)) - expert_counts.scatter_(1, topk_indices, 1) - tokens_per_expert = expert_counts.sum(dim=0) - - # Sort tokens by their assigned expert - sorted_indices = topk_indices.view(-1).argsort() - sorted_tokens = hidden_states[sorted_indices // num_topk] - tokens_per_expert = tokens_per_expert.cpu().numpy() + # Sort tokens by their assigned expert + sorted_indices = topk_indices.view(-1).argsort() + sorted_tokens = hidden_states[sorted_indices // num_topk] + tokens_per_expert = tokens_per_expert.cpu().numpy() - # Process tokens through their assigned experts - expert_outputs = [] - current_pos = 0 + # Process tokens through their assigned experts + expert_outputs = [] + current_pos = 0 for expert_idx, num_tokens in enumerate(tokens_per_expert): if num_tokens == 0: @@ -265,6 +266,33 @@ def rotate_half(x): return torch.cat((-x2, x1), dim=-1) +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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`, *optional*): + Deprecated and unused. + 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.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + 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, @@ -309,39 +337,6 @@ def yarn_get_mscale(scale=1, mscale=1): return 0.1 * mscale * math.log(scale) + 1.0 -def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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`, *optional*): - Deprecated and unused. - 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_length, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_length, 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_length, 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.unsqueeze(unsqueeze_dim) - sin = sin.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 DeepseekV3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -404,26 +399,22 @@ def forward( hidden_shape = (*input_shape, self.num_heads, -1) batch_size, seq_length = input_shape - q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) - q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) + q_rot, q_pass = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_pe = k_pe.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2) - kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(hidden_shape).transpose(1, 2) + k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(hidden_shape).transpose(1, 2) + k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - cos, sin = position_embeddings - q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin) + k_rot = k_rot.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2) - query_states = k_pe.new_empty(batch_size, self.num_heads, seq_length, self.q_head_dim) - query_states[:, :, :, : self.qk_nope_head_dim] = q_nope - query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + cos, sin = position_embeddings + q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - key_states = k_pe.new_empty(batch_size, self.num_heads, seq_length, self.q_head_dim) - key_states[:, :, :, : self.qk_nope_head_dim] = k_nope - key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + query_states = torch.cat(q_rot, q_pass, dim=-1) + key_states = torch.cat(k_rot, k_pass, dim=-1) if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) @@ -484,6 +475,7 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC @@ -511,11 +503,19 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) + + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = (torch.zeros((1,), device=hidden_states.device, dtype=torch.int64),) + hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) + if output_router_logits: + outputs += (router_logits,) return outputs @@ -642,6 +642,17 @@ def _init_weights(self, module): """ +def permute_for_rope(input_tensor, n_heads, dim1, dim2): + """ + When you go from the complex ROPE formulation to sin and cos one, you need + to permute the query and key weights (to avoid doing it on the fly) + """ + input_tensor = input_tensor.reshape(dim1, dim2) + input_tensor = input_tensor.view(n_heads, dim1 // n_heads // 2, 2, dim2) + input_tensor = input_tensor.transpose(1, 2).reshape(dim1, dim2) + return input_tensor + + @add_start_docstrings( "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", DEEPSEEK_V3_START_DOCSTRING, @@ -654,7 +665,7 @@ class DeepseekV3Model(DeepseekV3PreTrainedModel): config: DeepseekV3Config """ - def __init__(self, config: DeepseekV3Config): + def __init__(self, config): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size @@ -666,6 +677,7 @@ def __init__(self, config: DeepseekV3Config): self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = DeepseekV3RotaryEmbedding(config=config) self.gradient_checkpointing = False + self._register_load_state_dict_pre_hook(self.load_hook) # Initialize weights and apply final processing self.post_init() @@ -682,43 +694,48 @@ def forward( input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, **flash_attn_kwargs: Unpack[FlashAttentionKwargs], ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False if use_cache and past_key_values is None: past_key_values = DynamicCache() + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) - if position_ids is None: position_ids = cache_position.unsqueeze(0) @@ -734,8 +751,9 @@ def forward( # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None - for decoder_layer in self.layers[: self.config.num_hidden_layers]: + for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) @@ -747,6 +765,7 @@ def forward( position_ids, past_key_values, output_attentions, + output_router_logits, use_cache, cache_position, position_embeddings, @@ -758,6 +777,7 @@ def forward( position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, + output_router_logits=output_router_logits, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -769,17 +789,21 @@ def forward( if output_attentions: all_self_attns += (layer_outputs[1],) + if output_router_logits: + all_router_logits += (layer_outputs[-1],) + hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) - output = BaseModelOutputWithPast( + output = MoeModelOutputWithPast( last_hidden_state=hidden_states, - past_key_values=past_key_values if use_cache else None, + past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_self_attns, + router_logits=all_router_logits, ) return output if return_dict else output.to_tuple() @@ -792,7 +816,15 @@ def _update_causal_mask( output_attentions: bool, ): if self.config._attn_implementation == "flash_attention_2": - if attention_mask is not None and (attention_mask == 0.0).any(): + if attention_mask is not None and past_key_values is not None: + is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of DeepseekV3. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if attention_mask is not None and 0.0 in attention_mask: return attention_mask return None @@ -801,21 +833,30 @@ def _update_causal_mask( # to infer the attention mask. past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 using_static_cache = isinstance(past_key_values, StaticCache) + using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + if ( + self.config._attn_implementation == "sdpa" + and not (using_static_cache or using_sliding_window_cache) + and not output_attentions + ): if AttentionMaskConverter._ignore_causal_mask_sdpa( attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, is_training=self.training, ): return None dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min sequence_length = input_tensor.shape[1] - if using_static_cache: + # SlidingWindowCache or StaticCache + if using_sliding_window_cache or using_static_cache: target_length = past_key_values.get_max_cache_shape() + # DynamicCache or no cache else: target_length = ( attention_mask.shape[-1] @@ -832,6 +873,8 @@ def _update_causal_mask( device=device, cache_position=cache_position, batch_size=input_tensor.shape[0], + config=self.config, + past_key_values=past_key_values, ) if ( @@ -843,7 +886,6 @@ def _update_causal_mask( # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) return causal_mask @@ -857,7 +899,8 @@ def _prepare_4d_causal_attention_mask_with_cache_position( device: torch.device, cache_position: torch.Tensor, batch_size: int, - **kwargs, + config: DeepseekV3Config, + past_key_values: Cache, ): """ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape @@ -865,13 +908,11 @@ def _prepare_4d_causal_attention_mask_with_cache_position( Args: attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. sequence_length (`int`): The sequence length being processed. target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. + The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. dtype (`torch.dtype`): The dtype to use for the 4D attention mask. device (`torch.device`): @@ -880,6 +921,10 @@ def _prepare_4d_causal_attention_mask_with_cache_position( Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): Batch size. + config (`DeepseekV3Config`): + The model's configuration class + past_key_values (`Cache`): + The cache class that is being used currently to generate """ if attention_mask is not None and attention_mask.dim() == 4: # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. @@ -889,25 +934,128 @@ def _prepare_4d_causal_attention_mask_with_cache_position( causal_mask = torch.full( (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + if config.sliding_window is not None: + # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also + # the check is needed to verify is current checkpoint was trained with sliding window or not + if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: + sliding_attend_mask = torch.arange(target_length, device=device) <= ( + cache_position.reshape(-1, 1) - config.sliding_window + ) + diagonal_attend_mask.bitwise_or_(sliding_attend_mask) + causal_mask *= diagonal_attend_mask causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) if attention_mask is not None: causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.shape[-1] > target_length: + attention_mask = attention_mask[:, :target_length] mask_length = attention_mask.shape[-1] padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] padding_mask = padding_mask == 0 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( padding_mask, min_dtype ) - return causal_mask + def load_hook(self, state_dict, prefix, *args): + """ + Weights have to be permutted for correct rope formulation. We can't do this in the weights + as every other framework already uses the `Llama` orginal function (which is copyrighted btw). + And I am not even sure it's better.... anyways end of my rant + """ + for k in state_dict: + if "q_b_proj." in k: + weight = state_dict.pop(k[: self.qk_nope_head_dim]) + if "k_b_proj." in k: + weight = state_dict.pop(k[self.qk_nope_head_dim :]) + state_dict[k] = permute_for_rope(weight, weight.shape[0], weight.shape[1], weight.shape[2]) + class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... +def load_balancing_loss_func( + gate_logits: Union[torch.Tensor, Tuple[torch.Tensor], None], + num_experts: Optional[int] = None, + top_k=2, + attention_mask: Optional[torch.Tensor] = None, +) -> Union[torch.Tensor, int]: + r""" + Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. + + See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss + function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between + experts is too unbalanced. + + Args: + gate_logits: + Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of + shape [batch_size X sequence_length, num_experts]. + num_experts: + Number of experts + top_k: + The number of experts to route per-token, can be also interpreted as the `top-k` routing + parameter. + attention_mask (`torch.Tensor`, *optional*): + The attention_mask used in forward function + shape [batch_size X sequence_length] if not None. + + Returns: + The auxiliary loss. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + if isinstance(gate_logits, tuple): + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0 + ) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0 + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts + + class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] _tp_plan = {"lm_head": "colwise_rep"} @@ -917,6 +1065,9 @@ def __init__(self, config): self.model = DeepseekV3Model(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.router_aux_loss_coef = config.router_aux_loss_coef + self.num_experts = config.num_local_experts + self.num_experts_per_tok = config.num_experts_per_tok # Initialize weights and apply final processing self.post_init() @@ -947,12 +1098,13 @@ def forward( input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, @@ -979,8 +1131,8 @@ def forward( ```python >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM - >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") - >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") + >>> model = DeepseekV3ForCausalLM.from_pretrained("mistralai/DeepseekV3-8x7B-v0.1") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/DeepseekV3-8x7B-v0.1") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") @@ -990,7 +1142,12 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) @@ -1006,6 +1163,7 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, return_dict=return_dict, cache_position=cache_position, **kwargs, @@ -1018,18 +1176,33 @@ def forward( loss = None if labels is not None: - loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits if return_dict else outputs[-1], + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device if not return_dict: output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( + return MoeCausalLMOutputWithPast( loss=loss, + aux_loss=aux_loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, + router_logits=outputs.router_logits, ) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 8b20f068d489..9d76a96b6a87 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -13,15 +13,14 @@ from ...processing_utils import Unpack from ...utils import logging from ..llama.modeling_llama import ( - LlamaForCausalLM, LlamaForSequenceClassification, - LlamaModel, LlamaPreTrainedModel, LlamaRMSNorm, LlamaRotaryEmbedding, + apply_rotary_pos_emb, eager_attention_forward, - rotate_half, ) +from ..mixtral.modeling_mixtral import MixtralForCausalLM, MixtralModel from .configuration_deepseek_v3 import DeepseekV3Config @@ -42,39 +41,6 @@ def yarn_get_mscale(scale=1, mscale=1): return 0.1 * mscale * math.log(scale) + 1.0 -def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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`, *optional*): - Deprecated and unused. - 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_length, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_length, 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_length, 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.unsqueeze(unsqueeze_dim) - sin = sin.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): super().__init__() @@ -92,7 +58,7 @@ def forward(self, x): return down_proj -class MoEGate(nn.Module): +class DeepseekV3TopkRouter(nn.Module): def __init__(self, config): super().__init__() self.config = config @@ -101,7 +67,6 @@ def __init__(self, config): self.routed_scaling_factor = config.routed_scaling_factor self.n_group = config.n_group self.topk_group = config.topk_group - self.norm_topk_prob = config.norm_topk_prob self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) @@ -109,9 +74,9 @@ def __init__(self, config): def forward(self, hidden_states): batch_size, seq_length = hidden_states.shape[:-1] hidden_states = hidden_states.view(-1, self.config.hidden_size) - logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) - scores = logits.sigmoid() + scores = router_logits.sigmoid() scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) @@ -129,11 +94,10 @@ def forward(self, hidden_states): scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] _, topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) topk_weights = scores.gather(1, topk_indices) - if self.norm_topk_prob: - denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 - topk_weights /= denominator + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator topk_weights = topk_weights * self.routed_scaling_factor # must multiply the scaling factor - return topk_indices, topk_weights + return topk_indices, topk_weights, router_logits class DeepseekV3MoE(nn.Module): @@ -150,24 +114,21 @@ def __init__(self, config): for _ in range(config.n_routed_experts) ] ) - self.gate = MoEGate(config) - intermediate_size = config.moe_intermediate_size * config.n_shared_experts - self.shared_experts = DeepseekV3MLP(config=config, intermediate_size=intermediate_size) + self.gate = DeepseekV3TopkRouter(config) + self.shared_experts = DeepseekV3MLP(config=config, intermediate_size=config.moe_intermediate_size) def forward(self, hidden_states): - identity = hidden_states + residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) + topk_indices, topk_weights, router_logits = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - y = self.moe_infer(hidden_states, topk_indices, topk_weights).view(*orig_shape) - y = y + self.shared_experts(identity) - return y + hidden_states = self.moe_infer(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = hidden_states + self.shared_experts(residuals) + return hidden_states, router_logits - @torch.no_grad() def moe_infer(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): """ Perform inference using a Mixture of Experts (MoE) model. - Args: hidden_states (torch.Tensor): Input hidden states. topk_indices (torch.Tensor): Indices of the top-k experts for each token. @@ -178,20 +139,20 @@ def moe_infer(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, top """ num_experts = len(self.experts) batch_size, num_topk = topk_indices.shape + with torch.no_grad(): + # Count the number of tokens assigned to each expert + expert_counts = topk_indices.new_zeros((batch_size, num_experts)) + expert_counts.scatter_(1, topk_indices, 1) + tokens_per_expert = expert_counts.sum(dim=0) - # Count the number of tokens assigned to each expert - expert_counts = topk_indices.new_zeros((batch_size, num_experts)) - expert_counts.scatter_(1, topk_indices, 1) - tokens_per_expert = expert_counts.sum(dim=0) - - # Sort tokens by their assigned expert - sorted_indices = topk_indices.view(-1).argsort() - sorted_tokens = hidden_states[sorted_indices // num_topk] - tokens_per_expert = tokens_per_expert.cpu().numpy() + # Sort tokens by their assigned expert + sorted_indices = topk_indices.view(-1).argsort() + sorted_tokens = hidden_states[sorted_indices // num_topk] + tokens_per_expert = tokens_per_expert.cpu().numpy() - # Process tokens through their assigned experts - expert_outputs = [] - current_pos = 0 + # Process tokens through their assigned experts + expert_outputs = [] + current_pos = 0 for expert_idx, num_tokens in enumerate(tokens_per_expert): if num_tokens == 0: @@ -279,26 +240,22 @@ def forward( hidden_shape = (*input_shape, self.num_heads, -1) batch_size, seq_length = input_shape - q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) - q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) + q_rot, q_pass = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_pe = k_pe.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2) - kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(hidden_shape).transpose(1, 2) + k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(hidden_shape).transpose(1, 2) + k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - cos, sin = position_embeddings - q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin) + k_rot = k_rot.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2) - query_states = k_pe.new_empty(batch_size, self.num_heads, seq_length, self.q_head_dim) - query_states[:, :, :, : self.qk_nope_head_dim] = q_nope - query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + cos, sin = position_embeddings + q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - key_states = k_pe.new_empty(batch_size, self.num_heads, seq_length, self.q_head_dim) - key_states[:, :, :, : self.qk_nope_head_dim] = k_nope - key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + query_states = torch.cat(q_rot, q_pass, dim=-1) + key_states = torch.cat(k_rot, k_pass, dim=-1) if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) @@ -359,6 +316,7 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC @@ -386,11 +344,19 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) + + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = (torch.zeros((1,), device=hidden_states.device, dtype=torch.int64),) + hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) + if output_router_logits: + outputs += (router_logits,) return outputs @@ -399,11 +365,38 @@ class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): pass -class DeepseekV3Model(LlamaModel): - pass +def permute_for_rope(input_tensor, n_heads, dim1, dim2): + """ + When you go from the complex ROPE formulation to sin and cos one, you need + to permute the query and key weights (to avoid doing it on the fly) + """ + input_tensor = input_tensor.reshape(dim1, dim2) + input_tensor = input_tensor.view(n_heads, dim1 // n_heads // 2, 2, dim2) + input_tensor = input_tensor.transpose(1, 2).reshape(dim1, dim2) + return input_tensor + + +class DeepseekV3Model(MixtralModel): + def __init__(self, config): + super().__init__(config) + self._register_load_state_dict_pre_hook(self.load_hook) + self.post_init() + + def load_hook(self, state_dict, prefix, *args): + """ + Weights have to be permutted for correct rope formulation. We can't do this in the weights + as every other framework already uses the `Llama` orginal function (which is copyrighted btw). + And I am not even sure it's better.... anyways end of my rant + """ + for k in state_dict: + if "q_b_proj." in k: + weight = state_dict.pop(k[: self.qk_nope_head_dim]) + if "k_b_proj." in k: + weight = state_dict.pop(k[self.qk_nope_head_dim :]) + state_dict[k] = permute_for_rope(weight, weight.shape[0], weight.shape[1], weight.shape[2]) -class DeepseekV3ForCausalLM(LlamaForCausalLM): +class DeepseekV3ForCausalLM(MixtralForCausalLM): pass From 24bc8b2c279237c6faca54b421e8aade6fb65777 Mon Sep 17 00:00:00 2001 From: ryan u Date: Fri, 31 Jan 2025 22:44:28 +0900 Subject: [PATCH 23/86] revise load_hook to work properly; make moe func trainable; use llama instead of mixtral --- .../deepseek_v3/modeling_deepseek_v3.py | 336 +++++------------- .../models/deepseek_v3/modular_deepseek_v3.py | 125 +++---- 2 files changed, 140 insertions(+), 321 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index f0ee6bf50f83..6c59e03b63f5 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -12,17 +12,11 @@ from torch import nn from ...activations import ACT2FN -from ...cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache +from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationMixin from ...modeling_attn_mask_utils import AttentionMaskConverter from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import ( - BaseModelOutputWithPast, - CausalLMOutputWithPast, - MoeCausalLMOutputWithPast, - MoeModelOutputWithPast, - SequenceClassifierOutputWithPast, -) +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -203,60 +197,27 @@ def forward(self, hidden_states): orig_shape = hidden_states.shape topk_indices, topk_weights, router_logits = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe_infer(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states, router_logits - def moe_infer(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): - """ - Perform inference using a Mixture of Experts (MoE) model. - Args: - hidden_states (torch.Tensor): Input hidden states. - topk_indices (torch.Tensor): Indices of the top-k experts for each token. - topk_weights (torch.Tensor): Weights associated with the top-k experts. + def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) + expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) + expert_mask = expert_mask.permute(2, 0, 1) - Returns: - torch.Tensor: Output of the MoE model. - """ - num_experts = len(self.experts) - batch_size, num_topk = topk_indices.shape - with torch.no_grad(): - # Count the number of tokens assigned to each expert - expert_counts = topk_indices.new_zeros((batch_size, num_experts)) - expert_counts.scatter_(1, topk_indices, 1) - tokens_per_expert = expert_counts.sum(dim=0) - - # Sort tokens by their assigned expert - sorted_indices = topk_indices.view(-1).argsort() - sorted_tokens = hidden_states[sorted_indices // num_topk] - tokens_per_expert = tokens_per_expert.cpu().numpy() - - # Process tokens through their assigned experts - expert_outputs = [] - current_pos = 0 - - for expert_idx, num_tokens in enumerate(tokens_per_expert): - if num_tokens == 0: - continue - - next_pos = current_pos + num_tokens + for expert_idx in range(len(self.experts)): expert = self.experts[expert_idx] - expert_tokens = sorted_tokens[current_pos:next_pos] - expert_outputs.append(expert(expert_tokens)) - current_pos = next_pos + mask = expert_mask[expert_idx] + token_indices, weight_indices = torch.where(mask) - # Combine the outputs from all experts - expert_outputs = torch.cat(expert_outputs, dim=0) if expert_outputs else sorted_tokens.new_empty(0) - - # Reorder the outputs to match the original token sequence - reordered_outputs = torch.empty_like(expert_outputs) - reordered_outputs[sorted_indices] = expert_outputs - - # Reshape and apply the expert weights - reordered_outputs = reordered_outputs.view(batch_size, num_topk, -1).type(topk_weights.dtype) - moe_output = torch.matmul(topk_weights.unsqueeze(1), reordered_outputs) - moe_output = moe_output.sum(dim=1).type(hidden_states.dtype) - return moe_output + if token_indices.numel() > 0: + expert_weights = topk_weights[token_indices, weight_indices] + expert_input = hidden_states[token_indices] + expert_output = expert(expert_input) + weighted_output = expert_output * expert_weights.unsqueeze(-1) + final_hidden_states.index_add_(0, token_indices, weighted_output) + return final_hidden_states.type(hidden_states.dtype) def rotate_half(x): @@ -642,17 +603,6 @@ def _init_weights(self, module): """ -def permute_for_rope(input_tensor, n_heads, dim1, dim2): - """ - When you go from the complex ROPE formulation to sin and cos one, you need - to permute the query and key weights (to avoid doing it on the fly) - """ - input_tensor = input_tensor.reshape(dim1, dim2) - input_tensor = input_tensor.view(n_heads, dim1 // n_heads // 2, 2, dim2) - input_tensor = input_tensor.transpose(1, 2).reshape(dim1, dim2) - return input_tensor - - @add_start_docstrings( "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", DEEPSEEK_V3_START_DOCSTRING, @@ -694,48 +644,43 @@ def forward( input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, **flash_attn_kwargs: Unpack[FlashAttentionKwargs], ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - if use_cache and past_key_values is None: - past_key_values = DynamicCache() + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) + if position_ids is None: position_ids = cache_position.unsqueeze(0) @@ -751,9 +696,8 @@ def forward( # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None - all_router_logits = () if output_router_logits else None - for decoder_layer in self.layers: + for decoder_layer in self.layers[: self.config.num_hidden_layers]: if output_hidden_states: all_hidden_states += (hidden_states,) @@ -765,7 +709,6 @@ def forward( position_ids, past_key_values, output_attentions, - output_router_logits, use_cache, cache_position, position_embeddings, @@ -777,7 +720,6 @@ def forward( position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, - output_router_logits=output_router_logits, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -789,21 +731,17 @@ def forward( if output_attentions: all_self_attns += (layer_outputs[1],) - if output_router_logits: - all_router_logits += (layer_outputs[-1],) - hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) - output = MoeModelOutputWithPast( + output = BaseModelOutputWithPast( last_hidden_state=hidden_states, - past_key_values=past_key_values, + past_key_values=past_key_values if use_cache else None, hidden_states=all_hidden_states, attentions=all_self_attns, - router_logits=all_router_logits, ) return output if return_dict else output.to_tuple() @@ -816,15 +754,7 @@ def _update_causal_mask( output_attentions: bool, ): if self.config._attn_implementation == "flash_attention_2": - if attention_mask is not None and past_key_values is not None: - is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] - if is_padding_right: - raise ValueError( - "You are attempting to perform batched generation with padding_side='right'" - " this may lead to unexpected behaviour for Flash Attention version of DeepseekV3. Make sure to " - " call `tokenizer.padding_side = 'left'` before tokenizing the input. " - ) - if attention_mask is not None and 0.0 in attention_mask: + if attention_mask is not None and (attention_mask == 0.0).any(): return attention_mask return None @@ -833,30 +763,21 @@ def _update_causal_mask( # to infer the attention mask. past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 using_static_cache = isinstance(past_key_values, StaticCache) - using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if ( - self.config._attn_implementation == "sdpa" - and not (using_static_cache or using_sliding_window_cache) - and not output_attentions - ): + if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: if AttentionMaskConverter._ignore_causal_mask_sdpa( attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens, - sliding_window=self.config.sliding_window, is_training=self.training, ): return None dtype, device = input_tensor.dtype, input_tensor.device - min_dtype = torch.finfo(dtype).min sequence_length = input_tensor.shape[1] - # SlidingWindowCache or StaticCache - if using_sliding_window_cache or using_static_cache: + if using_static_cache: target_length = past_key_values.get_max_cache_shape() - # DynamicCache or no cache else: target_length = ( attention_mask.shape[-1] @@ -873,8 +794,6 @@ def _update_causal_mask( device=device, cache_position=cache_position, batch_size=input_tensor.shape[0], - config=self.config, - past_key_values=past_key_values, ) if ( @@ -886,6 +805,7 @@ def _update_causal_mask( # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. # Details: https://github.com/pytorch/pytorch/issues/110213 + min_dtype = torch.finfo(dtype).min causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) return causal_mask @@ -899,8 +819,7 @@ def _prepare_4d_causal_attention_mask_with_cache_position( device: torch.device, cache_position: torch.Tensor, batch_size: int, - config: DeepseekV3Config, - past_key_values: Cache, + **kwargs, ): """ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape @@ -908,11 +827,13 @@ def _prepare_4d_causal_attention_mask_with_cache_position( Args: attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape + `(batch_size, 1, query_length, key_value_length)`. sequence_length (`int`): The sequence length being processed. target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. + The target length: when generating with static cache, the mask should be as long as the static cache, + to account for the 0 padding, the part of the cache that is not filled yet. dtype (`torch.dtype`): The dtype to use for the 4D attention mask. device (`torch.device`): @@ -921,10 +842,6 @@ def _prepare_4d_causal_attention_mask_with_cache_position( Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): Batch size. - config (`DeepseekV3Config`): - The model's configuration class - past_key_values (`Cache`): - The cache class that is being used currently to generate """ if attention_mask is not None and attention_mask.dim() == 4: # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. @@ -934,128 +851,68 @@ def _prepare_4d_causal_attention_mask_with_cache_position( causal_mask = torch.full( (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device ) - diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) - if config.sliding_window is not None: - # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also - # the check is needed to verify is current checkpoint was trained with sliding window or not - if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: - sliding_attend_mask = torch.arange(target_length, device=device) <= ( - cache_position.reshape(-1, 1) - config.sliding_window - ) - diagonal_attend_mask.bitwise_or_(sliding_attend_mask) - causal_mask *= diagonal_attend_mask + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) if attention_mask is not None: causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - if attention_mask.shape[-1] > target_length: - attention_mask = attention_mask[:, :target_length] mask_length = attention_mask.shape[-1] padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] padding_mask = padding_mask == 0 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( padding_mask, min_dtype ) + return causal_mask def load_hook(self, state_dict, prefix, *args): """ - Weights have to be permutted for correct rope formulation. We can't do this in the weights - as every other framework already uses the `Llama` orginal function (which is copyrighted btw). + Weights have to be permuted for correct rope formulation. We can't do this in the weights + as every other framework already uses the `Llama` original function (which is copyrighted btw). And I am not even sure it's better.... anyways end of my rant """ + + def permute_for_rope(input_tensor): + """ + When you go from the complex ROPE formulation to sin and cos one, you need + to permute the query and key weights (to avoid doing it on the fly) + """ + n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] + input_tensor = input_tensor.reshape(n_heads * dim1, dim2) + input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) + input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) + return input_tensor + + def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): + weight = state_dict[key] + weight = weight.view(num_heads, head_dim, -1) + weight_rot = weight[:, -rope_dim:] + weight_rot = permute_for_rope(weight_rot) + weight[:, -rope_dim:] = weight_rot + weight = weight.view(-1, weight.shape[-1]) + state_dict[key] = weight + for k in state_dict: if "q_b_proj." in k: - weight = state_dict.pop(k[: self.qk_nope_head_dim]) - if "k_b_proj." in k: - weight = state_dict.pop(k[self.qk_nope_head_dim :]) - state_dict[k] = permute_for_rope(weight, weight.shape[0], weight.shape[1], weight.shape[2]) + permute_layer_for_rope( + k, + num_heads=self.config.num_attention_heads, + head_dim=self.config.q_head_dim, + rope_dim=self.config.qk_rope_head_dim, + ) + if "kv_a_proj_with_mqa." in k: + permute_layer_for_rope( + k, + num_heads=1, + head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, + rope_dim=self.config.qk_rope_head_dim, + ) class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... -def load_balancing_loss_func( - gate_logits: Union[torch.Tensor, Tuple[torch.Tensor], None], - num_experts: Optional[int] = None, - top_k=2, - attention_mask: Optional[torch.Tensor] = None, -) -> Union[torch.Tensor, int]: - r""" - Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. - - See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss - function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between - experts is too unbalanced. - - Args: - gate_logits: - Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of - shape [batch_size X sequence_length, num_experts]. - num_experts: - Number of experts - top_k: - The number of experts to route per-token, can be also interpreted as the `top-k` routing - parameter. - attention_mask (`torch.Tensor`, *optional*): - The attention_mask used in forward function - shape [batch_size X sequence_length] if not None. - - Returns: - The auxiliary loss. - """ - if gate_logits is None or not isinstance(gate_logits, tuple): - return 0 - - if isinstance(gate_logits, tuple): - compute_device = gate_logits[0].device - concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) - - routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) - - _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) - - expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) - - if attention_mask is None: - # Compute the percentage of tokens routed to each experts - tokens_per_expert = torch.mean(expert_mask.float(), dim=0) - - # Compute the average probability of routing to these experts - router_prob_per_expert = torch.mean(routing_weights, dim=0) - else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) - - # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask - expert_attention_mask = ( - attention_mask[None, :, :, None, None] - .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) - .reshape(-1, top_k, num_experts) - .to(compute_device) - ) - - # Compute the percentage of tokens routed to each experts - tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( - expert_attention_mask, dim=0 - ) - - # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert - router_per_expert_attention_mask = ( - attention_mask[None, :, :, None] - .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) - .reshape(-1, num_experts) - .to(compute_device) - ) - - # Compute the average probability of routing to these experts - router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( - router_per_expert_attention_mask, dim=0 - ) - - overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) - return overall_loss * num_experts - - class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] _tp_plan = {"lm_head": "colwise_rep"} @@ -1065,9 +922,6 @@ def __init__(self, config): self.model = DeepseekV3Model(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - self.router_aux_loss_coef = config.router_aux_loss_coef - self.num_experts = config.num_local_experts - self.num_experts_per_tok = config.num_experts_per_tok # Initialize weights and apply final processing self.post_init() @@ -1098,13 +952,12 @@ def forward( input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, @@ -1131,8 +984,8 @@ def forward( ```python >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM - >>> model = DeepseekV3ForCausalLM.from_pretrained("mistralai/DeepseekV3-8x7B-v0.1") - >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/DeepseekV3-8x7B-v0.1") + >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") @@ -1142,12 +995,7 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) @@ -1163,7 +1011,6 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - output_router_logits=output_router_logits, return_dict=return_dict, cache_position=cache_position, **kwargs, @@ -1176,33 +1023,18 @@ def forward( loss = None if labels is not None: - loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) - - aux_loss = None - if output_router_logits: - aux_loss = load_balancing_loss_func( - outputs.router_logits if return_dict else outputs[-1], - self.num_experts, - self.num_experts_per_tok, - attention_mask, - ) - if labels is not None: - loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: output = (logits,) + outputs[1:] - if output_router_logits: - output = (aux_loss,) + output return (loss,) + output if loss is not None else output - return MoeCausalLMOutputWithPast( + return CausalLMOutputWithPast( loss=loss, - aux_loss=aux_loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, - router_logits=outputs.router_logits, ) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 9d76a96b6a87..db3bb276742a 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -13,14 +13,15 @@ from ...processing_utils import Unpack from ...utils import logging from ..llama.modeling_llama import ( + LlamaForCausalLM, LlamaForSequenceClassification, + LlamaModel, LlamaPreTrainedModel, LlamaRMSNorm, LlamaRotaryEmbedding, apply_rotary_pos_emb, eager_attention_forward, ) -from ..mixtral.modeling_mixtral import MixtralForCausalLM, MixtralModel from .configuration_deepseek_v3 import DeepseekV3Config @@ -122,60 +123,27 @@ def forward(self, hidden_states): orig_shape = hidden_states.shape topk_indices, topk_weights, router_logits = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe_infer(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states, router_logits - def moe_infer(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): - """ - Perform inference using a Mixture of Experts (MoE) model. - Args: - hidden_states (torch.Tensor): Input hidden states. - topk_indices (torch.Tensor): Indices of the top-k experts for each token. - topk_weights (torch.Tensor): Weights associated with the top-k experts. - - Returns: - torch.Tensor: Output of the MoE model. - """ - num_experts = len(self.experts) - batch_size, num_topk = topk_indices.shape - with torch.no_grad(): - # Count the number of tokens assigned to each expert - expert_counts = topk_indices.new_zeros((batch_size, num_experts)) - expert_counts.scatter_(1, topk_indices, 1) - tokens_per_expert = expert_counts.sum(dim=0) - - # Sort tokens by their assigned expert - sorted_indices = topk_indices.view(-1).argsort() - sorted_tokens = hidden_states[sorted_indices // num_topk] - tokens_per_expert = tokens_per_expert.cpu().numpy() - - # Process tokens through their assigned experts - expert_outputs = [] - current_pos = 0 - - for expert_idx, num_tokens in enumerate(tokens_per_expert): - if num_tokens == 0: - continue - - next_pos = current_pos + num_tokens - expert = self.experts[expert_idx] - expert_tokens = sorted_tokens[current_pos:next_pos] - expert_outputs.append(expert(expert_tokens)) - current_pos = next_pos - - # Combine the outputs from all experts - expert_outputs = torch.cat(expert_outputs, dim=0) if expert_outputs else sorted_tokens.new_empty(0) + def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) + expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) + expert_mask = expert_mask.permute(2, 0, 1) - # Reorder the outputs to match the original token sequence - reordered_outputs = torch.empty_like(expert_outputs) - reordered_outputs[sorted_indices] = expert_outputs + for expert_idx in range(len(self.experts)): + expert = self.experts[expert_idx] + mask = expert_mask[expert_idx] + token_indices, weight_indices = torch.where(mask) - # Reshape and apply the expert weights - reordered_outputs = reordered_outputs.view(batch_size, num_topk, -1).type(topk_weights.dtype) - moe_output = torch.matmul(topk_weights.unsqueeze(1), reordered_outputs) - moe_output = moe_output.sum(dim=1).type(hidden_states.dtype) - return moe_output + if token_indices.numel() > 0: + expert_weights = topk_weights[token_indices, weight_indices] + expert_input = hidden_states[token_indices] + expert_output = expert(expert_input) + weighted_output = expert_output * expert_weights.unsqueeze(-1) + final_hidden_states.index_add_(0, token_indices, weighted_output) + return final_hidden_states.type(hidden_states.dtype) class DeepseekV3Attention(nn.Module): @@ -365,18 +333,7 @@ class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): pass -def permute_for_rope(input_tensor, n_heads, dim1, dim2): - """ - When you go from the complex ROPE formulation to sin and cos one, you need - to permute the query and key weights (to avoid doing it on the fly) - """ - input_tensor = input_tensor.reshape(dim1, dim2) - input_tensor = input_tensor.view(n_heads, dim1 // n_heads // 2, 2, dim2) - input_tensor = input_tensor.transpose(1, 2).reshape(dim1, dim2) - return input_tensor - - -class DeepseekV3Model(MixtralModel): +class DeepseekV3Model(LlamaModel): def __init__(self, config): super().__init__(config) self._register_load_state_dict_pre_hook(self.load_hook) @@ -384,19 +341,49 @@ def __init__(self, config): def load_hook(self, state_dict, prefix, *args): """ - Weights have to be permutted for correct rope formulation. We can't do this in the weights - as every other framework already uses the `Llama` orginal function (which is copyrighted btw). + Weights have to be permuted for correct rope formulation. We can't do this in the weights + as every other framework already uses the `Llama` original function (which is copyrighted btw). And I am not even sure it's better.... anyways end of my rant """ + + def permute_for_rope(input_tensor): + """ + When you go from the complex ROPE formulation to sin and cos one, you need + to permute the query and key weights (to avoid doing it on the fly) + """ + n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] + input_tensor = input_tensor.reshape(n_heads * dim1, dim2) + input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) + input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) + return input_tensor + + def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): + weight = state_dict[key] + weight = weight.view(num_heads, head_dim, -1) + weight_rot = weight[:, -rope_dim:] + weight_rot = permute_for_rope(weight_rot) + weight[:, -rope_dim:] = weight_rot + weight = weight.view(-1, weight.shape[-1]) + state_dict[key] = weight + for k in state_dict: if "q_b_proj." in k: - weight = state_dict.pop(k[: self.qk_nope_head_dim]) - if "k_b_proj." in k: - weight = state_dict.pop(k[self.qk_nope_head_dim :]) - state_dict[k] = permute_for_rope(weight, weight.shape[0], weight.shape[1], weight.shape[2]) + permute_layer_for_rope( + k, + num_heads=self.config.num_attention_heads, + head_dim=self.config.q_head_dim, + rope_dim=self.config.qk_rope_head_dim, + ) + if "kv_a_proj_with_mqa." in k: + permute_layer_for_rope( + k, + num_heads=1, + head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, + rope_dim=self.config.qk_rope_head_dim, + ) -class DeepseekV3ForCausalLM(MixtralForCausalLM): +class DeepseekV3ForCausalLM(LlamaForCausalLM): pass From 5c0cd917b02d5ba9d4fb373e01f3fe24b36c40ff Mon Sep 17 00:00:00 2001 From: ryan u Date: Fri, 31 Jan 2025 23:37:18 +0900 Subject: [PATCH 24/86] fix attention forward --- .../models/deepseek_v3/modeling_deepseek_v3.py | 7 ++++--- src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 6c59e03b63f5..0c5c510d2277 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -361,7 +361,7 @@ def forward( batch_size, seq_length = input_shape q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) - q_rot, q_pass = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) @@ -373,9 +373,10 @@ def forward( cos, sin = position_embeddings q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) + k_rot = k_rot.expand(batch_size, self.num_heads, seq_length, self.qk_rope_head_dim) - query_states = torch.cat(q_rot, q_pass, dim=-1) - key_states = torch.cat(k_rot, k_pass, dim=-1) + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index db3bb276742a..5ac409e3b7f0 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -209,7 +209,7 @@ def forward( batch_size, seq_length = input_shape q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) - q_rot, q_pass = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) @@ -221,9 +221,10 @@ def forward( cos, sin = position_embeddings q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) + k_rot = k_rot.expand(batch_size, self.num_heads, seq_length, self.qk_rope_head_dim) - query_states = torch.cat(q_rot, q_pass, dim=-1) - key_states = torch.cat(k_rot, k_pass, dim=-1) + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) From 8e994dd83a68f20081b60221d35a7e3b2c43ab61 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 1 Feb 2025 14:01:23 +0900 Subject: [PATCH 25/86] use -1 for not-changing dim when to use exapnd --- src/transformers/models/deepseek_v3/modeling_deepseek_v3.py | 3 +-- src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 0c5c510d2277..af55c5f668a4 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -358,7 +358,6 @@ def forward( ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, self.num_heads, -1) - batch_size, seq_length = input_shape q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) @@ -373,7 +372,7 @@ def forward( cos, sin = position_embeddings q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - k_rot = k_rot.expand(batch_size, self.num_heads, seq_length, self.qk_rope_head_dim) + k_rot = k_rot.expand(-1, self.num_heads, -1, -1) query_states = torch.cat((q_pass, q_rot), dim=-1) key_states = torch.cat((k_pass, k_rot), dim=-1) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 5ac409e3b7f0..6966a0de2132 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -206,7 +206,6 @@ def forward( ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, self.num_heads, -1) - batch_size, seq_length = input_shape q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) @@ -221,7 +220,7 @@ def forward( cos, sin = position_embeddings q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - k_rot = k_rot.expand(batch_size, self.num_heads, seq_length, self.qk_rope_head_dim) + k_rot = k_rot.expand(-1, self.num_heads, -1, -1) query_states = torch.cat((q_pass, q_rot), dim=-1) key_states = torch.cat((k_pass, k_rot), dim=-1) From 7405a95f71f26ad68fe7000367b77507f621ee66 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 1 Feb 2025 17:08:00 +0900 Subject: [PATCH 26/86] refactor DeepseekV3TopkRouter --- .../deepseek_v3/modeling_deepseek_v3.py | 35 +++++++++++-------- .../models/deepseek_v3/modular_deepseek_v3.py | 35 +++++++++++-------- 2 files changed, 40 insertions(+), 30 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index af55c5f668a4..781d7756de19 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -142,37 +142,42 @@ def __init__(self, config): self.routed_scaling_factor = config.routed_scaling_factor self.n_group = config.n_group self.topk_group = config.topk_group + self.norm_topk_prob = config.norm_topk_prob self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) def forward(self, hidden_states): - batch_size, seq_length = hidden_states.shape[:-1] hidden_states = hidden_states.view(-1, self.config.hidden_size) router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) - scores = router_logits.sigmoid() + topk_indices = self.get_topk_indices(scores) + topk_weights = scores.gather(1, topk_indices) + if self.norm_topk_prob: + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator + topk_weights = topk_weights * self.routed_scaling_factor + return topk_indices, topk_weights, router_logits + + @torch.no_grad() + def get_topk_indices(self, scores): scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) .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] + ) + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) score_mask = ( group_mask.unsqueeze(-1) - .expand(batch_size * seq_length, self.n_group, self.n_routed_experts // self.n_group) + .expand(-1, self.n_group, self.n_routed_experts // self.n_group) .reshape(-1, self.n_routed_experts) - ) # [n, e] - scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] - _, topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) - topk_weights = scores.gather(1, topk_indices) - denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 - topk_weights /= denominator - topk_weights = topk_weights * self.routed_scaling_factor # must multiply the scaling factor - return topk_indices, topk_weights, router_logits + ) + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) + topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + return topk_indices class DeepseekV3MoE(nn.Module): diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 6966a0de2132..83ad25684814 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -68,37 +68,42 @@ def __init__(self, config): self.routed_scaling_factor = config.routed_scaling_factor self.n_group = config.n_group self.topk_group = config.topk_group + self.norm_topk_prob = config.norm_topk_prob self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) def forward(self, hidden_states): - batch_size, seq_length = hidden_states.shape[:-1] hidden_states = hidden_states.view(-1, self.config.hidden_size) router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) - scores = router_logits.sigmoid() + topk_indices = self.get_topk_indices(scores) + topk_weights = scores.gather(1, topk_indices) + if self.norm_topk_prob: + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator + topk_weights = topk_weights * self.routed_scaling_factor + return topk_indices, topk_weights, router_logits + + @torch.no_grad() + def get_topk_indices(self, scores): scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) .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] + ) + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) score_mask = ( group_mask.unsqueeze(-1) - .expand(batch_size * seq_length, self.n_group, self.n_routed_experts // self.n_group) + .expand(-1, self.n_group, self.n_routed_experts // self.n_group) .reshape(-1, self.n_routed_experts) - ) # [n, e] - scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] - _, topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) - topk_weights = scores.gather(1, topk_indices) - denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 - topk_weights /= denominator - topk_weights = topk_weights * self.routed_scaling_factor # must multiply the scaling factor - return topk_indices, topk_weights, router_logits + ) + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) + topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + return topk_indices class DeepseekV3MoE(nn.Module): From ea3c9225546c2ac08622487e6b397cceacc93ff0 Mon Sep 17 00:00:00 2001 From: ryan u Date: Mon, 3 Feb 2025 16:59:37 +0900 Subject: [PATCH 27/86] use reshape_for_rope instead of load_hook; revise attention forward for TP; rename q_head_dim with qk_head_dim --- .../deepseek_v3/configuration_deepseek_v3.py | 5 +- .../deepseek_v3/modeling_deepseek_v3.py | 83 ++++++------------ .../models/deepseek_v3/modular_deepseek_v3.py | 85 ++++++------------- 3 files changed, 52 insertions(+), 121 deletions(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 8d1c56343638..5943aa12648e 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -138,6 +138,9 @@ class DeepseekV3Config(PretrainedConfig): "layers.*.gate_proj": "colwise", "layers.*.up_proj": "colwise", "layers.*.down_proj": "rowwise", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", } def __init__( @@ -194,7 +197,7 @@ def __init__( self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim - self.q_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim self.head_dim = qk_rope_head_dim self.n_group = n_group self.topk_group = topk_group diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 781d7756de19..ef6be2f6add5 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -319,12 +319,12 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): 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.q_head_dim + self.qk_head_dim = config.qk_head_dim self.is_causal = True self.q_a_proj = nn.Linear(config.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.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) self.kv_a_proj_with_mqa = nn.Linear( config.hidden_size, @@ -334,7 +334,7 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): self.kv_a_layernorm = DeepseekV3RMSNorm(self.kv_lora_rank) self.kv_b_proj = nn.Linear( self.kv_lora_rank, - self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, ) @@ -344,7 +344,7 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): bias=config.attention_bias, ) - self.scaling = self.q_head_dim ** (-0.5) + self.scaling = self.qk_head_dim ** (-0.5) if self.config.rope_scaling is not None: mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) scaling_factor = self.config.rope_scaling["factor"] @@ -361,29 +361,31 @@ def forward( cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, self.num_heads, -1) + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.qk_head_dim) + key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim) - q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(query_shape).transpose(1, 2) q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(hidden_shape).transpose(1, 2) + k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - k_rot = k_rot.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2) + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings + q_rot, k_rot = self.reshape_for_rope(q_rot, k_rot) q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - k_rot = k_rot.expand(-1, self.num_heads, -1, -1) + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) query_states = torch.cat((q_pass, q_rot), dim=-1) key_states = torch.cat((k_pass, k_rot), dim=-1) - if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache @@ -411,13 +413,20 @@ def forward( **kwargs, ) - if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: + if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: attn_output = attn_output[:, :, :, : self.v_head_dim] - attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights + def reshape_for_rope(self, q, k): + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(-1, -2).reshape(b, h, s, d) + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(-1, -2).reshape(b, h, s, d) + return q, k + class DeepseekV3DecoderLayer(nn.Module): def __init__(self, config: DeepseekV3Config, layer_idx: int): @@ -620,7 +629,7 @@ class DeepseekV3Model(DeepseekV3PreTrainedModel): config: DeepseekV3Config """ - def __init__(self, config): + def __init__(self, config: DeepseekV3Config): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size @@ -632,7 +641,6 @@ def __init__(self, config): self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = DeepseekV3RotaryEmbedding(config=config) self.gradient_checkpointing = False - self._register_load_state_dict_pre_hook(self.load_hook) # Initialize weights and apply final processing self.post_init() @@ -871,49 +879,6 @@ def _prepare_4d_causal_attention_mask_with_cache_position( return causal_mask - def load_hook(self, state_dict, prefix, *args): - """ - Weights have to be permuted for correct rope formulation. We can't do this in the weights - as every other framework already uses the `Llama` original function (which is copyrighted btw). - And I am not even sure it's better.... anyways end of my rant - """ - - def permute_for_rope(input_tensor): - """ - When you go from the complex ROPE formulation to sin and cos one, you need - to permute the query and key weights (to avoid doing it on the fly) - """ - n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] - input_tensor = input_tensor.reshape(n_heads * dim1, dim2) - input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) - input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) - return input_tensor - - def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): - weight = state_dict[key] - weight = weight.view(num_heads, head_dim, -1) - weight_rot = weight[:, -rope_dim:] - weight_rot = permute_for_rope(weight_rot) - weight[:, -rope_dim:] = weight_rot - weight = weight.view(-1, weight.shape[-1]) - state_dict[key] = weight - - for k in state_dict: - if "q_b_proj." in k: - permute_layer_for_rope( - k, - num_heads=self.config.num_attention_heads, - head_dim=self.config.q_head_dim, - rope_dim=self.config.qk_rope_head_dim, - ) - if "kv_a_proj_with_mqa." in k: - permute_layer_for_rope( - k, - num_heads=1, - head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, - rope_dim=self.config.qk_rope_head_dim, - ) - class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 83ad25684814..55d1ac5c691f 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -167,12 +167,12 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): 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.q_head_dim + self.qk_head_dim = config.qk_head_dim self.is_causal = True self.q_a_proj = nn.Linear(config.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.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) self.kv_a_proj_with_mqa = nn.Linear( config.hidden_size, @@ -182,7 +182,7 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): self.kv_a_layernorm = DeepseekV3RMSNorm(self.kv_lora_rank) self.kv_b_proj = nn.Linear( self.kv_lora_rank, - self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, ) @@ -192,7 +192,7 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): bias=config.attention_bias, ) - self.scaling = self.q_head_dim ** (-0.5) + self.scaling = self.qk_head_dim ** (-0.5) if self.config.rope_scaling is not None: mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) scaling_factor = self.config.rope_scaling["factor"] @@ -209,29 +209,31 @@ def forward( cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, self.num_heads, -1) + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.qk_head_dim) + key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim) - q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2) + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(query_shape).transpose(1, 2) q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(hidden_shape).transpose(1, 2) + k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - k_rot = k_rot.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2) + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings + q_rot, k_rot = self.reshape_for_rope(q_rot, k_rot) q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - k_rot = k_rot.expand(-1, self.num_heads, -1, -1) + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) query_states = torch.cat((q_pass, q_rot), dim=-1) key_states = torch.cat((k_pass, k_rot), dim=-1) - if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache @@ -259,13 +261,20 @@ def forward( **kwargs, ) - if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: + if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: attn_output = attn_output[:, :, :, : self.v_head_dim] - attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights + def reshape_for_rope(self, q, k): + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(-1, -2).reshape(b, h, s, d) + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(-1, -2).reshape(b, h, s, d) + return q, k + class DeepseekV3DecoderLayer(nn.Module): def __init__(self, config: DeepseekV3Config, layer_idx: int): @@ -339,53 +348,7 @@ class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): class DeepseekV3Model(LlamaModel): - def __init__(self, config): - super().__init__(config) - self._register_load_state_dict_pre_hook(self.load_hook) - self.post_init() - - def load_hook(self, state_dict, prefix, *args): - """ - Weights have to be permuted for correct rope formulation. We can't do this in the weights - as every other framework already uses the `Llama` original function (which is copyrighted btw). - And I am not even sure it's better.... anyways end of my rant - """ - - def permute_for_rope(input_tensor): - """ - When you go from the complex ROPE formulation to sin and cos one, you need - to permute the query and key weights (to avoid doing it on the fly) - """ - n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] - input_tensor = input_tensor.reshape(n_heads * dim1, dim2) - input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) - input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) - return input_tensor - - def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): - weight = state_dict[key] - weight = weight.view(num_heads, head_dim, -1) - weight_rot = weight[:, -rope_dim:] - weight_rot = permute_for_rope(weight_rot) - weight[:, -rope_dim:] = weight_rot - weight = weight.view(-1, weight.shape[-1]) - state_dict[key] = weight - - for k in state_dict: - if "q_b_proj." in k: - permute_layer_for_rope( - k, - num_heads=self.config.num_attention_heads, - head_dim=self.config.q_head_dim, - rope_dim=self.config.qk_rope_head_dim, - ) - if "kv_a_proj_with_mqa." in k: - permute_layer_for_rope( - k, - num_heads=1, - head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, - rope_dim=self.config.qk_rope_head_dim, - ) + pass class DeepseekV3ForCausalLM(LlamaForCausalLM): From c8132687dfa51416e4f84f017034948bee79260c Mon Sep 17 00:00:00 2001 From: ryan u Date: Mon, 3 Feb 2025 19:05:15 +0900 Subject: [PATCH 28/86] register pre_hook and hook both --- .../deepseek_v3/configuration_deepseek_v3.py | 6 +- .../deepseek_v3/modeling_deepseek_v3.py | 58 +++++++++++++++--- .../models/deepseek_v3/modular_deepseek_v3.py | 60 ++++++++++++++++--- 3 files changed, 103 insertions(+), 21 deletions(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 5943aa12648e..c0b412dde023 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -135,12 +135,12 @@ class DeepseekV3Config(PretrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `DeepseekV3Model` base_model_tp_plan = { - "layers.*.gate_proj": "colwise", - "layers.*.up_proj": "colwise", - "layers.*.down_proj": "rowwise", "layers.*.self_attn.q_b_proj": "colwise", "layers.*.self_attn.kv_b_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", + "layers.*.gate_proj": "colwise", + "layers.*.up_proj": "colwise", + "layers.*.down_proj": "rowwise", } def __init__( diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index ef6be2f6add5..3757a26cacbc 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -377,7 +377,6 @@ def forward( k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings - q_rot, k_rot = self.reshape_for_rope(q_rot, k_rot) q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) @@ -420,13 +419,6 @@ def forward( attn_output = self.o_proj(attn_output) return attn_output, attn_weights - def reshape_for_rope(self, q, k): - b, h, s, d = q.shape - q = q.view(b, h, s, d // 2, 2).transpose(-1, -2).reshape(b, h, s, d) - b, h, s, d = k.shape - k = k.view(b, h, s, d // 2, 2).transpose(-1, -2).reshape(b, h, s, d) - return q, k - class DeepseekV3DecoderLayer(nn.Module): def __init__(self, config: DeepseekV3Config, layer_idx: int): @@ -629,7 +621,7 @@ class DeepseekV3Model(DeepseekV3PreTrainedModel): config: DeepseekV3Config """ - def __init__(self, config: DeepseekV3Config): + def __init__(self, config): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size @@ -641,6 +633,8 @@ def __init__(self, config: DeepseekV3Config): self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = DeepseekV3RotaryEmbedding(config=config) self.gradient_checkpointing = False + self._register_load_state_dict_pre_hook(self.load_pre_hook) + self._register_state_dict_hook(self.load_hook) # Initialize weights and apply final processing self.post_init() @@ -879,6 +873,52 @@ def _prepare_4d_causal_attention_mask_with_cache_position( return causal_mask + def load_pre_hook(self, state_dict, prefix, *args): + """ + Weights have to be permuted for correct rope formulation. We can't do this in the weights + as every other framework already uses the `Llama` original function (which is copyrighted btw). + And I am not even sure it's better.... anyways end of my rant + """ + + def permute_for_rope(input_tensor): + """ + When you go from the complex ROPE formulation to sin and cos one, you need + to permute the query and key weights (to avoid doing it on the fly) + """ + n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] + input_tensor = input_tensor.reshape(n_heads * dim1, dim2) + input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) + input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) + return input_tensor + + def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): + weight = state_dict[key] + weight = weight.view(num_heads, head_dim, -1) + weight_rot = weight[:, -rope_dim:] + weight_rot = permute_for_rope(weight_rot) + weight[:, -rope_dim:] = weight_rot + weight = weight.view(-1, weight.shape[-1]) + state_dict[key] = weight + + for k in state_dict: + if "q_b_proj." in k: + permute_layer_for_rope( + k, + num_heads=self.config.num_attention_heads, + head_dim=self.config.qk_head_dim, + rope_dim=self.config.qk_rope_head_dim, + ) + if "kv_a_proj_with_mqa." in k: + permute_layer_for_rope( + k, + num_heads=1, + head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, + rope_dim=self.config.qk_rope_head_dim, + ) + + def load_hook(self, module, state_dict, prefix, *args): + self.load_pre_hook(state_dict, prefix, *args) + class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 55d1ac5c691f..e41ded87f2bd 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -225,7 +225,6 @@ def forward( k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings - q_rot, k_rot = self.reshape_for_rope(q_rot, k_rot) q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) @@ -268,13 +267,6 @@ def forward( attn_output = self.o_proj(attn_output) return attn_output, attn_weights - def reshape_for_rope(self, q, k): - b, h, s, d = q.shape - q = q.view(b, h, s, d // 2, 2).transpose(-1, -2).reshape(b, h, s, d) - b, h, s, d = k.shape - k = k.view(b, h, s, d // 2, 2).transpose(-1, -2).reshape(b, h, s, d) - return q, k - class DeepseekV3DecoderLayer(nn.Module): def __init__(self, config: DeepseekV3Config, layer_idx: int): @@ -348,7 +340,57 @@ class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): class DeepseekV3Model(LlamaModel): - pass + def __init__(self, config): + super().__init__(config) + self._register_load_state_dict_pre_hook(self.load_pre_hook) + self._register_state_dict_hook(self.load_hook) + self.post_init() + + def load_pre_hook(self, state_dict, prefix, *args): + """ + Weights have to be permuted for correct rope formulation. We can't do this in the weights + as every other framework already uses the `Llama` original function (which is copyrighted btw). + And I am not even sure it's better.... anyways end of my rant + """ + + def permute_for_rope(input_tensor): + """ + When you go from the complex ROPE formulation to sin and cos one, you need + to permute the query and key weights (to avoid doing it on the fly) + """ + n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] + input_tensor = input_tensor.reshape(n_heads * dim1, dim2) + input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) + input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) + return input_tensor + + def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): + weight = state_dict[key] + weight = weight.view(num_heads, head_dim, -1) + weight_rot = weight[:, -rope_dim:] + weight_rot = permute_for_rope(weight_rot) + weight[:, -rope_dim:] = weight_rot + weight = weight.view(-1, weight.shape[-1]) + state_dict[key] = weight + + for k in state_dict: + if "q_b_proj." in k: + permute_layer_for_rope( + k, + num_heads=self.config.num_attention_heads, + head_dim=self.config.qk_head_dim, + rope_dim=self.config.qk_rope_head_dim, + ) + if "kv_a_proj_with_mqa." in k: + permute_layer_for_rope( + k, + num_heads=1, + head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, + rope_dim=self.config.qk_rope_head_dim, + ) + + def load_hook(self, module, state_dict, prefix, *args): + self.load_pre_hook(state_dict, prefix, *args) class DeepseekV3ForCausalLM(LlamaForCausalLM): From 4ab2f9e8f1a2922da699535afbd53be9593fb418 Mon Sep 17 00:00:00 2001 From: ryan u Date: Mon, 3 Feb 2025 19:06:07 +0900 Subject: [PATCH 29/86] make style --- src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index e41ded87f2bd..98173ac9f237 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -388,7 +388,7 @@ def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, rope_dim=self.config.qk_rope_head_dim, ) - + def load_hook(self, module, state_dict, prefix, *args): self.load_pre_hook(state_dict, prefix, *args) From c5429ec7e58ec7478cf2e1a3b5d17cd5dff5dd3a Mon Sep 17 00:00:00 2001 From: ryan u Date: Mon, 10 Feb 2025 11:05:49 +0900 Subject: [PATCH 30/86] use n_shared_experts --- src/transformers/models/deepseek_v3/modeling_deepseek_v3.py | 4 +++- src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 3757a26cacbc..b7c34226a0dc 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -195,7 +195,9 @@ def __init__(self, config): ] ) self.gate = DeepseekV3TopkRouter(config) - self.shared_experts = DeepseekV3MLP(config=config, intermediate_size=config.moe_intermediate_size) + self.shared_experts = DeepseekV3MLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) def forward(self, hidden_states): residuals = hidden_states diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 98173ac9f237..af83524adca2 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -121,7 +121,9 @@ def __init__(self, config): ] ) self.gate = DeepseekV3TopkRouter(config) - self.shared_experts = DeepseekV3MLP(config=config, intermediate_size=config.moe_intermediate_size) + self.shared_experts = DeepseekV3MLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) def forward(self, hidden_states): residuals = hidden_states From 4df42f0b59619b63d693f272dc04d1c00b0be1b0 Mon Sep 17 00:00:00 2001 From: Minho Ryu Date: Fri, 14 Feb 2025 16:47:42 +0900 Subject: [PATCH 31/86] Update src/transformers/models/deepseek_v3/configuration_deepseek_v3.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- .../models/deepseek_v3/configuration_deepseek_v3.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index c0b412dde023..e459b47b4bc4 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -142,7 +142,11 @@ class DeepseekV3Config(PretrainedConfig): "layers.*.up_proj": "colwise", "layers.*.down_proj": "rowwise", } - + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } def __init__( self, vocab_size=129280, From ba21b7c1c50d09101278e8f3cb1d940cd8b49719 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 17:08:53 +0900 Subject: [PATCH 32/86] add test file --- tests/models/deepseek_v3/__init__.py | 0 .../deepseek_v3/test_modeling_deepseek_v3.py | 1015 +++++++++++++++++ 2 files changed, 1015 insertions(+) create mode 100644 tests/models/deepseek_v3/__init__.py create mode 100644 tests/models/deepseek_v3/test_modeling_deepseek_v3.py diff --git a/tests/models/deepseek_v3/__init__.py b/tests/models/deepseek_v3/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py new file mode 100644 index 000000000000..340af2293c09 --- /dev/null +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -0,0 +1,1015 @@ +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# 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. +"""Testing suite for the PyTorch DiffLlama model.""" + +import gc +import tempfile +import unittest + +import pytest +from packaging import version +from parameterized import parameterized + +from transformers import AutoTokenizer, DeepseekV3Config, StaticCache, is_torch_available, set_seed +from transformers.testing_utils import ( + backend_empty_cache, + require_bitsandbytes, + require_flash_attn, + require_read_token, + require_torch, + require_torch_accelerator, + require_torch_gpu, + require_torch_sdpa, + slow, + torch_device, +) + +from ...generation.test_utils import GenerationTesterMixin +from ...test_configuration_common import ConfigTester +from ...test_modeling_common import ModelTesterMixin, ids_tensor +from ...test_pipeline_mixin import PipelineTesterMixin + + +if is_torch_available(): + import torch + + from transformers import ( + DeepseekV3ForCausalLM, + DeepseekV3ForSequenceClassification, + DeepseekV3Model, + ) + from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3RotaryEmbedding, + ) + + +class DeepseekV3ModelTester: + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_input_mask=True, + use_token_type_ids=False, + use_labels=True, + vocab_size=99, + hidden_size=32, + intermediate_size=37, + moe_intermediate_size=12, + num_hidden_layers=5, + num_attention_heads=4, + num_key_value_heads=4, + n_shared_experts=1, + n_routed_experts=8, + routed_scaling_factor=2.5, + kv_lora_rank=16, + q_lora_rank=32, + qk_rope_head_dim=16, + v_head_dim=32, + qk_nope_head_dim=32, + n_group=2, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=2, + norm_topk_prob=True, + aux_loss_alpha=0.001, + hidden_act="silu", + max_position_embeddings=512, + initializer_range=0.02, + attention_probs_dropout_prob=0.1, + type_vocab_size=16, + type_sequence_label_size=2, + num_labels=3, + num_choices=4, + pad_token_id=0, + scope=None, + ): + self.parent = parent + self.batch_size = batch_size + self.seq_length = seq_length + self.is_training = is_training + self.use_input_mask = use_input_mask + self.use_token_type_ids = use_token_type_ids + self.use_labels = use_labels + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.aux_loss_alpha = aux_loss_alpha + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.type_vocab_size = type_vocab_size + self.type_sequence_label_size = type_sequence_label_size + self.num_labels = num_labels + self.num_choices = num_choices + self.pad_token_id = pad_token_id + self.scope = scope + + def prepare_config_and_inputs(self): + input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + + input_mask = None + if self.use_input_mask: + input_mask = torch.tril(torch.ones_like(input_ids).to(torch_device)) + + token_type_ids = None + if self.use_token_type_ids: + token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size) + + sequence_labels = None + token_labels = None + choice_labels = None + if self.use_labels: + sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size) + token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels) + choice_labels = ids_tensor([self.batch_size], self.num_choices) + + config = self.get_config() + + return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels + + def get_config(self): + return DeepseekV3Config( + vocab_size=self.vocab_size, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + moe_intermediate_size=self.moe_intermediate_size, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + n_shared_experts=self.n_shared_experts, + n_routed_experts=self.n_routed_experts, + routed_scaling_factor=self.routed_scaling_factor, + kv_lora_rank=self.kv_lora_rank, + q_lora_rank=self.q_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + qk_nope_head_dim=self.qk_nope_head_dim, + n_group=self.n_group, + topk_group=self.topk_group, + num_experts_per_tok=self.num_experts_per_tok, + first_k_dense_replace=self.first_k_dense_replace, + norm_topk_prob=self.norm_topk_prob, + aux_loss_alpha=self.aux_loss_alpha, + hidden_act=self.hidden_act, + max_position_embeddings=self.max_position_embeddings, + initializer_range=self.initializer_range, + use_cache=True, + pad_token_id=self.pad_token_id, + attention_dropout=self.attention_probs_dropout_prob, + ) + + def create_and_check_model( + self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels + ): + model = DeepseekV3Model(config=config) + model.to(torch_device) + model.eval() + result = model(input_ids, attention_mask=input_mask) + result = model(input_ids) + self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size)) + + def create_and_check_model_as_decoder( + self, + config, + input_ids, + token_type_ids, + input_mask, + sequence_labels, + token_labels, + choice_labels, + encoder_hidden_states, + encoder_attention_mask, + ): + config.add_cross_attention = True + model = DeepseekV3Model(config) + model.to(torch_device) + model.eval() + result = model( + input_ids, + attention_mask=input_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ) + result = model( + input_ids, + attention_mask=input_mask, + encoder_hidden_states=encoder_hidden_states, + ) + result = model(input_ids, attention_mask=input_mask) + self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size)) + + def create_and_check_for_causal_lm( + self, + config, + input_ids, + token_type_ids, + input_mask, + sequence_labels, + token_labels, + choice_labels, + encoder_hidden_states, + encoder_attention_mask, + ): + model = DeepseekV3ForCausalLM(config=config) + model.to(torch_device) + model.eval() + result = model(input_ids, attention_mask=input_mask, labels=token_labels) + self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size)) + + def create_and_check_decoder_model_past_large_inputs( + self, + config, + input_ids, + token_type_ids, + input_mask, + sequence_labels, + token_labels, + choice_labels, + encoder_hidden_states, + encoder_attention_mask, + ): + config.is_decoder = True + config.add_cross_attention = True + model = DeepseekV3ForCausalLM(config=config) + model.to(torch_device) + model.eval() + + # first forward pass + outputs = model( + input_ids, + attention_mask=input_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=True, + ) + past_key_values = outputs.past_key_values + + # create hypothetical multiple next token and extent to next_input_ids + next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size) + next_mask = ids_tensor((self.batch_size, 3), vocab_size=2) + + # append to next input_ids and + next_input_ids = torch.cat([input_ids, next_tokens], dim=-1) + next_attention_mask = torch.cat([input_mask, next_mask], dim=-1) + + output_from_no_past = model( + next_input_ids, + attention_mask=next_attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_hidden_states=True, + )["hidden_states"][0] + output_from_past = model( + next_tokens, + attention_mask=next_attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + output_hidden_states=True, + )["hidden_states"][0] + + # select random slice + random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item() + output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx].detach() + output_from_past_slice = output_from_past[:, :, random_slice_idx].detach() + + self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]) + + # test that outputs are equal for slice + self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3)) + + def prepare_config_and_inputs_for_common(self): + config_and_inputs = self.prepare_config_and_inputs() + ( + config, + input_ids, + token_type_ids, + input_mask, + sequence_labels, + token_labels, + choice_labels, + ) = config_and_inputs + inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask} + return config, inputs_dict + + +@require_torch +class DiffLlamaModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase): + all_model_classes = ( + ( + DeepseekV3Model, + DeepseekV3ForCausalLM, + DeepseekV3ForSequenceClassification, + ) + if is_torch_available() + else () + ) + all_generative_model_classes = (DeepseekV3ForCausalLM,) if is_torch_available() else () + pipeline_model_mapping = ( + { + "feature-extraction": DeepseekV3Model, + "text-classification": DeepseekV3ForSequenceClassification, + "text-generation": DeepseekV3ForCausalLM, + "zero-shot": DeepseekV3ForSequenceClassification, + } + if is_torch_available() + else {} + ) + test_headmasking = False + test_pruning = False + fx_compatible = False + + # Need to use `0.8` instead of `0.9` for `test_cpu_offload` + # This is because we are hitting edge cases with the causal_mask buffer + model_split_percents = [0.5, 0.7, 0.8] + + # used in `test_torch_compile_for_training` + _torch_compile_train_cls = DeepseekV3ForCausalLM if is_torch_available() else None + + def setUp(self): + self.model_tester = DeepseekV3ModelTester(self) + self.config_tester = ConfigTester(self, config_class=DeepseekV3Config, hidden_size=37) + + def test_config(self): + self.config_tester.run_common_tests() + + def test_model(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_model(*config_and_inputs) + + def test_model_various_embeddings(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + for type in ["absolute", "relative_key", "relative_key_query"]: + config_and_inputs[0].position_embedding_type = type + self.model_tester.create_and_check_model(*config_and_inputs) + + def test_diffllama_sequence_classification_model(self): + config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.num_labels = 3 + input_ids = input_dict["input_ids"] + attention_mask = input_ids.ne(1).to(torch_device) + sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.type_sequence_label_size) + model = DeepseekV3ForSequenceClassification(config) + model.to(torch_device) + model.eval() + result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) + self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) + + def test_diffllama_sequence_classification_model_for_single_label(self): + config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.num_labels = 3 + config.problem_type = "single_label_classification" + input_ids = input_dict["input_ids"] + attention_mask = input_ids.ne(1).to(torch_device) + sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.type_sequence_label_size) + model = DeepseekV3ForSequenceClassification(config) + model.to(torch_device) + model.eval() + result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) + self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) + + def test_diffllama_sequence_classification_model_for_multi_label(self): + config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.num_labels = 3 + config.problem_type = "multi_label_classification" + input_ids = input_dict["input_ids"] + attention_mask = input_ids.ne(1).to(torch_device) + sequence_labels = ids_tensor( + [self.model_tester.batch_size, config.num_labels], self.model_tester.type_sequence_label_size + ).to(torch.float) + model = DeepseekV3ForSequenceClassification(config) + model.to(torch_device) + model.eval() + result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) + self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) + + @unittest.skip(reason="DiffLlama buffers include complex numbers, which breaks this test") + def test_save_load_fast_init_from_base(self): + pass + + @parameterized.expand([("linear",), ("dynamic",), ("yarn",)]) + def test_model_rope_scaling_from_config(self, scaling_type): + config, _ = self.model_tester.prepare_config_and_inputs_for_common() + short_input = ids_tensor([1, 10], config.vocab_size) + long_input = ids_tensor([1, int(config.max_position_embeddings * 1.5)], config.vocab_size) + + set_seed(42) # Fixed seed at init time so the two models get the same random weights + original_model = DeepseekV3Model(config) + original_model.to(torch_device) + original_model.eval() + original_short_output = original_model(short_input).last_hidden_state + original_long_output = original_model(long_input).last_hidden_state + + set_seed(42) # Fixed seed at init time so the two models get the same random weights + config.rope_scaling = {"type": scaling_type, "factor": 10.0} + scaled_model = DeepseekV3Model(config) + scaled_model.to(torch_device) + scaled_model.eval() + scaled_short_output = scaled_model(short_input).last_hidden_state + scaled_long_output = scaled_model(long_input).last_hidden_state + + # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original + # maximum sequence length, so the outputs for the short input should match. + if scaling_type == "dynamic": + torch.testing.assert_close(original_short_output, scaled_short_output, rtol=1e-5, atol=1e-5) + else: + self.assertFalse(torch.allclose(original_short_output, scaled_short_output, atol=1e-5)) + + # The output should be different for long inputs + self.assertFalse(torch.allclose(original_long_output, scaled_long_output, atol=1e-5)) + + def test_model_rope_scaling(self): + config, _ = self.model_tester.prepare_config_and_inputs_for_common() + scaling_factor = 10 + short_input_length = 10 + long_input_length = int(config.max_position_embeddings * 1.5) + + # Inputs + x = torch.randn(1, dtype=torch.float32, device=torch_device) # used exlusively to get the dtype and the device + position_ids_short = torch.arange(short_input_length, dtype=torch.long, device=torch_device) + position_ids_short = position_ids_short.unsqueeze(0) + position_ids_long = torch.arange(long_input_length, dtype=torch.long, device=torch_device) + position_ids_long = position_ids_long.unsqueeze(0) + + # Sanity check original RoPE + original_rope = DeepseekV3RotaryEmbedding(config=config).to(torch_device) + original_cos_short, original_sin_short = original_rope(x, position_ids_short) + original_cos_long, original_sin_long = original_rope(x, position_ids_long) + torch.testing.assert_close(original_cos_short, original_cos_long[:, :short_input_length, :]) + torch.testing.assert_close(original_sin_short, original_sin_long[:, :short_input_length, :]) + + # Sanity check linear RoPE scaling + # New position "x" should match original position with index "x/scaling_factor" + config.rope_scaling = {"type": "linear", "factor": scaling_factor} + linear_scaling_rope = DeepseekV3RotaryEmbedding(config=config).to(torch_device) + linear_cos_short, linear_sin_short = linear_scaling_rope(x, position_ids_short) + linear_cos_long, linear_sin_long = linear_scaling_rope(x, position_ids_long) + torch.testing.assert_close(linear_cos_short, linear_cos_long[:, :short_input_length, :]) + torch.testing.assert_close(linear_sin_short, linear_sin_long[:, :short_input_length, :]) + for new_position in range(0, long_input_length, scaling_factor): + original_position = int(new_position // scaling_factor) + torch.testing.assert_close(linear_cos_long[:, new_position, :], original_cos_long[:, original_position, :]) + torch.testing.assert_close(linear_sin_long[:, new_position, :], original_sin_long[:, original_position, :]) + + # Sanity check Dynamic NTK RoPE scaling + # Scaling should only be observed after a long input is fed. We can observe that the frequencies increase + # with scaling_factor (or that `inv_freq` decreases) + config.rope_scaling = {"type": "dynamic", "factor": scaling_factor} + ntk_scaling_rope = DeepseekV3RotaryEmbedding(config=config).to(torch_device) + ntk_cos_short, ntk_sin_short = ntk_scaling_rope(x, position_ids_short) + ntk_cos_long, ntk_sin_long = ntk_scaling_rope(x, position_ids_long) + torch.testing.assert_close(ntk_cos_short, original_cos_short) + torch.testing.assert_close(ntk_sin_short, original_sin_short) + with self.assertRaises(AssertionError): + torch.testing.assert_close(ntk_cos_long, original_cos_long) + with self.assertRaises(AssertionError): + torch.testing.assert_close(ntk_sin_long, original_sin_long) + self.assertTrue((ntk_scaling_rope.inv_freq <= original_rope.inv_freq).all()) + + # Sanity check Yarn RoPE scaling + # Scaling should be over the entire input + config.rope_scaling = {"type": "yarn", "factor": scaling_factor} + yarn_scaling_rope = DeepseekV3RotaryEmbedding(config=config).to(torch_device) + yarn_cos_short, yarn_sin_short = yarn_scaling_rope(x, position_ids_short) + yarn_cos_long, yarn_sin_long = yarn_scaling_rope(x, position_ids_long) + torch.testing.assert_close(yarn_cos_short, yarn_cos_long[:, :short_input_length, :]) + torch.testing.assert_close(yarn_sin_short, yarn_sin_long[:, :short_input_length, :]) + with self.assertRaises(AssertionError): + torch.testing.assert_close(yarn_cos_short, original_cos_short) + with self.assertRaises(AssertionError): + torch.testing.assert_close(yarn_sin_short, original_sin_short) + with self.assertRaises(AssertionError): + torch.testing.assert_close(yarn_cos_long, original_cos_long) + with self.assertRaises(AssertionError): + torch.testing.assert_close(yarn_sin_long, original_sin_long) + + def test_model_loading_old_rope_configs(self): + def _reinitialize_config(base_config, new_kwargs): + # Reinitialize the config with the new kwargs, forcing the config to go through its __init__ validation + # steps. + base_config_dict = base_config.to_dict() + new_config = DeepseekV3Config.from_dict(config_dict={**base_config_dict, **new_kwargs}) + return new_config + + # from untouched config -> βœ… + base_config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common() + original_model = DeepseekV3ForCausalLM(base_config).to(torch_device) + original_model(**model_inputs) + + # from a config with the expected rope configuration -> βœ… + config = _reinitialize_config(base_config, {"rope_scaling": {"rope_type": "linear", "factor": 10.0}}) + original_model = DeepseekV3ForCausalLM(config).to(torch_device) + original_model(**model_inputs) + + # from a config with the old rope configuration ('type' instead of 'rope_type') -> βœ… we gracefully handle BC + config = _reinitialize_config(base_config, {"rope_scaling": {"type": "linear", "factor": 10.0}}) + original_model = DeepseekV3ForCausalLM(config).to(torch_device) + original_model(**model_inputs) + + # from a config with both 'type' and 'rope_type' -> βœ… they can coexist (and both are present in the config) + config = _reinitialize_config( + base_config, {"rope_scaling": {"type": "linear", "rope_type": "linear", "factor": 10.0}} + ) + self.assertTrue(config.rope_scaling["type"] == "linear") + self.assertTrue(config.rope_scaling["rope_type"] == "linear") + original_model = DeepseekV3ForCausalLM(config).to(torch_device) + original_model(**model_inputs) + + # from a config with parameters in a bad range ('factor' should be >= 1.0) -> ⚠️ throws a warning + with self.assertLogs("transformers.modeling_rope_utils", level="WARNING") as logs: + config = _reinitialize_config(base_config, {"rope_scaling": {"rope_type": "linear", "factor": -999.0}}) + original_model = DeepseekV3ForCausalLM(config).to(torch_device) + original_model(**model_inputs) + self.assertEqual(len(logs.output), 1) + self.assertIn("factor field", logs.output[0]) + + # from a config with unknown parameters ('foo' isn't a rope option) -> ⚠️ throws a warning + with self.assertLogs("transformers.modeling_rope_utils", level="WARNING") as logs: + config = _reinitialize_config( + base_config, {"rope_scaling": {"rope_type": "linear", "factor": 10.0, "foo": "bar"}} + ) + original_model = DeepseekV3ForCausalLM(config).to(torch_device) + original_model(**model_inputs) + self.assertEqual(len(logs.output), 1) + self.assertIn("Unrecognized keys", logs.output[0]) + + # from a config with specific rope type but missing one of its mandatory parameters -> ❌ throws exception + with self.assertRaises(KeyError): + config = _reinitialize_config(base_config, {"rope_scaling": {"rope_type": "linear"}}) # missing "factor" + + @require_flash_attn + @require_torch_gpu + @require_bitsandbytes + @pytest.mark.flash_attn_test + @require_read_token + @slow + def test_flash_attn_2_generate_padding_right(self): + """ + Overwritting the common test as the test is flaky on tiny models + """ + model = DeepseekV3ForCausalLM.from_pretrained( + "bzantium/deepseek-v3-test", + load_in_4bit=True, + device_map={"": 0}, + ) + + tokenizer = AutoTokenizer.from_pretrained("bzantium/deepseek-v3-test") + + texts = ["hi", "Hello this is a very long sentence"] + + tokenizer.padding_side = "right" + tokenizer.pad_token = tokenizer.eos_token + + inputs = tokenizer(texts, return_tensors="pt", padding=True).to(0) + + output_native = model.generate(**inputs, max_new_tokens=20, do_sample=False) + output_native = tokenizer.batch_decode(output_native) + + model = DeepseekV3ForCausalLM.from_pretrained( + "bzantium/deepseek-v3-test", + load_in_4bit=True, + device_map={"": 0}, + attn_implementation="flash_attention_2", + ) + + output_fa_2 = model.generate(**inputs, max_new_tokens=20, do_sample=False) + output_fa_2 = tokenizer.batch_decode(output_fa_2) + + self.assertListEqual(output_native, output_fa_2) + + @require_flash_attn + @require_torch_gpu + @slow + @pytest.mark.flash_attn_test + def test_use_flash_attention_2_true(self): + """ + NOTE: this is the only test testing that the legacy `use_flash_attention=2` argument still works as intended. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + for model_class in self.all_model_classes: + with tempfile.TemporaryDirectory() as tmp_dir: + model = model_class(config) + model.save_pretrained(tmp_dir) + + new_model = DeepseekV3ForCausalLM.from_pretrained( + tmp_dir, use_flash_attention_2=True, torch_dtype=torch.float16 + ).to("cuda") + + self.assertTrue(new_model.config._attn_implementation == "flash_attention_2") + + has_flash = False + for name, submodule in new_model.named_modules(): + if "FlashAttention" in submodule.__class__.__name__: + has_flash = True + break + if not has_flash: + raise ValueError("The flash model should have flash attention layers") + + @require_torch_sdpa + @slow + def test_eager_matches_sdpa_generate(self): + """ + Overwritting the common test as the test is flaky on tiny models + """ + max_new_tokens = 30 + + tokenizer = AutoTokenizer.from_pretrained("bzantium/deepseek-v3-test") + + model_sdpa = DeepseekV3ForCausalLM.from_pretrained( + "bzantium/deepseek-v3-test", + torch_dtype=torch.float16, + low_cpu_mem_usage=True, + ).to(torch_device) + + self.assertTrue(model_sdpa.config._attn_implementation == "sdpa") + + model_eager = DeepseekV3ForCausalLM.from_pretrained( + "bzantium/deepseek-v3-test", + torch_dtype=torch.float16, + low_cpu_mem_usage=True, + attn_implementation="eager", + ).to(torch_device) + + self.assertTrue(model_eager.config._attn_implementation == "eager") + + for name, submodule in model_eager.named_modules(): + if "SdpaAttention" in submodule.__class__.__name__: + raise ValueError("The eager model should not have SDPA attention layers") + + has_sdpa = False + for name, submodule in model_sdpa.named_modules(): + if "SdpaAttention" in submodule.__class__.__name__: + has_sdpa = True + break + if not has_sdpa: + raise ValueError("The SDPA model should have SDPA attention layers") + + texts = [ + "hi here's a longer context, getting longer and", + "Hello this is a very long sentence my friend, very long for real", + "Today I am in Paris and", + ] + + for padding_side in ["left", "right"]: + tokenizer.padding_side = padding_side + tokenizer.pad_token = tokenizer.eos_token + + inputs = tokenizer(texts, return_tensors="pt", padding=True).to(torch_device) + + res_eager = model_eager.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) + res_sdpa = model_sdpa.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) + + with self.subTest(f"{padding_side}"): + torch.testing.assert_close( + res_eager, + res_sdpa, + msg=f"\n{tokenizer.batch_decode(res_eager)} \nvs\n{tokenizer.batch_decode(res_sdpa)}", + ) + + +@require_torch_accelerator +class DiffLlamaIntegrationTest(unittest.TestCase): + # This variable is used to determine which CUDA device are we using for our runners (A10 or T4) + # Depending on the hardware we get different logits / generations + cuda_compute_capability_major_version = None + + @classmethod + def setUpClass(cls): + if is_torch_available() and torch.cuda.is_available(): + # 8 is for A100 / A10 and 7 for T4 + cls.cuda_compute_capability_major_version = torch.cuda.get_device_capability()[0] + + @slow + @require_torch_accelerator + @require_read_token + def test_compile_static_cache(self): + # `torch==2.2` will throw an error on this test (as in other compilation tests), but torch==2.1.2 and torch>2.2 + # work as intended. See https://github.com/pytorch/pytorch/issues/121943 + if version.parse(torch.__version__) < version.parse("2.3.0"): + self.skipTest(reason="This test requires torch >= 2.3 to run.") + + NUM_TOKENS_TO_GENERATE = 40 + # Note on `EXPECTED_TEXT_COMPLETION`'s diff: the current value matches the original test if the original test + # was changed to have a cache of 53 tokens (as opposed to 4096), on Ampere GPUs. + EXPECTED_TEXT_COMPLETION = [ + "Simply put, the theory of relativity states that 1) the speed of light is constant in all inertial " + "reference frames, and 2) the laws of physics are the same for all inertial reference frames.\nThe " + "theory of relativ", + "My favorite all time favorite condiment is ketchup. I love it on everything. I love it on my eggs, " + "my fries, my chicken, my burgers, my hot dogs, my sandwiches, my salads, my p", + ] + + prompts = [ + "Simply put, the theory of relativity states that ", + "My favorite all time favorite condiment is ketchup.", + ] + tokenizer = AutoTokenizer.from_pretrained( + "bzantium/deepseek-v3-test", pad_token="", padding_side="right" + ) + model = DeepseekV3ForCausalLM.from_pretrained( + "bzantium/deepseek-v3-test", device_map=torch_device, torch_dtype=torch.float16 + ) + inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device) + + # Dynamic Cache + generated_ids = model.generate(**inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False) + dynamic_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) + self.assertEqual(EXPECTED_TEXT_COMPLETION, dynamic_text) + + # Static Cache + generated_ids = model.generate( + **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, cache_implementation="static" + ) + static_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) + self.assertEqual(EXPECTED_TEXT_COMPLETION, static_text) + + # Static Cache + compile + model._cache = None # clear cache object, initialized when we pass `cache_implementation="static"` + model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True) + generated_ids = model.generate( + **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, cache_implementation="static" + ) + static_compiled_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) + self.assertEqual(EXPECTED_TEXT_COMPLETION, static_compiled_text) + + +@slow +@require_torch_accelerator +class Mask4DTestHard(unittest.TestCase): + def tearDown(self): + gc.collect() + backend_empty_cache(torch_device) + + def setUp(self): + model_name = "bzantium/deepseek-v3-test" + self.model_dtype = torch.float32 + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + self.model = DeepseekV3ForCausalLM.from_pretrained(model_name, torch_dtype=self.model_dtype).to(torch_device) + + def get_test_data(self): + template = "my favorite {}" + items = ("pet is a", "artist plays a", "name is L") # same number of tokens in each item + + batch_separate = [template.format(x) for x in items] # 3 separate lines + batch_shared_prefix = template.format(" ".join(items)) # 1 line with options concatenated + + input_ids = self.tokenizer(batch_separate, return_tensors="pt").input_ids.to(torch_device) + input_ids_shared_prefix = self.tokenizer(batch_shared_prefix, return_tensors="pt").input_ids.to(torch_device) + + mask_shared_prefix = torch.tensor( + [ + [ + [ + [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0], + [1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0], + [1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1], + ] + ] + ], + device=torch_device, + ) + + position_ids = torch.arange(input_ids.shape[1]).tile(input_ids.shape[0], 1).to(torch_device) + + # building custom positions ids based on custom mask + position_ids_shared_prefix = (mask_shared_prefix.sum(dim=-1) - 1).reshape(1, -1) + # effectively: position_ids_shared_prefix = torch.tensor([[0, 1, 2, 3, 4, 5, 3, 4, 5, 3, 4, 5]]).to(device) + + # inverting the mask + min_dtype = torch.finfo(self.model_dtype).min + mask_shared_prefix = (mask_shared_prefix.eq(0.0)).to(dtype=self.model_dtype) * min_dtype + + return input_ids, position_ids, input_ids_shared_prefix, mask_shared_prefix, position_ids_shared_prefix + + def test_stacked_causal_mask(self): + ( + input_ids, + position_ids, + input_ids_shared_prefix, + mask_shared_prefix, + position_ids_shared_prefix, + ) = self.get_test_data() + + # regular batch + logits = self.model.forward(input_ids, position_ids=position_ids).logits + logits_last = logits[:, -1, :] # last tokens in each batch line + decoded = [self.tokenizer.decode(t) for t in logits_last.argmax(dim=-1)] + + # single forward run with 4D custom mask + logits_shared_prefix = self.model.forward( + input_ids_shared_prefix, attention_mask=mask_shared_prefix, position_ids=position_ids_shared_prefix + ).logits + logits_shared_prefix_last = logits_shared_prefix[ + 0, torch.where(position_ids_shared_prefix == position_ids_shared_prefix.max())[1], : + ] # last three tokens + decoded_shared_prefix = [self.tokenizer.decode(t) for t in logits_shared_prefix_last.argmax(dim=-1)] + + self.assertEqual(decoded, decoded_shared_prefix) + + def test_partial_stacked_causal_mask(self): + # Same as the test above, but the input is passed in two groups. It tests that we can pass partial 4D attention masks + + ( + input_ids, + position_ids, + input_ids_shared_prefix, + mask_shared_prefix, + position_ids_shared_prefix, + ) = self.get_test_data() + + # regular batch + logits = self.model.forward(input_ids, position_ids=position_ids).logits + logits_last = logits[:, -1, :] # last tokens in each batch line + decoded = [self.tokenizer.decode(t) for t in logits_last.argmax(dim=-1)] + + # 2 forward runs with custom 4D masks + part_a = 3 # split point + + input_1a = input_ids_shared_prefix[:, :part_a] + position_ids_1a = position_ids_shared_prefix[:, :part_a] + mask_1a = mask_shared_prefix[:, :, :part_a, :part_a] + + outs_1a = self.model.forward(input_1a, attention_mask=mask_1a, position_ids=position_ids_1a) + past_key_values_a = outs_1a["past_key_values"] + + # Case 1: we pass a 4D attention mask regarding the current sequence length (i.e. [..., seq_len, full_len]) + input_1b = input_ids_shared_prefix[:, part_a:] + position_ids_1b = position_ids_shared_prefix[:, part_a:] + mask_1b = mask_shared_prefix[:, :, part_a:, :] + outs_1b = self.model.forward( + input_1b, + attention_mask=mask_1b, + position_ids=position_ids_1b, + past_key_values=past_key_values_a, + ) + decoded_1b = [ + self.tokenizer.decode(t) + for t in outs_1b.logits.argmax(-1)[ + 0, torch.where(position_ids_shared_prefix == position_ids_shared_prefix.max())[1] - part_a + ] + ] + self.assertEqual(decoded, decoded_1b) + + def test_stacked_causal_mask_static_cache(self): + """same as above but with StaticCache""" + ( + input_ids, + position_ids, + input_ids_shared_prefix, + mask_shared_prefix, + position_ids_shared_prefix, + ) = self.get_test_data() + + # regular batch + logits = self.model.forward(input_ids, position_ids=position_ids).logits + logits_last = logits[:, -1, :] # last tokens in each batch line + decoded = [self.tokenizer.decode(t) for t in logits_last.argmax(dim=-1)] + + # upgrade the model with StaticCache + max_cache_len = 16 # note that max_cache_len is greater than the attention_mask.shape[-1] + past_key_values = StaticCache( + config=self.model.config, + batch_size=1, + max_cache_len=max_cache_len, + device=torch_device, + dtype=self.model.dtype, + ) + + padded_attention_mask = torch.nn.functional.pad( + input=mask_shared_prefix, + pad=(0, max_cache_len - mask_shared_prefix.shape[-1]), + mode="constant", + value=torch.finfo(self.model_dtype).min, + ) + + # single forward run with 4D custom mask + logits_shared_prefix = self.model.forward( + input_ids_shared_prefix, + attention_mask=padded_attention_mask, + position_ids=position_ids_shared_prefix, + cache_position=torch.arange(input_ids_shared_prefix.shape[-1], device=torch_device), + past_key_values=past_key_values, + ).logits + logits_shared_prefix_last = logits_shared_prefix[ + 0, torch.where(position_ids_shared_prefix == position_ids_shared_prefix.max())[1], : + ] # last three tokens + decoded_shared_prefix = [self.tokenizer.decode(t) for t in logits_shared_prefix_last.argmax(dim=-1)] + + self.assertEqual(decoded, decoded_shared_prefix) + + def test_partial_stacked_causal_mask_static_cache(self): + # Same as the test above, but the input is passed in two groups. It tests that we can pass partial 4D attention masks + # we pass a 4D attention mask shaped [..., seq_len, full_static_cache_len]) + ( + input_ids, + position_ids, + input_ids_shared_prefix, + mask_shared_prefix, + position_ids_shared_prefix, + ) = self.get_test_data() + + # regular batch + logits = self.model.forward(input_ids, position_ids=position_ids).logits + logits_last = logits[:, -1, :] # last tokens in each batch line + decoded = [self.tokenizer.decode(t) for t in logits_last.argmax(dim=-1)] + + # upgrade the model with StaticCache + max_cache_len = 16 # note that max_cache_len is greater than the attention_mask.shape[-1] + past_key_values = StaticCache( + config=self.model.config, + batch_size=1, + max_cache_len=max_cache_len, + device=torch_device, + dtype=self.model.dtype, + ) + + # forward run for the first part of input + part_a = 3 # split point + + input_1a = input_ids_shared_prefix[:, :part_a] + position_ids_1a = position_ids_shared_prefix[:, :part_a] + mask_1a = mask_shared_prefix[:, :, :part_a, :part_a] + + padded_mask_1a = torch.nn.functional.pad( + input=mask_1a, + pad=(0, max_cache_len - mask_1a.shape[-1]), + mode="constant", + value=torch.finfo(self.model_dtype).min, + ) + + _ = self.model.forward( + input_1a, + attention_mask=padded_mask_1a, + position_ids=position_ids_1a, + cache_position=torch.arange(part_a, device=torch_device), + past_key_values=past_key_values, + ) + + # forward run for the second part of input + input_1b = input_ids_shared_prefix[:, part_a:] + position_ids_1b = position_ids_shared_prefix[:, part_a:] + mask_1b = mask_shared_prefix[:, :, part_a:, :] + + padded_mask_1b = torch.nn.functional.pad( + input=mask_1b, pad=(0, max_cache_len - mask_1b.shape[-1]), mode="constant", value=0 + ) + + outs_1b = self.model.forward( + input_1b, + attention_mask=padded_mask_1b, + position_ids=position_ids_1b, + cache_position=torch.arange( + part_a, + input_ids_shared_prefix.shape[-1], + device=torch_device, + ), + past_key_values=past_key_values, + ) + decoded_1b = [ + self.tokenizer.decode(t) + for t in outs_1b.logits.argmax(-1)[ + 0, torch.where(position_ids_shared_prefix == position_ids_shared_prefix.max())[1] - part_a + ] + ] + self.assertEqual(decoded, decoded_1b) From b5f420b75f6f754ee1df3e7d5be2adedc182001d Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 17:18:45 +0900 Subject: [PATCH 33/86] update modeling_file according to modular file --- .../deepseek_v3/modeling_deepseek_v3.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index b7c34226a0dc..b7abe24d0fc8 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -808,7 +808,7 @@ def _update_causal_mask( if ( self.config._attn_implementation == "sdpa" and attention_mask is not None - and attention_mask.device.type == "cuda" + and attention_mask.device.type in ["cuda", "xpu"] and not output_attentions ): # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when @@ -867,7 +867,9 @@ def _prepare_4d_causal_attention_mask_with_cache_position( if attention_mask is not None: causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( + causal_mask.device + ) padding_mask = padding_mask == 0 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( padding_mask, min_dtype @@ -928,6 +930,7 @@ class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): super().__init__(config) @@ -1125,17 +1128,20 @@ def forward( if self.config.pad_token_id is None and batch_size != 1: raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") if self.config.pad_token_id is None: - sequence_lengths = -1 + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) else: - if input_ids is not None: - # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility - sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 - sequence_lengths = sequence_lengths % input_ids.shape[-1] - sequence_lengths = sequence_lengths.to(logits.device) - else: - sequence_lengths = -1 + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) - pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] loss = None if labels is not None: From 6bd75a9ad34c94c4720fb81a617351d96d48fdd2 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 17:24:49 +0900 Subject: [PATCH 34/86] make style --- .../models/deepseek_v3/configuration_deepseek_v3.py | 1 + tests/models/deepseek_v3/test_modeling_deepseek_v3.py | 4 +--- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index e459b47b4bc4..8787e63f0871 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -147,6 +147,7 @@ class DeepseekV3Config(PretrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + def __init__( self, vocab_size=129280, diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index 340af2293c09..b939e637e40b 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -735,9 +735,7 @@ def test_compile_static_cache(self): "Simply put, the theory of relativity states that ", "My favorite all time favorite condiment is ketchup.", ] - tokenizer = AutoTokenizer.from_pretrained( - "bzantium/deepseek-v3-test", pad_token="", padding_side="right" - ) + tokenizer = AutoTokenizer.from_pretrained("bzantium/deepseek-v3-test", pad_token="", padding_side="right") model = DeepseekV3ForCausalLM.from_pretrained( "bzantium/deepseek-v3-test", device_map=torch_device, torch_dtype=torch.float16 ) From 6ccbc6631af9557da8c0b5761243bd843adc853b Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 17:34:49 +0900 Subject: [PATCH 35/86] add mapping for DeepseekV3ForSequenceClassification --- src/transformers/models/auto/modeling_auto.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 0409c5d96fad..72006d5f0479 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -991,6 +991,7 @@ ("data2vec-text", "Data2VecTextForSequenceClassification"), ("deberta", "DebertaForSequenceClassification"), ("deberta-v2", "DebertaV2ForSequenceClassification"), + ("deepseek_v3", "DeepseekV3ForSequenceClassification"), ("diffllama", "DiffLlamaForSequenceClassification"), ("distilbert", "DistilBertForSequenceClassification"), ("electra", "ElectraForSequenceClassification"), From a1c6274362d44bdcd49e70b408f89879262ce872 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 17:46:35 +0900 Subject: [PATCH 36/86] remove aux_loss_alpha --- .../models/deepseek_v3/configuration_deepseek_v3.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 8787e63f0871..24c2fb66cbbb 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -28,7 +28,7 @@ class DeepseekV3Config(PretrainedConfig): This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the DeepSeek-V3. - + e.g. [bzantium/tiny-deepseek-v3](https://huggingface.co/bzantium/tiny-deepseek-v3) Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. @@ -82,9 +82,6 @@ class DeepseekV3Config(PretrainedConfig): \--k dense layers--/ norm_topk_prob (`bool`, *optional*, defaults to `True`): Whether to normalize the weights of the routed experts. - aux_loss_alpha (`float`, *optional*, defaults to 0.001): - Auxiliary loss weight coefficient. - Whether to compute the auxiliary loss for each individual sample. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. max_position_embeddings (`int`, *optional*, defaults to 4096): @@ -170,7 +167,6 @@ def __init__( num_experts_per_tok=8, first_k_dense_replace=3, norm_topk_prob=True, - aux_loss_alpha=0.001, hidden_act="silu", max_position_embeddings=4096, initializer_range=0.02, @@ -209,7 +205,6 @@ def __init__( self.num_experts_per_tok = num_experts_per_tok self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob - self.aux_loss_alpha = aux_loss_alpha # for backward compatibility if num_key_value_heads is None: From a80462b34258cba28737d7b271a739b7fcf2bad4 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 17:56:37 +0900 Subject: [PATCH 37/86] add deepseek_v3 for perf --- docs/source/en/perf_infer_gpu_one.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/en/perf_infer_gpu_one.md b/docs/source/en/perf_infer_gpu_one.md index 59b686436018..7715d606e45f 100644 --- a/docs/source/en/perf_infer_gpu_one.md +++ b/docs/source/en/perf_infer_gpu_one.md @@ -47,6 +47,7 @@ FlashAttention-2 is currently supported for the following architectures: * [Cohere2](https://huggingface.co/docs/transformers/model_doc/cohere2#transformers.Cohere2Model) * [GLM](https://huggingface.co/docs/transformers/model_doc/glm#transformers.GLMModel) * [Dbrx](https://huggingface.co/docs/transformers/model_doc/dbrx#transformers.DbrxModel) +* [DeepseekV3](https://huggingface.co/docs/transformers/model_doc/deepseek_v3#transformers.DeepseekV3Model) * [DiffLlama](https://huggingface.co/docs/transformers/model_doc/diffllama#transformers.DiffLlamaModel) * [DistilBert](https://huggingface.co/docs/transformers/model_doc/distilbert#transformers.DistilBertModel) * [Emu3](https://huggingface.co/docs/transformers/model_doc/emu3) @@ -245,6 +246,7 @@ For now, Transformers supports SDPA inference and training for the following arc * [Dbrx](https://huggingface.co/docs/transformers/model_doc/dbrx#transformers.DbrxModel) * [DeiT](https://huggingface.co/docs/transformers/model_doc/deit#transformers.DeiTModel) * [DepthPro](https://huggingface.co/docs/transformers/model_doc/depth_pro#transformers.DepthProModel) +* [DeepseekV3](https://huggingface.co/docs/transformers/model_doc/deepseek_v3#transformers.DeepseekV3Model) * [DiffLlama](https://huggingface.co/docs/transformers/model_doc/diffllama#transformers.DiffLlamaModel) * [Dinov2](https://huggingface.co/docs/transformers/en/model_doc/dinov2) * [Dinov2_with_registers](https://huggingface.co/docs/transformers/en/model_doc/dinov2) From dd78f48c0ef0a23f4638cefa40b5da599f9572d8 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 17:58:35 +0900 Subject: [PATCH 38/86] add deepseek_v3 --- docs/source/en/_toctree.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index dc259103ae2e..1c8c2ad48916 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -395,6 +395,8 @@ title: DeBERTa - local: model_doc/deberta-v2 title: DeBERTa-v2 + - local: model_doc/deepseek_v3 + title: DeepSeek-V3 - local: model_doc/dialogpt title: DialoGPT - local: model_doc/diffllama From 54481ef59a25f1b81fb2e9f7ec1a89012e8f3867 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 18:10:14 +0900 Subject: [PATCH 39/86] rename test as deepseekv3 --- .../deepseek_v3/test_modeling_deepseek_v3.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index b939e637e40b..5e99b5812ccf 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -12,7 +12,7 @@ # 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. -"""Testing suite for the PyTorch DiffLlama model.""" +"""Testing suite for the PyTorch DeepseekV3 model.""" import gc import tempfile @@ -326,7 +326,7 @@ def prepare_config_and_inputs_for_common(self): @require_torch -class DiffLlamaModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase): +class DeepseekV3ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase): all_model_classes = ( ( DeepseekV3Model, @@ -375,7 +375,7 @@ def test_model_various_embeddings(self): config_and_inputs[0].position_embedding_type = type self.model_tester.create_and_check_model(*config_and_inputs) - def test_diffllama_sequence_classification_model(self): + def test_deepseek_v3_sequence_classification_model(self): config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() config.num_labels = 3 input_ids = input_dict["input_ids"] @@ -387,7 +387,7 @@ def test_diffllama_sequence_classification_model(self): result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) - def test_diffllama_sequence_classification_model_for_single_label(self): + def test_deepseek_v3_sequence_classification_model_for_single_label(self): config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() config.num_labels = 3 config.problem_type = "single_label_classification" @@ -400,7 +400,7 @@ def test_diffllama_sequence_classification_model_for_single_label(self): result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) - def test_diffllama_sequence_classification_model_for_multi_label(self): + def test_deepseek_v3_sequence_classification_model_for_multi_label(self): config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() config.num_labels = 3 config.problem_type = "multi_label_classification" @@ -415,7 +415,7 @@ def test_diffllama_sequence_classification_model_for_multi_label(self): result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) - @unittest.skip(reason="DiffLlama buffers include complex numbers, which breaks this test") + @unittest.skip(reason="DeepseekV3 buffers include complex numbers, which breaks this test") def test_save_load_fast_init_from_base(self): pass @@ -700,7 +700,7 @@ def test_eager_matches_sdpa_generate(self): @require_torch_accelerator -class DiffLlamaIntegrationTest(unittest.TestCase): +class DeepseekV3IntegrationTest(unittest.TestCase): # This variable is used to determine which CUDA device are we using for our runners (A10 or T4) # Depending on the hardware we get different logits / generations cuda_compute_capability_major_version = None From e0f1c2dd649d29a780d275323fdd335678bd2928 Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 18:19:57 +0900 Subject: [PATCH 40/86] use tiny-deepseek-v3 --- .../deepseek_v3/test_modeling_deepseek_v3.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index 5e99b5812ccf..a9e3475dfa5b 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -580,12 +580,12 @@ def test_flash_attn_2_generate_padding_right(self): Overwritting the common test as the test is flaky on tiny models """ model = DeepseekV3ForCausalLM.from_pretrained( - "bzantium/deepseek-v3-test", + "bzantium/tiny-deepseek-v3", load_in_4bit=True, device_map={"": 0}, ) - tokenizer = AutoTokenizer.from_pretrained("bzantium/deepseek-v3-test") + tokenizer = AutoTokenizer.from_pretrained("bzantium/tiny-deepseek-v3") texts = ["hi", "Hello this is a very long sentence"] @@ -598,7 +598,7 @@ def test_flash_attn_2_generate_padding_right(self): output_native = tokenizer.batch_decode(output_native) model = DeepseekV3ForCausalLM.from_pretrained( - "bzantium/deepseek-v3-test", + "bzantium/tiny-deepseek-v3", load_in_4bit=True, device_map={"": 0}, attn_implementation="flash_attention_2", @@ -645,10 +645,10 @@ def test_eager_matches_sdpa_generate(self): """ max_new_tokens = 30 - tokenizer = AutoTokenizer.from_pretrained("bzantium/deepseek-v3-test") + tokenizer = AutoTokenizer.from_pretrained("bzantium/tiny-deepseek-v3") model_sdpa = DeepseekV3ForCausalLM.from_pretrained( - "bzantium/deepseek-v3-test", + "bzantium/tiny-deepseek-v3", torch_dtype=torch.float16, low_cpu_mem_usage=True, ).to(torch_device) @@ -656,7 +656,7 @@ def test_eager_matches_sdpa_generate(self): self.assertTrue(model_sdpa.config._attn_implementation == "sdpa") model_eager = DeepseekV3ForCausalLM.from_pretrained( - "bzantium/deepseek-v3-test", + "bzantium/tiny-deepseek-v3", torch_dtype=torch.float16, low_cpu_mem_usage=True, attn_implementation="eager", @@ -735,9 +735,9 @@ def test_compile_static_cache(self): "Simply put, the theory of relativity states that ", "My favorite all time favorite condiment is ketchup.", ] - tokenizer = AutoTokenizer.from_pretrained("bzantium/deepseek-v3-test", pad_token="", padding_side="right") + tokenizer = AutoTokenizer.from_pretrained("bzantium/tiny-deepseek-v3", pad_token="", padding_side="right") model = DeepseekV3ForCausalLM.from_pretrained( - "bzantium/deepseek-v3-test", device_map=torch_device, torch_dtype=torch.float16 + "bzantium/tiny-deepseek-v3", device_map=torch_device, torch_dtype=torch.float16 ) inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device) @@ -771,7 +771,7 @@ def tearDown(self): backend_empty_cache(torch_device) def setUp(self): - model_name = "bzantium/deepseek-v3-test" + model_name = "bzantium/tiny-deepseek-v3" self.model_dtype = torch.float32 self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = DeepseekV3ForCausalLM.from_pretrained(model_name, torch_dtype=self.model_dtype).to(torch_device) From 52147415a70cf008ed2eebc1bcde4a5e2d5c754a Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 18:39:00 +0900 Subject: [PATCH 41/86] remove DeepseekV3ForSequenceClassification --- docs/source/en/model_doc/deepseek_v3.md | 5 - src/transformers/__init__.py | 2 - src/transformers/models/auto/modeling_auto.py | 1 - .../deepseek_v3/modeling_deepseek_v3.py | 116 +------ .../models/deepseek_v3/modular_deepseek_v3.py | 6 - src/transformers/utils/dummy_pt_objects.py | 7 - .../deepseek_v3/test_modeling_deepseek_v3.py | 298 +----------------- 7 files changed, 3 insertions(+), 432 deletions(-) diff --git a/docs/source/en/model_doc/deepseek_v3.md b/docs/source/en/model_doc/deepseek_v3.md index 596805cdb70f..e44d78739069 100644 --- a/docs/source/en/model_doc/deepseek_v3.md +++ b/docs/source/en/model_doc/deepseek_v3.md @@ -39,8 +39,3 @@ The model uses Multi-head Latent Attention (MLA) and DeepSeekMoE architectures f [[autodoc]] DeepseekV3ForCausalLM - forward - -## DeepseekV3ForSequenceClassification - -[[autodoc]] DeepseekV3ForSequenceClassification - - forward diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 4f305546e194..a2c0bce5dc77 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -1989,7 +1989,6 @@ _import_structure["models.deepseek_v3"].extend( [ "DeepseekV3ForCausalLM", - "DeepseekV3ForSequenceClassification", "DeepseekV3Model", "DeepseekV3PreTrainedModel", ] @@ -7058,7 +7057,6 @@ ) from .models.deepseek_v3 import ( DeepseekV3ForCausalLM, - DeepseekV3ForSequenceClassification, DeepseekV3Model, DeepseekV3PreTrainedModel, ) diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index fc0aabb515e2..4a9f88258b95 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -993,7 +993,6 @@ ("data2vec-text", "Data2VecTextForSequenceClassification"), ("deberta", "DebertaForSequenceClassification"), ("deberta-v2", "DebertaV2ForSequenceClassification"), - ("deepseek_v3", "DeepseekV3ForSequenceClassification"), ("diffllama", "DiffLlamaForSequenceClassification"), ("distilbert", "DistilBertForSequenceClassification"), ("electra", "ElectraForSequenceClassification"), diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index b7abe24d0fc8..e5707050f5e7 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -16,7 +16,7 @@ from ...generation import GenerationMixin from ...modeling_attn_mask_utils import AttentionMaskConverter from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -1053,116 +1053,4 @@ def forward( ) -@add_start_docstrings( - """ - The DeepseekV3 Model transformer with a sequence classification head on top (linear layer). - - [`DeepseekV3ForSequenceClassification`] uses the last token in order to do the classification, as other causal models - (e.g. GPT-2) do. - - Since it does classification on the last token, it requires to know the position of the last token. If a - `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If - no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the - padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in - each row of the batch). - """, - DEEPSEEK_V3_START_DOCSTRING, -) -class DeepseekV3ForSequenceClassification(DeepseekV3PreTrainedModel): - def __init__(self, config): - super().__init__(config) - self.num_labels = config.num_labels - self.model = DeepseekV3Model(config) - self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) - - # Initialize weights and apply final processing - self.post_init() - - def get_input_embeddings(self): - return self.model.embed_tokens - - def set_input_embeddings(self, value): - self.model.embed_tokens = value - - @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING) - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, SequenceClassifierOutputWithPast]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If - `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - transformer_outputs = self.model( - input_ids, - 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, - ) - hidden_states = transformer_outputs[0] - logits = self.score(hidden_states) - - if input_ids is not None: - batch_size = input_ids.shape[0] - else: - batch_size = inputs_embeds.shape[0] - - if self.config.pad_token_id is None and batch_size != 1: - raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") - if self.config.pad_token_id is None: - last_non_pad_token = -1 - elif input_ids is not None: - # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id - non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) - token_indices = torch.arange(input_ids.shape[-1], device=logits.device) - last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) - else: - last_non_pad_token = -1 - logger.warning_once( - f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " - "unexpected if using padding tokens in conjunction with `inputs_embeds.`" - ) - - pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] - - loss = None - if labels is not None: - loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) - - if not return_dict: - output = (pooled_logits,) + transformer_outputs[1:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutputWithPast( - loss=loss, - logits=pooled_logits, - past_key_values=transformer_outputs.past_key_values, - hidden_states=transformer_outputs.hidden_states, - attentions=transformer_outputs.attentions, - ) - - -__all__ = [ - "DeepseekV3PreTrainedModel", - "DeepseekV3Model", - "DeepseekV3ForCausalLM", - "DeepseekV3ForSequenceClassification", -] +__all__ = ["DeepseekV3PreTrainedModel", "DeepseekV3Model", "DeepseekV3ForCausalLM"] diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index af83524adca2..a1bb96b9a021 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -14,7 +14,6 @@ from ...utils import logging from ..llama.modeling_llama import ( LlamaForCausalLM, - LlamaForSequenceClassification, LlamaModel, LlamaPreTrainedModel, LlamaRMSNorm, @@ -399,13 +398,8 @@ class DeepseekV3ForCausalLM(LlamaForCausalLM): pass -class DeepseekV3ForSequenceClassification(LlamaForSequenceClassification): - pass - - __all__ = [ "DeepseekV3PreTrainedModel", "DeepseekV3Model", "DeepseekV3ForCausalLM", - "DeepseekV3ForSequenceClassification", ] diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index 060916e85bbf..1f005fb4fe11 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -2790,13 +2790,6 @@ def __init__(self, *args, **kwargs): requires_backends(self, ["torch"]) -class DeepseekV3ForSequenceClassification(metaclass=DummyObject): - _backends = ["torch"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch"]) - - class DeepseekV3Model(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index a9e3475dfa5b..28f5c2fb9990 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -14,7 +14,6 @@ # limitations under the License. """Testing suite for the PyTorch DeepseekV3 model.""" -import gc import tempfile import unittest @@ -22,9 +21,8 @@ from packaging import version from parameterized import parameterized -from transformers import AutoTokenizer, DeepseekV3Config, StaticCache, is_torch_available, set_seed +from transformers import AutoTokenizer, DeepseekV3Config, is_torch_available, set_seed from transformers.testing_utils import ( - backend_empty_cache, require_bitsandbytes, require_flash_attn, require_read_token, @@ -47,7 +45,6 @@ from transformers import ( DeepseekV3ForCausalLM, - DeepseekV3ForSequenceClassification, DeepseekV3Model, ) from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( @@ -331,7 +328,6 @@ class DeepseekV3ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTeste ( DeepseekV3Model, DeepseekV3ForCausalLM, - DeepseekV3ForSequenceClassification, ) if is_torch_available() else () @@ -340,9 +336,7 @@ class DeepseekV3ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTeste pipeline_model_mapping = ( { "feature-extraction": DeepseekV3Model, - "text-classification": DeepseekV3ForSequenceClassification, "text-generation": DeepseekV3ForCausalLM, - "zero-shot": DeepseekV3ForSequenceClassification, } if is_torch_available() else {} @@ -375,46 +369,6 @@ def test_model_various_embeddings(self): config_and_inputs[0].position_embedding_type = type self.model_tester.create_and_check_model(*config_and_inputs) - def test_deepseek_v3_sequence_classification_model(self): - config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() - config.num_labels = 3 - input_ids = input_dict["input_ids"] - attention_mask = input_ids.ne(1).to(torch_device) - sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.type_sequence_label_size) - model = DeepseekV3ForSequenceClassification(config) - model.to(torch_device) - model.eval() - result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) - self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) - - def test_deepseek_v3_sequence_classification_model_for_single_label(self): - config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() - config.num_labels = 3 - config.problem_type = "single_label_classification" - input_ids = input_dict["input_ids"] - attention_mask = input_ids.ne(1).to(torch_device) - sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.type_sequence_label_size) - model = DeepseekV3ForSequenceClassification(config) - model.to(torch_device) - model.eval() - result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) - self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) - - def test_deepseek_v3_sequence_classification_model_for_multi_label(self): - config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() - config.num_labels = 3 - config.problem_type = "multi_label_classification" - input_ids = input_dict["input_ids"] - attention_mask = input_ids.ne(1).to(torch_device) - sequence_labels = ids_tensor( - [self.model_tester.batch_size, config.num_labels], self.model_tester.type_sequence_label_size - ).to(torch.float) - model = DeepseekV3ForSequenceClassification(config) - model.to(torch_device) - model.eval() - result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) - self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) - @unittest.skip(reason="DeepseekV3 buffers include complex numbers, which breaks this test") def test_save_load_fast_init_from_base(self): pass @@ -761,253 +715,3 @@ def test_compile_static_cache(self): ) static_compiled_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) self.assertEqual(EXPECTED_TEXT_COMPLETION, static_compiled_text) - - -@slow -@require_torch_accelerator -class Mask4DTestHard(unittest.TestCase): - def tearDown(self): - gc.collect() - backend_empty_cache(torch_device) - - def setUp(self): - model_name = "bzantium/tiny-deepseek-v3" - self.model_dtype = torch.float32 - self.tokenizer = AutoTokenizer.from_pretrained(model_name) - self.model = DeepseekV3ForCausalLM.from_pretrained(model_name, torch_dtype=self.model_dtype).to(torch_device) - - def get_test_data(self): - template = "my favorite {}" - items = ("pet is a", "artist plays a", "name is L") # same number of tokens in each item - - batch_separate = [template.format(x) for x in items] # 3 separate lines - batch_shared_prefix = template.format(" ".join(items)) # 1 line with options concatenated - - input_ids = self.tokenizer(batch_separate, return_tensors="pt").input_ids.to(torch_device) - input_ids_shared_prefix = self.tokenizer(batch_shared_prefix, return_tensors="pt").input_ids.to(torch_device) - - mask_shared_prefix = torch.tensor( - [ - [ - [ - [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], - [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0], - [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0], - [1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0], - [1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0], - [1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0], - [1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1], - ] - ] - ], - device=torch_device, - ) - - position_ids = torch.arange(input_ids.shape[1]).tile(input_ids.shape[0], 1).to(torch_device) - - # building custom positions ids based on custom mask - position_ids_shared_prefix = (mask_shared_prefix.sum(dim=-1) - 1).reshape(1, -1) - # effectively: position_ids_shared_prefix = torch.tensor([[0, 1, 2, 3, 4, 5, 3, 4, 5, 3, 4, 5]]).to(device) - - # inverting the mask - min_dtype = torch.finfo(self.model_dtype).min - mask_shared_prefix = (mask_shared_prefix.eq(0.0)).to(dtype=self.model_dtype) * min_dtype - - return input_ids, position_ids, input_ids_shared_prefix, mask_shared_prefix, position_ids_shared_prefix - - def test_stacked_causal_mask(self): - ( - input_ids, - position_ids, - input_ids_shared_prefix, - mask_shared_prefix, - position_ids_shared_prefix, - ) = self.get_test_data() - - # regular batch - logits = self.model.forward(input_ids, position_ids=position_ids).logits - logits_last = logits[:, -1, :] # last tokens in each batch line - decoded = [self.tokenizer.decode(t) for t in logits_last.argmax(dim=-1)] - - # single forward run with 4D custom mask - logits_shared_prefix = self.model.forward( - input_ids_shared_prefix, attention_mask=mask_shared_prefix, position_ids=position_ids_shared_prefix - ).logits - logits_shared_prefix_last = logits_shared_prefix[ - 0, torch.where(position_ids_shared_prefix == position_ids_shared_prefix.max())[1], : - ] # last three tokens - decoded_shared_prefix = [self.tokenizer.decode(t) for t in logits_shared_prefix_last.argmax(dim=-1)] - - self.assertEqual(decoded, decoded_shared_prefix) - - def test_partial_stacked_causal_mask(self): - # Same as the test above, but the input is passed in two groups. It tests that we can pass partial 4D attention masks - - ( - input_ids, - position_ids, - input_ids_shared_prefix, - mask_shared_prefix, - position_ids_shared_prefix, - ) = self.get_test_data() - - # regular batch - logits = self.model.forward(input_ids, position_ids=position_ids).logits - logits_last = logits[:, -1, :] # last tokens in each batch line - decoded = [self.tokenizer.decode(t) for t in logits_last.argmax(dim=-1)] - - # 2 forward runs with custom 4D masks - part_a = 3 # split point - - input_1a = input_ids_shared_prefix[:, :part_a] - position_ids_1a = position_ids_shared_prefix[:, :part_a] - mask_1a = mask_shared_prefix[:, :, :part_a, :part_a] - - outs_1a = self.model.forward(input_1a, attention_mask=mask_1a, position_ids=position_ids_1a) - past_key_values_a = outs_1a["past_key_values"] - - # Case 1: we pass a 4D attention mask regarding the current sequence length (i.e. [..., seq_len, full_len]) - input_1b = input_ids_shared_prefix[:, part_a:] - position_ids_1b = position_ids_shared_prefix[:, part_a:] - mask_1b = mask_shared_prefix[:, :, part_a:, :] - outs_1b = self.model.forward( - input_1b, - attention_mask=mask_1b, - position_ids=position_ids_1b, - past_key_values=past_key_values_a, - ) - decoded_1b = [ - self.tokenizer.decode(t) - for t in outs_1b.logits.argmax(-1)[ - 0, torch.where(position_ids_shared_prefix == position_ids_shared_prefix.max())[1] - part_a - ] - ] - self.assertEqual(decoded, decoded_1b) - - def test_stacked_causal_mask_static_cache(self): - """same as above but with StaticCache""" - ( - input_ids, - position_ids, - input_ids_shared_prefix, - mask_shared_prefix, - position_ids_shared_prefix, - ) = self.get_test_data() - - # regular batch - logits = self.model.forward(input_ids, position_ids=position_ids).logits - logits_last = logits[:, -1, :] # last tokens in each batch line - decoded = [self.tokenizer.decode(t) for t in logits_last.argmax(dim=-1)] - - # upgrade the model with StaticCache - max_cache_len = 16 # note that max_cache_len is greater than the attention_mask.shape[-1] - past_key_values = StaticCache( - config=self.model.config, - batch_size=1, - max_cache_len=max_cache_len, - device=torch_device, - dtype=self.model.dtype, - ) - - padded_attention_mask = torch.nn.functional.pad( - input=mask_shared_prefix, - pad=(0, max_cache_len - mask_shared_prefix.shape[-1]), - mode="constant", - value=torch.finfo(self.model_dtype).min, - ) - - # single forward run with 4D custom mask - logits_shared_prefix = self.model.forward( - input_ids_shared_prefix, - attention_mask=padded_attention_mask, - position_ids=position_ids_shared_prefix, - cache_position=torch.arange(input_ids_shared_prefix.shape[-1], device=torch_device), - past_key_values=past_key_values, - ).logits - logits_shared_prefix_last = logits_shared_prefix[ - 0, torch.where(position_ids_shared_prefix == position_ids_shared_prefix.max())[1], : - ] # last three tokens - decoded_shared_prefix = [self.tokenizer.decode(t) for t in logits_shared_prefix_last.argmax(dim=-1)] - - self.assertEqual(decoded, decoded_shared_prefix) - - def test_partial_stacked_causal_mask_static_cache(self): - # Same as the test above, but the input is passed in two groups. It tests that we can pass partial 4D attention masks - # we pass a 4D attention mask shaped [..., seq_len, full_static_cache_len]) - ( - input_ids, - position_ids, - input_ids_shared_prefix, - mask_shared_prefix, - position_ids_shared_prefix, - ) = self.get_test_data() - - # regular batch - logits = self.model.forward(input_ids, position_ids=position_ids).logits - logits_last = logits[:, -1, :] # last tokens in each batch line - decoded = [self.tokenizer.decode(t) for t in logits_last.argmax(dim=-1)] - - # upgrade the model with StaticCache - max_cache_len = 16 # note that max_cache_len is greater than the attention_mask.shape[-1] - past_key_values = StaticCache( - config=self.model.config, - batch_size=1, - max_cache_len=max_cache_len, - device=torch_device, - dtype=self.model.dtype, - ) - - # forward run for the first part of input - part_a = 3 # split point - - input_1a = input_ids_shared_prefix[:, :part_a] - position_ids_1a = position_ids_shared_prefix[:, :part_a] - mask_1a = mask_shared_prefix[:, :, :part_a, :part_a] - - padded_mask_1a = torch.nn.functional.pad( - input=mask_1a, - pad=(0, max_cache_len - mask_1a.shape[-1]), - mode="constant", - value=torch.finfo(self.model_dtype).min, - ) - - _ = self.model.forward( - input_1a, - attention_mask=padded_mask_1a, - position_ids=position_ids_1a, - cache_position=torch.arange(part_a, device=torch_device), - past_key_values=past_key_values, - ) - - # forward run for the second part of input - input_1b = input_ids_shared_prefix[:, part_a:] - position_ids_1b = position_ids_shared_prefix[:, part_a:] - mask_1b = mask_shared_prefix[:, :, part_a:, :] - - padded_mask_1b = torch.nn.functional.pad( - input=mask_1b, pad=(0, max_cache_len - mask_1b.shape[-1]), mode="constant", value=0 - ) - - outs_1b = self.model.forward( - input_1b, - attention_mask=padded_mask_1b, - position_ids=position_ids_1b, - cache_position=torch.arange( - part_a, - input_ids_shared_prefix.shape[-1], - device=torch_device, - ), - past_key_values=past_key_values, - ) - decoded_1b = [ - self.tokenizer.decode(t) - for t in outs_1b.logits.argmax(-1)[ - 0, torch.where(position_ids_shared_prefix == position_ids_shared_prefix.max())[1] - part_a - ] - ] - self.assertEqual(decoded, decoded_1b) From 67f1f0ca15ac8ef21090c576a0b0c4a489a2888f Mon Sep 17 00:00:00 2001 From: ryan u Date: Sat, 15 Feb 2025 18:47:00 +0900 Subject: [PATCH 42/86] cache before padding --- src/transformers/models/deepseek_v3/modeling_deepseek_v3.py | 6 +++--- src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index e5707050f5e7..78cec74a103d 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -385,14 +385,14 @@ def forward( query_states = torch.cat((q_pass, q_rot), dim=-1) key_states = torch.cat((k_pass, k_rot), dim=-1) - if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) - if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index a1bb96b9a021..e395ab07e77b 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -232,14 +232,14 @@ def forward( query_states = torch.cat((q_pass, q_rot), dim=-1) key_states = torch.cat((k_pass, k_rot), dim=-1) - if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) - if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): From f264f800d04950390db8413b9efb24cef8186330 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 18 Feb 2025 19:48:25 +0900 Subject: [PATCH 43/86] remote output_router_logits --- .../models/deepseek_v3/modeling_deepseek_v3.py | 15 +++------------ .../models/deepseek_v3/modular_deepseek_v3.py | 15 +++------------ 2 files changed, 6 insertions(+), 24 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 78cec74a103d..e443954b79af 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -157,7 +157,7 @@ def forward(self, hidden_states): denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 topk_weights /= denominator topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights, router_logits + return topk_indices, topk_weights @torch.no_grad() def get_topk_indices(self, scores): @@ -202,11 +202,11 @@ def __init__(self, config): def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights, router_logits = self.gate(hidden_states) + topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states, router_logits + return hidden_states def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) @@ -444,7 +444,6 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, - output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC @@ -472,19 +471,11 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) - - if isinstance(hidden_states, tuple): - hidden_states, router_logits = hidden_states - else: - router_logits = (torch.zeros((1,), device=hidden_states.device, dtype=torch.int64),) - hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) - if output_router_logits: - outputs += (router_logits,) return outputs diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index e395ab07e77b..cc8895665515 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -82,7 +82,7 @@ def forward(self, hidden_states): denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 topk_weights /= denominator topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights, router_logits + return topk_indices, topk_weights @torch.no_grad() def get_topk_indices(self, scores): @@ -127,11 +127,11 @@ def __init__(self, config): def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights, router_logits = self.gate(hidden_states) + topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states, router_logits + return hidden_states def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) @@ -291,7 +291,6 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, - output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC @@ -319,19 +318,11 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) - - if isinstance(hidden_states, tuple): - hidden_states, router_logits = hidden_states - else: - router_logits = (torch.zeros((1,), device=hidden_states.device, dtype=torch.int64),) - hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) - if output_router_logits: - outputs += (router_logits,) return outputs From d4c6a1bd1475e0c53b1242a64b689cc311860bd7 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 18 Feb 2025 19:49:12 +0900 Subject: [PATCH 44/86] Revert "remote output_router_logits" This reverts commit f264f800d04950390db8413b9efb24cef8186330. --- .../models/deepseek_v3/modeling_deepseek_v3.py | 15 ++++++++++++--- .../models/deepseek_v3/modular_deepseek_v3.py | 15 ++++++++++++--- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index e443954b79af..78cec74a103d 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -157,7 +157,7 @@ def forward(self, hidden_states): denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 topk_weights /= denominator topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights + return topk_indices, topk_weights, router_logits @torch.no_grad() def get_topk_indices(self, scores): @@ -202,11 +202,11 @@ def __init__(self, config): def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) + topk_indices, topk_weights, router_logits = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states + return hidden_states, router_logits def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) @@ -444,6 +444,7 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC @@ -471,11 +472,19 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) + + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = (torch.zeros((1,), device=hidden_states.device, dtype=torch.int64),) + hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) + if output_router_logits: + outputs += (router_logits,) return outputs diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index cc8895665515..e395ab07e77b 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -82,7 +82,7 @@ def forward(self, hidden_states): denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 topk_weights /= denominator topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights + return topk_indices, topk_weights, router_logits @torch.no_grad() def get_topk_indices(self, scores): @@ -127,11 +127,11 @@ def __init__(self, config): def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) + topk_indices, topk_weights, router_logits = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states + return hidden_states, router_logits def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) @@ -291,6 +291,7 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC @@ -318,11 +319,19 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) + + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = (torch.zeros((1,), device=hidden_states.device, dtype=torch.int64),) + hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) + if output_router_logits: + outputs += (router_logits,) return outputs From c7c8d766b124bae0ee76d17e963a3c62c481f9ac Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 18 Feb 2025 19:50:56 +0900 Subject: [PATCH 45/86] remove output_router_logits --- .../models/deepseek_v3/modeling_deepseek_v3.py | 15 +++------------ .../models/deepseek_v3/modular_deepseek_v3.py | 15 +++------------ 2 files changed, 6 insertions(+), 24 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 78cec74a103d..e443954b79af 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -157,7 +157,7 @@ def forward(self, hidden_states): denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 topk_weights /= denominator topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights, router_logits + return topk_indices, topk_weights @torch.no_grad() def get_topk_indices(self, scores): @@ -202,11 +202,11 @@ def __init__(self, config): def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights, router_logits = self.gate(hidden_states) + topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states, router_logits + return hidden_states def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) @@ -444,7 +444,6 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, - output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC @@ -472,19 +471,11 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) - - if isinstance(hidden_states, tuple): - hidden_states, router_logits = hidden_states - else: - router_logits = (torch.zeros((1,), device=hidden_states.device, dtype=torch.int64),) - hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) - if output_router_logits: - outputs += (router_logits,) return outputs diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index e395ab07e77b..cc8895665515 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -82,7 +82,7 @@ def forward(self, hidden_states): denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 topk_weights /= denominator topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights, router_logits + return topk_indices, topk_weights @torch.no_grad() def get_topk_indices(self, scores): @@ -127,11 +127,11 @@ def __init__(self, config): def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights, router_logits = self.gate(hidden_states) + topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states, router_logits + return hidden_states def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) @@ -291,7 +291,6 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, - output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC @@ -319,19 +318,11 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) - - if isinstance(hidden_states, tuple): - hidden_states, router_logits = hidden_states - else: - router_logits = (torch.zeros((1,), device=hidden_states.device, dtype=torch.int64),) - hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) - if output_router_logits: - outputs += (router_logits,) return outputs From ba6f7d404e48bd0cd9f8e86426f91841dcac1ccb Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 18 Feb 2025 21:16:11 +0900 Subject: [PATCH 46/86] make e_score_correction_bias as buffer --- src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index cc8895665515..fd40f4676252 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -70,7 +70,7 @@ def __init__(self, config): self.norm_topk_prob = config.norm_topk_prob self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) - self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) + self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts))) def forward(self, hidden_states): hidden_states = hidden_states.view(-1, self.config.hidden_size) From d7931b32fff6ebdeabed1a32fefe3fb2a371a0e9 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 18 Feb 2025 21:17:28 +0900 Subject: [PATCH 47/86] skip tests not compatible --- .../deepseek_v3/test_modeling_deepseek_v3.py | 187 ++++++------------ 1 file changed, 61 insertions(+), 126 deletions(-) diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index 28f5c2fb9990..3f0582cc011c 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -356,6 +356,63 @@ def setUp(self): self.model_tester = DeepseekV3ModelTester(self) self.config_tester = ConfigTester(self, config_class=DeepseekV3Config, hidden_size=37) + @unittest.skip("Failing because of unique cache (HybridCache)") + def test_model_outputs_equivalence(self, **kwargs): + pass + + @parameterized.expand([("random",), ("same",)]) + @unittest.skip("DeepseekV3 has HybridCache which is not compatible with assisted decoding") + def test_assisted_decoding_matches_greedy_search(self, assistant_type): + pass + + @unittest.skip("DeepseekV3 has HybridCache which is not compatible with assisted decoding") + def test_prompt_lookup_decoding_matches_greedy_search(self, assistant_type): + pass + + @unittest.skip("DeepseekV3 has HybridCache which is not compatible with assisted decoding") + def test_assisted_decoding_sample(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache which is not compatible with dola decoding") + def test_dola_decoding_sample(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache and doesn't support continue from past kv") + def test_generate_continue_from_past_key_values(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache and doesn't support low_memory generation") + def test_beam_search_low_memory(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache and doesn't support contrastive generation") + def test_contrastive_generate(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache and doesn't support contrastive generation") + def test_contrastive_generate_dict_outputs_use_cache(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache and doesn't support contrastive generation") + def test_contrastive_generate_low_memory(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support.") + def test_generate_with_static_cache(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support.") + def test_generate_from_inputs_embeds_with_static_cache(self): + pass + + @unittest.skip("DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support.") + def test_generate_continue_from_inputs_embeds(self): + pass + + @unittest.skip("DeepseekV3's eager attn/sdpa attn outputs are expected to be different") + def test_sdpa_equivalence(self): + pass + def test_config(self): self.config_tester.run_common_tests() @@ -369,11 +426,7 @@ def test_model_various_embeddings(self): config_and_inputs[0].position_embedding_type = type self.model_tester.create_and_check_model(*config_and_inputs) - @unittest.skip(reason="DeepseekV3 buffers include complex numbers, which breaks this test") - def test_save_load_fast_init_from_base(self): - pass - - @parameterized.expand([("linear",), ("dynamic",), ("yarn",)]) + @parameterized.expand([("yarn",)]) def test_model_rope_scaling_from_config(self, scaling_type): config, _ = self.model_tester.prepare_config_and_inputs_for_common() short_input = ids_tensor([1, 10], config.vocab_size) @@ -469,127 +522,9 @@ def test_model_rope_scaling(self): with self.assertRaises(AssertionError): torch.testing.assert_close(yarn_sin_long, original_sin_long) - def test_model_loading_old_rope_configs(self): - def _reinitialize_config(base_config, new_kwargs): - # Reinitialize the config with the new kwargs, forcing the config to go through its __init__ validation - # steps. - base_config_dict = base_config.to_dict() - new_config = DeepseekV3Config.from_dict(config_dict={**base_config_dict, **new_kwargs}) - return new_config - - # from untouched config -> βœ… - base_config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common() - original_model = DeepseekV3ForCausalLM(base_config).to(torch_device) - original_model(**model_inputs) - - # from a config with the expected rope configuration -> βœ… - config = _reinitialize_config(base_config, {"rope_scaling": {"rope_type": "linear", "factor": 10.0}}) - original_model = DeepseekV3ForCausalLM(config).to(torch_device) - original_model(**model_inputs) - - # from a config with the old rope configuration ('type' instead of 'rope_type') -> βœ… we gracefully handle BC - config = _reinitialize_config(base_config, {"rope_scaling": {"type": "linear", "factor": 10.0}}) - original_model = DeepseekV3ForCausalLM(config).to(torch_device) - original_model(**model_inputs) - - # from a config with both 'type' and 'rope_type' -> βœ… they can coexist (and both are present in the config) - config = _reinitialize_config( - base_config, {"rope_scaling": {"type": "linear", "rope_type": "linear", "factor": 10.0}} - ) - self.assertTrue(config.rope_scaling["type"] == "linear") - self.assertTrue(config.rope_scaling["rope_type"] == "linear") - original_model = DeepseekV3ForCausalLM(config).to(torch_device) - original_model(**model_inputs) - - # from a config with parameters in a bad range ('factor' should be >= 1.0) -> ⚠️ throws a warning - with self.assertLogs("transformers.modeling_rope_utils", level="WARNING") as logs: - config = _reinitialize_config(base_config, {"rope_scaling": {"rope_type": "linear", "factor": -999.0}}) - original_model = DeepseekV3ForCausalLM(config).to(torch_device) - original_model(**model_inputs) - self.assertEqual(len(logs.output), 1) - self.assertIn("factor field", logs.output[0]) - - # from a config with unknown parameters ('foo' isn't a rope option) -> ⚠️ throws a warning - with self.assertLogs("transformers.modeling_rope_utils", level="WARNING") as logs: - config = _reinitialize_config( - base_config, {"rope_scaling": {"rope_type": "linear", "factor": 10.0, "foo": "bar"}} - ) - original_model = DeepseekV3ForCausalLM(config).to(torch_device) - original_model(**model_inputs) - self.assertEqual(len(logs.output), 1) - self.assertIn("Unrecognized keys", logs.output[0]) - - # from a config with specific rope type but missing one of its mandatory parameters -> ❌ throws exception - with self.assertRaises(KeyError): - config = _reinitialize_config(base_config, {"rope_scaling": {"rope_type": "linear"}}) # missing "factor" - - @require_flash_attn - @require_torch_gpu - @require_bitsandbytes - @pytest.mark.flash_attn_test - @require_read_token - @slow - def test_flash_attn_2_generate_padding_right(self): - """ - Overwritting the common test as the test is flaky on tiny models - """ - model = DeepseekV3ForCausalLM.from_pretrained( - "bzantium/tiny-deepseek-v3", - load_in_4bit=True, - device_map={"": 0}, - ) - - tokenizer = AutoTokenizer.from_pretrained("bzantium/tiny-deepseek-v3") - - texts = ["hi", "Hello this is a very long sentence"] - - tokenizer.padding_side = "right" - tokenizer.pad_token = tokenizer.eos_token - - inputs = tokenizer(texts, return_tensors="pt", padding=True).to(0) - - output_native = model.generate(**inputs, max_new_tokens=20, do_sample=False) - output_native = tokenizer.batch_decode(output_native) - - model = DeepseekV3ForCausalLM.from_pretrained( - "bzantium/tiny-deepseek-v3", - load_in_4bit=True, - device_map={"": 0}, - attn_implementation="flash_attention_2", - ) - - output_fa_2 = model.generate(**inputs, max_new_tokens=20, do_sample=False) - output_fa_2 = tokenizer.batch_decode(output_fa_2) - - self.assertListEqual(output_native, output_fa_2) - - @require_flash_attn - @require_torch_gpu - @slow - @pytest.mark.flash_attn_test - def test_use_flash_attention_2_true(self): - """ - NOTE: this is the only test testing that the legacy `use_flash_attention=2` argument still works as intended. - """ - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - for model_class in self.all_model_classes: - with tempfile.TemporaryDirectory() as tmp_dir: - model = model_class(config) - model.save_pretrained(tmp_dir) - - new_model = DeepseekV3ForCausalLM.from_pretrained( - tmp_dir, use_flash_attention_2=True, torch_dtype=torch.float16 - ).to("cuda") - - self.assertTrue(new_model.config._attn_implementation == "flash_attention_2") - - has_flash = False - for name, submodule in new_model.named_modules(): - if "FlashAttention" in submodule.__class__.__name__: - has_flash = True - break - if not has_flash: - raise ValueError("The flash model should have flash attention layers") + @unittest.skip(reason="Deepseek-V3 uses MLA on all models so the KV cache is a non standard format") + def test_past_key_values_format(self): + pass @require_torch_sdpa @slow From 92bd99cb59701659443059f3b24263347c6238d9 Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 18 Feb 2025 21:21:05 +0900 Subject: [PATCH 48/86] make style --- .../deepseek_v3/test_modeling_deepseek_v3.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index 3f0582cc011c..03017630eecc 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -14,21 +14,16 @@ # limitations under the License. """Testing suite for the PyTorch DeepseekV3 model.""" -import tempfile import unittest -import pytest from packaging import version from parameterized import parameterized from transformers import AutoTokenizer, DeepseekV3Config, is_torch_available, set_seed from transformers.testing_utils import ( - require_bitsandbytes, - require_flash_attn, require_read_token, require_torch, require_torch_accelerator, - require_torch_gpu, require_torch_sdpa, slow, torch_device, @@ -397,15 +392,21 @@ def test_contrastive_generate_dict_outputs_use_cache(self): def test_contrastive_generate_low_memory(self): pass - @unittest.skip("DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support.") + @unittest.skip( + "DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support." + ) def test_generate_with_static_cache(self): pass - @unittest.skip("DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support.") + @unittest.skip( + "DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support." + ) def test_generate_from_inputs_embeds_with_static_cache(self): pass - @unittest.skip("DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support.") + @unittest.skip( + "DeepseekV3 has HybridCache and doesn't support StaticCache. Though it could, it shouldn't support." + ) def test_generate_continue_from_inputs_embeds(self): pass From 7d81efea7feda6b50af637220bc3146e87286b2e Mon Sep 17 00:00:00 2001 From: ryan u Date: Tue, 18 Feb 2025 21:26:13 +0900 Subject: [PATCH 49/86] make e_score_correction_bias as buffer --- src/transformers/models/deepseek_v3/modeling_deepseek_v3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index e443954b79af..2e52d1104656 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -145,7 +145,7 @@ def __init__(self, config): self.norm_topk_prob = config.norm_topk_prob self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) - self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts))) + self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts))) def forward(self, hidden_states): hidden_states = hidden_states.view(-1, self.config.hidden_size) From b33fdb5b1b73528c32397f874fcccb4f6dcc836e Mon Sep 17 00:00:00 2001 From: ryan u Date: Wed, 19 Feb 2025 11:30:16 +0900 Subject: [PATCH 50/86] use rope_interleave instead of load_hook --- .../deepseek_v3/configuration_deepseek_v3.py | 2 + .../deepseek_v3/modeling_deepseek_v3.py | 89 +++++------ .../models/deepseek_v3/modular_deepseek_v3.py | 143 +++++++++++------- 3 files changed, 132 insertions(+), 102 deletions(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 24c2fb66cbbb..1c73c078cd27 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -179,6 +179,7 @@ def __init__( tie_word_embeddings=False, rope_theta=10000.0, rope_scaling=None, + rope_interleave=True, attention_bias=False, attention_dropout=0.0, **kwargs, @@ -205,6 +206,7 @@ def __init__( self.num_experts_per_tok = num_experts_per_tok self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob + self.rope_interleave = rope_interleave # for backward compatibility if num_key_value_heads is None: diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 2e52d1104656..fb20b75cc8d2 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -299,6 +299,40 @@ def eager_attention_forward( return attn_output, attn_weights +def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, 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.unsqueeze(unsqueeze_dim) + sin = sin.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 + + def yarn_get_mscale(scale=1, mscale=1): if scale <= 1: return 1.0 @@ -379,7 +413,10 @@ def forward( k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings - q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) + if self.config.rope_interleave: + q_rot, k_rot = apply_rotary_pos_emb_interleave(q_pass, k_pass, cos, sin) + else: + q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) query_states = torch.cat((q_pass, q_rot), dim=-1) @@ -614,7 +651,7 @@ class DeepseekV3Model(DeepseekV3PreTrainedModel): config: DeepseekV3Config """ - def __init__(self, config): + def __init__(self, config: DeepseekV3Config): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size @@ -626,8 +663,6 @@ def __init__(self, config): self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = DeepseekV3RotaryEmbedding(config=config) self.gradient_checkpointing = False - self._register_load_state_dict_pre_hook(self.load_pre_hook) - self._register_state_dict_hook(self.load_hook) # Initialize weights and apply final processing self.post_init() @@ -868,52 +903,6 @@ def _prepare_4d_causal_attention_mask_with_cache_position( return causal_mask - def load_pre_hook(self, state_dict, prefix, *args): - """ - Weights have to be permuted for correct rope formulation. We can't do this in the weights - as every other framework already uses the `Llama` original function (which is copyrighted btw). - And I am not even sure it's better.... anyways end of my rant - """ - - def permute_for_rope(input_tensor): - """ - When you go from the complex ROPE formulation to sin and cos one, you need - to permute the query and key weights (to avoid doing it on the fly) - """ - n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] - input_tensor = input_tensor.reshape(n_heads * dim1, dim2) - input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) - input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) - return input_tensor - - def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): - weight = state_dict[key] - weight = weight.view(num_heads, head_dim, -1) - weight_rot = weight[:, -rope_dim:] - weight_rot = permute_for_rope(weight_rot) - weight[:, -rope_dim:] = weight_rot - weight = weight.view(-1, weight.shape[-1]) - state_dict[key] = weight - - for k in state_dict: - if "q_b_proj." in k: - permute_layer_for_rope( - k, - num_heads=self.config.num_attention_heads, - head_dim=self.config.qk_head_dim, - rope_dim=self.config.qk_rope_head_dim, - ) - if "kv_a_proj_with_mqa." in k: - permute_layer_for_rope( - k, - num_heads=1, - head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, - rope_dim=self.config.qk_rope_head_dim, - ) - - def load_hook(self, module, state_dict, prefix, *args): - self.load_pre_hook(state_dict, prefix, *args) - class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index fd40f4676252..cc8f7e00908a 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -20,6 +20,7 @@ LlamaRotaryEmbedding, apply_rotary_pos_emb, eager_attention_forward, + rotate_half, ) from .configuration_deepseek_v3 import DeepseekV3Config @@ -35,6 +36,40 @@ class DeepseekV3RotaryEmbedding(LlamaRotaryEmbedding): pass +def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, 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.unsqueeze(unsqueeze_dim) + sin = sin.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 + + def yarn_get_mscale(scale=1, mscale=1): if scale <= 1: return 1.0 @@ -226,7 +261,10 @@ def forward( k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings - q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) + if self.config.rope_interleave: + q_rot, k_rot = apply_rotary_pos_emb_interleave(q_pass, k_pass, cos, sin) + else: + q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) query_states = torch.cat((q_pass, q_rot), dim=-1) @@ -332,57 +370,58 @@ class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): class DeepseekV3Model(LlamaModel): - def __init__(self, config): - super().__init__(config) - self._register_load_state_dict_pre_hook(self.load_pre_hook) - self._register_state_dict_hook(self.load_hook) - self.post_init() - - def load_pre_hook(self, state_dict, prefix, *args): - """ - Weights have to be permuted for correct rope formulation. We can't do this in the weights - as every other framework already uses the `Llama` original function (which is copyrighted btw). - And I am not even sure it's better.... anyways end of my rant - """ - - def permute_for_rope(input_tensor): - """ - When you go from the complex ROPE formulation to sin and cos one, you need - to permute the query and key weights (to avoid doing it on the fly) - """ - n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] - input_tensor = input_tensor.reshape(n_heads * dim1, dim2) - input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) - input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) - return input_tensor - - def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): - weight = state_dict[key] - weight = weight.view(num_heads, head_dim, -1) - weight_rot = weight[:, -rope_dim:] - weight_rot = permute_for_rope(weight_rot) - weight[:, -rope_dim:] = weight_rot - weight = weight.view(-1, weight.shape[-1]) - state_dict[key] = weight - - for k in state_dict: - if "q_b_proj." in k: - permute_layer_for_rope( - k, - num_heads=self.config.num_attention_heads, - head_dim=self.config.qk_head_dim, - rope_dim=self.config.qk_rope_head_dim, - ) - if "kv_a_proj_with_mqa." in k: - permute_layer_for_rope( - k, - num_heads=1, - head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, - rope_dim=self.config.qk_rope_head_dim, - ) - - def load_hook(self, module, state_dict, prefix, *args): - self.load_pre_hook(state_dict, prefix, *args) + pass + # def __init__(self, config): + # super().__init__(config) + # self._register_load_state_dict_pre_hook(self.load_pre_hook) + # self._register_state_dict_hook(self.load_hook) + # self.post_init() + + # def load_pre_hook(self, state_dict, prefix, *args): + # """ + # Weights have to be permuted for correct rope formulation. We can't do this in the weights + # as every other framework already uses the `Llama` original function (which is copyrighted btw). + # And I am not even sure it's better.... anyways end of my rant + # """ + + # def permute_for_rope(input_tensor): + # """ + # When you go from the complex ROPE formulation to sin and cos one, you need + # to permute the query and key weights (to avoid doing it on the fly) + # """ + # n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] + # input_tensor = input_tensor.reshape(n_heads * dim1, dim2) + # input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) + # input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) + # return input_tensor + + # def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): + # weight = state_dict[key] + # weight = weight.view(num_heads, head_dim, -1) + # weight_rot = weight[:, -rope_dim:] + # weight_rot = permute_for_rope(weight_rot) + # weight[:, -rope_dim:] = weight_rot + # weight = weight.view(-1, weight.shape[-1]) + # state_dict[key] = weight + + # for k in state_dict: + # if "q_b_proj." in k: + # permute_layer_for_rope( + # k, + # num_heads=self.config.num_attention_heads, + # head_dim=self.config.qk_head_dim, + # rope_dim=self.config.qk_rope_head_dim, + # ) + # if "kv_a_proj_with_mqa." in k: + # permute_layer_for_rope( + # k, + # num_heads=1, + # head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, + # rope_dim=self.config.qk_rope_head_dim, + # ) + + # def load_hook(self, module, state_dict, prefix, *args): + # self.load_pre_hook(state_dict, prefix, *args) class DeepseekV3ForCausalLM(LlamaForCausalLM): From 7f859f8de4b832b091ff46cda3ccc4bfdac3a24d Mon Sep 17 00:00:00 2001 From: ryan u Date: Wed, 19 Feb 2025 11:30:35 +0900 Subject: [PATCH 51/86] skip tests not compatible with MLA --- .../deepseek_v3/test_modeling_deepseek_v3.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index 03017630eecc..c5efee3f9a6b 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -414,6 +414,22 @@ def test_generate_continue_from_inputs_embeds(self): def test_sdpa_equivalence(self): pass + @unittest.skip("Deepseek-V3 uses MLA so it is not compatible with the standard cache format") + def test_beam_search_generate_dict_outputs_use_cache(self): + pass + + @unittest.skip("Deepseek-V3 uses MLA so it is not compatible with the standard cache format") + def test_generate_compilation_all_outputs(self): + pass + + @unittest.skip("Deepseek-V3 uses MLA so it is not compatible with the standard cache format") + def test_generate_compile_model_forward(self): + pass + + @unittest.skip("Deepseek-V3 uses MLA so it is not compatible with the standard cache format") + def test_greedy_generate_dict_outputs_use_cache(self): + pass + def test_config(self): self.config_tester.run_common_tests() From 397ecf3c636492727b52e44baf106f587b1317a0 Mon Sep 17 00:00:00 2001 From: ryan u Date: Wed, 19 Feb 2025 11:32:41 +0900 Subject: [PATCH 52/86] add doc for rope_interleave --- .../models/deepseek_v3/configuration_deepseek_v3.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 1c73c078cd27..ccfc97832c1b 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -113,6 +113,8 @@ class DeepseekV3Config(PretrainedConfig): strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update `max_position_embeddings` to the expected new maximum. + rope_interleave (`bool`, *optional*, defaults to `True`): + Whether to interleave the rotary position embeddings. attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): Whether to use a bias in the query, key, value and output projection layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): From 2628438a435664ce0a01367745c0e2b412550115 Mon Sep 17 00:00:00 2001 From: ryan u Date: Wed, 19 Feb 2025 13:48:19 +0900 Subject: [PATCH 53/86] fix typo --- src/transformers/models/deepseek_v3/modeling_deepseek_v3.py | 2 +- src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index fb20b75cc8d2..84850462ce48 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -414,7 +414,7 @@ def forward( cos, sin = position_embeddings if self.config.rope_interleave: - q_rot, k_rot = apply_rotary_pos_emb_interleave(q_pass, k_pass, cos, sin) + q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) else: q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index cc8f7e00908a..8bd002d31dda 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -262,7 +262,7 @@ def forward( cos, sin = position_embeddings if self.config.rope_interleave: - q_rot, k_rot = apply_rotary_pos_emb_interleave(q_pass, k_pass, cos, sin) + q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) else: q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) From af3d32855e168b3f4c5a1a28788455f9fb4d42be Mon Sep 17 00:00:00 2001 From: ryan u Date: Wed, 19 Feb 2025 13:59:21 +0900 Subject: [PATCH 54/86] remove torch.no_grad for selecting topk --- src/transformers/models/deepseek_v3/modeling_deepseek_v3.py | 1 - src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 84850462ce48..872eb89e77c1 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -159,7 +159,6 @@ def forward(self, hidden_states): topk_weights = topk_weights * self.routed_scaling_factor return topk_indices, topk_weights - @torch.no_grad() def get_topk_indices(self, scores): scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 8bd002d31dda..fe5ea154d8d0 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -119,7 +119,6 @@ def forward(self, hidden_states): topk_weights = topk_weights * self.routed_scaling_factor return topk_indices, topk_weights - @torch.no_grad() def get_topk_indices(self, scores): scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( From 14e7d4e933df911354d3d826d16927d7f797427a Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 24 Mar 2025 14:22:32 +0000 Subject: [PATCH 55/86] fix post merge issue --- src/transformers/modeling_rope_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_rope_utils.py b/src/transformers/modeling_rope_utils.py index 6c66c669a0a6..7f7fbb3df469 100644 --- a/src/transformers/modeling_rope_utils.py +++ b/src/transformers/modeling_rope_utils.py @@ -449,7 +449,7 @@ def _validate_yarn_parameters(config: PretrainedConfig, ignore_keys: Optional[se "original_max_position_embeddings", "mscale", "mscale_all_dim", - , "original_max_position_embeddings"} + } received_keys = set(rope_scaling.keys()) _check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys) From 1d8516d5ae0769e0ad42c477e9626c26d9be2dac Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 24 Mar 2025 17:42:19 +0100 Subject: [PATCH 56/86] mrege with main and simplify --- src/transformers/modeling_utils.py | 33 +++++++++---------- .../models/llama/modeling_llama.py | 1 - tests/utils/test_modeling_utils.py | 12 ++++++- 3 files changed, 27 insertions(+), 19 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 7f1dd76c5800..c7b7ab61e402 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -28,7 +28,6 @@ import tempfile import warnings from collections import defaultdict -from collections import defaultdict from contextlib import contextmanager from dataclasses import dataclass from enum import Enum @@ -890,25 +889,25 @@ def update_key_name(keys): """ Updates a dictionary of keys to pack layers together as layer.{0, 1, 4} instead of layers.0, layers.1, layers.4. """ - key_dict = defaultdict(set) + key_dict = defaultdict(list) for key in keys: - modified_key = r"" - for char in key: - if char.isdigit() and modified_key[-1] == ".": - if modified_key != "": - modified_key += r"(\d+)" - modified_key = modified_key.replace(".", r"\.") - key_dict[modified_key].add(int(char)) - modified_key = r"" - else: - modified_key += char + all_digits = re.findall(r".(\d+).", key) + for i, k in enumerate(all_digits): + if len(key_dict[re.sub(r"(\d+)", "*", key)]) <= i: + key_dict[re.sub(r"(\d+)", "*", key)].append(set()) + key_dict[re.sub(r"(\d+)", "*", key)][i].add(int(k)) + final_keys = set() for key in keys: - text = key - for pattern, values in key_dict.items(): - if len(values) > 1: - text = re.sub(pattern, lambda match: match.group(0)[: -len(match.group(1))] + str(values), text) - final_keys.add(text) + text = re.sub(r"(\d+)", "*", key) + pattern = key_dict[text] + final_text = "" + for i, part in enumerate(text.split("*")): + if len(pattern) <= i: + final_text += part + else: + final_text += part + "{" + ",".join([str(i) for i in sorted(list(pattern[i]))]) + "}" + final_keys.add(final_text) return final_keys diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 2d2cf3eecb8d..eb37b6052913 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -80,7 +80,6 @@ def forward(self, hidden_states): hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) - def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 8c03d676dbdf..36ba44fd2cd7 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -1737,11 +1737,21 @@ def test_update_key_name(self): "model.language_model.1.self_attn.0.mlp.layer.1.weight", "model.language_model.2.self_attn.0.mlp.layer.1.weight", "model.language_model.3.self_attn.0.mlp.layer.1.weight", + "model.language_model.4.self_attn.0.mlp.layer.1.weight", + "model.language_model.5.self_attn.0.mlp.layer.1.weight", + "model.language_model.6.self_attn.0.mlp.layer.1.weight", + "model.language_model.7.self_attn.0.mlp.layer.1.weight", + "model.language_model.8.self_attn.0.mlp.layer.1.weight", + "model.language_model.9.self_attn.0.mlp.layer.1.weight", + "model.language_model.10.self_attn.0.mlp.layer.1.weight", + "model.language_model.11.self_attn.0.mlp.layer.1.weight", + "model.language_model.12.self_attn.0.mlp.layer.1.weight", + "model.language_model.13.self_attn.0.mlp.layer.1.weigh,t", ] new_keys = update_key_name(original_keys) expected_new_keys = [ + "model.language_model.{0, ..., 13}.self_attn.{0, 1}.mlp.layer.1.weight", "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.conv1.weight", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.1.weight", "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp0", "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.conv1.weight", "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.conv1.weight", From 9b4f43330d386358ebab941c271b5248f3d4e250 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 24 Mar 2025 17:48:51 +0100 Subject: [PATCH 57/86] nits --- src/transformers/modeling_utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index c7b7ab61e402..9712fec0b212 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -906,7 +906,12 @@ def update_key_name(keys): if len(pattern) <= i: final_text += part else: - final_text += part + "{" + ",".join([str(i) for i in sorted(list(pattern[i]))]) + "}" + data = [str(i) for i in sorted(pattern[i])] + result = ", ".join(data) # If there are only 1 or 2 elements, show them all + if len(data) > 1: + final_text += part + "{" + result + "}" + else: + final_text += part + data[0] final_keys.add(final_text) return final_keys From abffdfeb8d01d74edf322e4693295350a4383291 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 24 Mar 2025 17:58:57 +0100 Subject: [PATCH 58/86] final --- src/transformers/modeling_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 9712fec0b212..173743dbb09f 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -907,7 +907,10 @@ def update_key_name(keys): final_text += part else: data = [str(i) for i in sorted(pattern[i])] - result = ", ".join(data) # If there are only 1 or 2 elements, show them all + if len(data) > 10: + result = f"{data[0]}...{data[-1]}" + else: + result = ", ".join(data) # If there are only 1 or 2 elements, show them all if len(data) > 1: final_text += part + "{" + result + "}" else: From 6c7eaa5197d87361d7875d5908aab60e23dd259a Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 24 Mar 2025 18:15:17 +0100 Subject: [PATCH 59/86] small fixes --- src/transformers/modeling_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 173743dbb09f..ce05bc853183 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -893,13 +893,13 @@ def update_key_name(keys): for key in keys: all_digits = re.findall(r".(\d+).", key) for i, k in enumerate(all_digits): - if len(key_dict[re.sub(r"(\d+)", "*", key)]) <= i: - key_dict[re.sub(r"(\d+)", "*", key)].append(set()) - key_dict[re.sub(r"(\d+)", "*", key)][i].add(int(k)) + if len(key_dict[re.sub(r".(\d+).", ".*.", key)]) <= i: + key_dict[re.sub(r".(\d+).", ".*.", key)].append(set()) + key_dict[re.sub(r".(\d+).", ".*.", key)][i].add(int(k)) final_keys = set() for key in keys: - text = re.sub(r"(\d+)", "*", key) + text = re.sub(r".(\d+).", ".*.", key) pattern = key_dict[text] final_text = "" for i, part in enumerate(text.split("*")): @@ -916,7 +916,7 @@ def update_key_name(keys): else: final_text += part + data[0] final_keys.add(final_text) - return final_keys + return list(final_keys) def _get_resolved_checkpoint_files( From 9e4965ad4a1a98b21cbffd1fd3eaeb2712efc921 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 24 Mar 2025 18:19:20 +0100 Subject: [PATCH 60/86] fix --- src/transformers/modeling_utils.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index ce05bc853183..5b9bfc937d96 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -5023,16 +5023,12 @@ def _load_pretrained_model( "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method." ) raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}") - - missing_keys = update_key_name(missing_keys) - unexpected_keys = update_key_name(unexpected_keys) - if len(unexpected_keys) > 0: archs = [] if model.config.architectures is None else model.config.architectures warner = logger.warning if model.__class__.__name__ in archs else logger.info warner( f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" - f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" + f" initializing {model.__class__.__name__}: {update_key_name(unexpected_keys)}\n- This IS expected if you are" f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" " with another architecture (e.g. initializing a BertForSequenceClassification model from a" " BertForPreTraining model).\n- This IS NOT expected if you are initializing" @@ -5044,7 +5040,7 @@ def _load_pretrained_model( if len(missing_keys) > 0: logger.warning( f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" + f" {pretrained_model_name_or_path} and are newly initialized: {update_key_name(missing_keys)}\nYou should probably" " TRAIN this model on a down-stream task to be able to use it for predictions and inference." ) elif len(mismatched_keys) == 0: From 6bb8802d1ab0440f36200c538bb33d139c58840b Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Tue, 25 Mar 2025 08:22:04 +0000 Subject: [PATCH 61/86] support TP better --- src/transformers/integrations/fbgemm_fp8.py | 2 +- src/transformers/integrations/tensor_parallel.py | 5 ++++- src/transformers/modeling_utils.py | 12 +++++++----- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/src/transformers/integrations/fbgemm_fp8.py b/src/transformers/integrations/fbgemm_fp8.py index 71c2b570cc0a..98bc919caa15 100644 --- a/src/transformers/integrations/fbgemm_fp8.py +++ b/src/transformers/integrations/fbgemm_fp8.py @@ -28,7 +28,7 @@ logger = logging.get_logger(__name__) -class FbgemmFp8Linear(torch.nn.Module): +class FbgemmFp8Linear(torch.nn.Linear): def __init__(self, in_features, out_features, bias, weight_dtype=torch.float32): super().__init__() self.in_features = in_features diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index b3fab0fe19c3..7df881b9bac7 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -484,7 +484,10 @@ def __init__(self): # 1. We add hooks to the layer being loaded: if current_module_plan is not None: tp_layer = translate_to_torch_parallel_style(current_module_plan) - tp_layer.prepare_module_tp(module, device_mesh) + try: + tp_layer.prepare_module_tp(module, device_mesh) + except NotImplementedError as e: + print(f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP plan") # 2. We add hooks to the parrent module if needed if "." in layer_name: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 7747003aeb4a..b81183d4da2d 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -807,7 +807,7 @@ def _load_state_dict_into_meta_model( param_name, casting_dtype, to_contiguous, - tensor_device, # the rank + int(os.environ["RANK"]), # the rank device_mesh, ) else: @@ -4096,7 +4096,7 @@ def from_pretrained( if not torch.distributed.is_initialized(): try: logger.warning("Tensor Parallel requires torch.distributed to be initialized first.") - rank = int(os.environ["RANK"]) + rank = int(os.environ["LOCAL_RANK"]) world_size = int(os.environ["WORLD_SIZE"]) torch.distributed.init_process_group("nccl", rank=rank, world_size=world_size) torch.cuda.set_device(rank) @@ -4110,10 +4110,12 @@ def from_pretrained( device_type = torch._C._get_accelerator().type device_module = torch.get_device_module(device_type) # Get device with index assuming equal number of devices per host - tp_device = torch.device(device_type, torch.distributed.get_rank() % device_module.device_count()) + tp_device = torch.device(device_type, torch.cuda.current_device()) + if tp_device.index > 0: + import sys + sys.stdout = open(os.devnull, "w") # This is the easiest way to dispatch to the current process device device_map = tp_device - # Assuming sharding the model onto the world world_size = torch.distributed.get_world_size() device_mesh = torch.distributed.init_device_mesh(tp_device.type, (world_size,)) @@ -5848,7 +5850,7 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict): if _torch_distributed_available and torch.distributed.is_initialized() else None ) - + print(accelerator_device_map) total_byte_count = defaultdict(lambda: 0) for param_name, device in accelerator_device_map.items(): param = model.get_parameter_or_buffer(param_name) From 426d9413771e08235669ded64e969d2a1f668750 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Tue, 25 Mar 2025 08:58:40 +0000 Subject: [PATCH 62/86] stash --- src/transformers/modeling_utils.py | 1 + src/transformers/models/deepseek_v3/modeling_deepseek_v3.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index b81183d4da2d..b49206b1473c 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -5850,6 +5850,7 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict): if _torch_distributed_available and torch.distributed.is_initialized() else None ) + print("Preallocating memory with cache allocator to improve loading perfs") print(accelerator_device_map) total_byte_count = defaultdict(lambda: 0) for param_name, device in accelerator_device_map.items(): diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 872eb89e77c1..f1bbb1e22807 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -98,7 +98,7 @@ def forward(self, x, position_ids): self._dynamic_frequency_update(position_ids, device=x.device) # Core RoPE block - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(position_ids.device) position_ids_expanded = position_ids[:, None, :].float() # Force float32 (see https://github.com/huggingface/transformers/pull/29285) device_type = x.device.type From f0a838917090ce2a057327e7440d645a094a9661 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Tue, 25 Mar 2025 11:01:26 +0000 Subject: [PATCH 63/86] changes currently requires --- .../integrations/finegrained_fp8.py | 8 +++---- .../integrations/tensor_parallel.py | 3 ++- src/transformers/modeling_utils.py | 23 +++++++++++-------- .../deepseek_v3/configuration_deepseek_v3.py | 14 ++++++++--- .../deepseek_v3/modeling_deepseek_v3.py | 2 +- 5 files changed, 31 insertions(+), 19 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 0815df434453..9f5001f152b8 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -291,7 +291,7 @@ def w8a8_block_fp8_matmul_compile( return output.to(output_dtype) -class FP8Linear(nn.Module): +class FP8Linear(nn.Linear): dtype = torch.float8_e4m3fn def __init__( @@ -304,11 +304,11 @@ def __init__( device=None, activation_scheme="dynamic", ): - super().__init__() + super().__init__(in_features,out_features) self.in_features = in_features self.out_features = out_features - self.register_buffer("weight", torch.empty(out_features, in_features, dtype=FP8Linear.dtype, device=device)) + self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=FP8Linear.dtype, device=device)) scale_out_features = (out_features + block_size[0] - 1) // block_size[0] scale_in_features = (in_features + block_size[1] - 1) // block_size[1] @@ -404,7 +404,7 @@ def replace_with_fp8_linear( if quantization_config.modules_to_not_convert is not None: modules_to_not_convert.extend(quantization_config.modules_to_not_convert) modules_to_not_convert = list(set(modules_to_not_convert)) - + print(modules_to_not_convert) model, has_been_replaced = _replace_with_fp8_linear( model, modules_to_not_convert=modules_to_not_convert, diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 7df881b9bac7..17fc96d1af67 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -487,7 +487,7 @@ def __init__(self): try: tp_layer.prepare_module_tp(module, device_mesh) except NotImplementedError as e: - print(f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP plan") + print(f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP plan: {e}") # 2. We add hooks to the parrent module if needed if "." in layer_name: @@ -534,6 +534,7 @@ def shard_and_distribute_module( param, empty_param, param_type, param_casting_dtype, is_contiguous, rank, device_mesh ) else: + print("NO PLAN FOR: ", parameter_name) param = param[...].to(param_casting_dtype) if is_contiguous: param = param.contiguous() diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index d1ad111e240d..21508933884a 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -773,7 +773,7 @@ def _load_state_dict_into_meta_model( device_map_regex = "|".join([re.escape(k) for k in sorted(device_map.keys(), reverse=True)]) is_quantized = hf_quantizer is not None - is_meta_state_dict = shard_file.endswith(".safetensors") and not is_quantized + is_meta_state_dict = shard_file.endswith(".safetensors")# and not is_quantized file_pointer = None if is_meta_state_dict: @@ -789,7 +789,7 @@ def _load_state_dict_into_meta_model( serialized_param_name = reverse_renaming_mapping[param_name] param = file_pointer.get_slice(serialized_param_name) else: - param = empty_param # It is actually not empty! + param = empty_param.to(tensor_device) # It is actually not empty! to_contiguous, casting_dtype = _infer_parameter_dtype( model, @@ -1362,8 +1362,10 @@ def _get_device_map( max_memory = hf_quantizer.adjust_max_memory(max_memory) device_map_kwargs["max_memory"] = max_memory + print(device_map) device_map = infer_auto_device_map(model, dtype=target_dtype, **device_map_kwargs) - + print(device_map) + exit(0) if hf_quantizer is not None: hf_quantizer.validate_environment(device_map=device_map) @@ -4145,11 +4147,12 @@ def from_pretrained( device_module = torch.get_device_module(device_type) # Get device with index assuming equal number of devices per host tp_device = torch.device(device_type, torch.cuda.current_device()) - if tp_device.index > 0: - import sys - sys.stdout = open(os.devnull, "w") + # if tp_device.index > 0: + # import sys + # sys.stdout = open(os.devnull, "w") # This is the easiest way to dispatch to the current process device device_map = tp_device + print("DEVICE MAP:", device_map) # Assuming sharding the model onto the world world_size = torch.distributed.get_world_size() device_mesh = torch.distributed.init_device_mesh(tp_device.type, (world_size,)) @@ -5030,7 +5033,7 @@ def _load_pretrained_model( warner = logger.warning if model.__class__.__name__ in archs else logger.info warner( f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" - f" initializing {model.__class__.__name__}: {update_key_name(unexpected_keys)}\n- This IS expected if you are" + f" initializing {model.__class__.__name__}: {"\n\t".join(update_key_name(unexpected_keys))}\n- This IS expected if you are" f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" " with another architecture (e.g. initializing a BertForSequenceClassification model from a" " BertForPreTraining model).\n- This IS NOT expected if you are initializing" @@ -5041,8 +5044,8 @@ def _load_pretrained_model( logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") if len(missing_keys) > 0: logger.warning( - f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized: {update_key_name(missing_keys)}\nYou should probably" + f"\nSome weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized: {"\n\t".join(update_key_name(missing_keys))}\nYou should probably" " TRAIN this model on a down-stream task to be able to use it for predictions and inference." ) elif len(mismatched_keys) == 0: @@ -5060,7 +5063,7 @@ def _load_pretrained_model( ] ) logger.warning( - f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f"\nSome weights of {model.__class__.__name__} were not initialized from the model checkpoint at" f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able" " to use it for predictions and inference." diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index ccfc97832c1b..13e34fd891df 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -135,11 +135,19 @@ class DeepseekV3Config(PretrainedConfig): # Default tensor parallel plan for base model `DeepseekV3Model` base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.q_a_proj": "colwise", + "layers.*.self_attn.kv_a_proj_with_mqa": "colwise", "layers.*.self_attn.kv_b_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", - "layers.*.gate_proj": "colwise", - "layers.*.up_proj": "colwise", - "layers.*.down_proj": "rowwise", + "layers.*.mlp.experts.*.gate_proj": "colwise", + "layers.*.mlp.experts.*.up_proj": "colwise", + "layers.*.mlp.experts.*.down_proj": "rowwise", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index f1bbb1e22807..36c2208b64e8 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -657,7 +657,7 @@ def __init__(self, config: DeepseekV3Config): self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( - [DeepseekV3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + [DeepseekV3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers + 1)] ) self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = DeepseekV3RotaryEmbedding(config=config) From 4b8a8578fc1b89776d7125043e610591b3604abe Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Tue, 25 Mar 2025 11:27:46 +0000 Subject: [PATCH 64/86] remove synch --- src/transformers/integrations/finegrained_fp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 9f5001f152b8..14b6deb18863 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -334,7 +334,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: qinput, scale = act_quant(input, self.block_size[1]) # Blocks the CPU until all CUDA operations on the specified device are complete. It is used to ensure that the results of the # preceding operations are ready before proceeding - torch.cuda.synchronize(device=input.device) + # torch.cuda.synchronize(device=self.weight.device) with torch.cuda.device(input.device): output = w8a8_block_fp8_matmul_triton( qinput, @@ -344,7 +344,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: self.block_size, output_dtype=input.dtype, ) - torch.cuda.synchronize(device=input.device) + # torch.cuda.synchronize(device=self.weight.device) if self.bias is not None: output = output + self.bias return output.to(dtype=input.dtype) From eedbf59935d71703580882dbc810f87985806bdc Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Tue, 25 Mar 2025 15:30:27 +0000 Subject: [PATCH 65/86] more fixes for TP --- .../integrations/finegrained_fp8.py | 16 +++++++++++----- .../deepseek_v3/configuration_deepseek_v3.py | 17 +++++++++-------- .../models/deepseek_v3/modeling_deepseek_v3.py | 4 ++++ 3 files changed, 24 insertions(+), 13 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 14b6deb18863..c575dcf3de54 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -310,11 +310,12 @@ def __init__( self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=FP8Linear.dtype, device=device)) - scale_out_features = (out_features + block_size[0] - 1) // block_size[0] - scale_in_features = (in_features + block_size[1] - 1) // block_size[1] - self.register_buffer( - "weight_scale_inv", torch.empty(scale_out_features, scale_in_features, dtype=torch.float32, device=device) - ) + if self.weight.element_size() == 1: + scale_out_features = (out_features + block_size[0] - 1) // block_size[0] + scale_in_features = (in_features + block_size[1] - 1) // block_size[1] + self.weight_scale_inv = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float32, device=device)) + else: + self.register_parameter("weight_scale_inv", None) self.block_size = block_size @@ -352,6 +353,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: def _replace_with_fp8_linear( model, + tp_plan=None, modules_to_not_convert=None, current_key_name=None, quantization_config=None, @@ -378,10 +380,13 @@ def _replace_with_fp8_linear( block_size=quantization_config.weight_block_size, ) has_been_replaced = True + # import re + # tp_plan[re.sub("\d+", "*", current_key_name_str) + ".weight_scale_inv"] = tp_plan[re.sub("\d+", "*", current_key_name_str)] if len(list(module.children())) > 0: _, has_been_replaced = _replace_with_fp8_linear( module, + tp_plan, modules_to_not_convert, current_key_name, quantization_config, @@ -407,6 +412,7 @@ def replace_with_fp8_linear( print(modules_to_not_convert) model, has_been_replaced = _replace_with_fp8_linear( model, + tp_plan=model._tp_plan, modules_to_not_convert=modules_to_not_convert, quantization_config=quantization_config, ) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 13e34fd891df..c1cf25a2bc9f 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -134,14 +134,15 @@ class DeepseekV3Config(PretrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `DeepseekV3Model` base_model_tp_plan = { - "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.q_a_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "colwise", - "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.*.gate_proj": "colwise", - "layers.*.mlp.experts.*.up_proj": "colwise", - "layers.*.mlp.experts.*.down_proj": "rowwise", + # "layers.*.self_attn.q_b_proj": "colwise", + # "layers.*.self_attn.q_a_proj": "colwise", + # "layers.*.self_attn.kv_a_proj_with_mqa": "colwise", + # "layers.*.self_attn.kv_b_proj": "colwise", + # "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.experts.*.gate_proj": "local_colwise", + "layers.*.mlp.experts.*.up_proj": "local_colwise", + "layers.*.mlp.experts.*.down_proj": "local_rowwise", + "layers.*.mlp.experts": "isolated_parallel", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", "layers.*.mlp.shared_experts.down_proj": "rowwise", diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 36c2208b64e8..6ec116b8d724 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -223,6 +223,10 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig expert_output = expert(expert_input) weighted_output = expert_output * expert_weights.unsqueeze(-1) final_hidden_states.index_add_(0, token_indices, weighted_output) + + # in original deepseek, the output of the experts are gathered once we leave this module + # thus the moe module is itelsf an IsolatedParallel module + # and all expert are "local" meaning we shard but we don't gather return final_hidden_states.type(hidden_states.dtype) From 409f3412b0de40d6426d329fca228496012b3dd7 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Wed, 26 Mar 2025 16:56:58 +0000 Subject: [PATCH 66/86] temp fix for TP : some attention layers's FP8 scales are too small + shared is local colwise and anything is local if FP8 because weights are used --- src/transformers/integrations/finegrained_fp8.py | 6 +++--- .../deepseek_v3/configuration_deepseek_v3.py | 15 ++++++++------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index c575dcf3de54..e65cc2e82c92 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -331,8 +331,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return F.linear(input, self.weight, self.bias) else: # Context manager used to switch among the available cuda devices - with torch.cuda.device(input.device): - qinput, scale = act_quant(input, self.block_size[1]) + # with torch.cuda.device(input.device): + qinput, scale = act_quant(input, self.block_size[1]) # Blocks the CPU until all CUDA operations on the specified device are complete. It is used to ensure that the results of the # preceding operations are ready before proceeding # torch.cuda.synchronize(device=self.weight.device) @@ -345,7 +345,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: self.block_size, output_dtype=input.dtype, ) - # torch.cuda.synchronize(device=self.weight.device) + torch.cuda.synchronize() if self.bias is not None: output = output + self.bias return output.to(dtype=input.dtype) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index c1cf25a2bc9f..ecedef5f5023 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -142,13 +142,14 @@ class DeepseekV3Config(PretrainedConfig): "layers.*.mlp.experts.*.gate_proj": "local_colwise", "layers.*.mlp.experts.*.up_proj": "local_colwise", "layers.*.mlp.experts.*.down_proj": "local_rowwise", - "layers.*.mlp.experts": "isolated_parallel", - "layers.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.experts": "gather", + "layers.*.mlp.shared_experts.gate_proj": "local_colwise", + "layers.*.mlp.shared_experts.up_proj": "local_colwise", + "layers.*.mlp.shared_experts.down_proj": "local_rowwise", + "layers.*.mlp.shared_experts": "gather", # local's job + "layers.*.mlp.gate_proj": "local_colwise", + "layers.*.mlp.up_proj": "local_colwise", + "layers.*.mlp.down_proj": "local_rowwise", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), From 3fb9bea5ff382afba9ef69e10806ed03a5e1252a Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Thu, 27 Mar 2025 15:55:55 +0000 Subject: [PATCH 67/86] updates to have generation work! --- .../integrations/finegrained_fp8.py | 4 +++- .../integrations/tensor_parallel.py | 6 ++--- src/transformers/modeling_utils.py | 23 ++++++++----------- .../deepseek_v3/configuration_deepseek_v3.py | 5 ++-- 4 files changed, 19 insertions(+), 19 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index e65cc2e82c92..7e3f1df057f1 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -381,7 +381,9 @@ def _replace_with_fp8_linear( ) has_been_replaced = True # import re - # tp_plan[re.sub("\d+", "*", current_key_name_str) + ".weight_scale_inv"] = tp_plan[re.sub("\d+", "*", current_key_name_str)] + # TODO local should be determined here + # tp_plan[re.sub("\d+", "*", current_key_name_str)] = "local_" + tp_plan[re.sub("\d+", "*", current_key_name_str)] + # tp_plan[re.sub("\d+", "*", current_key_name_str).rsplit(".")[0]] = "gather" if len(list(module.children())) > 0: _, has_been_replaced = _replace_with_fp8_linear( diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 17fc96d1af67..1a0c017f6f09 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -231,8 +231,8 @@ def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module: distribute_module( module, device_mesh, - partial(self._prepare_input_fn), - partial(self._prepare_output_fn), + partial(self._prepare_input_fn, None, None), + partial(self._prepare_output_fn, None, None), ) @@ -534,7 +534,7 @@ def shard_and_distribute_module( param, empty_param, param_type, param_casting_dtype, is_contiguous, rank, device_mesh ) else: - print("NO PLAN FOR: ", parameter_name) + # TODO log no plan modules in set param = param[...].to(param_casting_dtype) if is_contiguous: param = param.contiguous() diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 21508933884a..5f2578371527 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4132,10 +4132,10 @@ def from_pretrained( if not torch.distributed.is_initialized(): try: logger.warning("Tensor Parallel requires torch.distributed to be initialized first.") - rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) - torch.distributed.init_process_group("nccl", rank=rank, world_size=world_size) - torch.cuda.set_device(rank) + torch.distributed.init_process_group("nccl", rank=rank, world_size=world_size, init_method="env://") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) except Exception as e: raise EnvironmentError( "We tried to initialize torch.distributed for you, but it failed, make" @@ -4147,12 +4147,11 @@ def from_pretrained( device_module = torch.get_device_module(device_type) # Get device with index assuming equal number of devices per host tp_device = torch.device(device_type, torch.cuda.current_device()) - # if tp_device.index > 0: - # import sys - # sys.stdout = open(os.devnull, "w") + if tp_device.index > 0: + import sys + sys.stdout = open(os.devnull, "w") # This is the easiest way to dispatch to the current process device device_map = tp_device - print("DEVICE MAP:", device_map) # Assuming sharding the model onto the world world_size = torch.distributed.get_world_size() device_mesh = torch.distributed.init_device_mesh(tp_device.type, (world_size,)) @@ -4893,7 +4892,7 @@ def _load_pretrained_model( expected_keys = hf_quantizer.update_expected_keys(model_to_load, expected_keys, checkpoint_keys) # Warmup cuda to load the weights much faster on devices - if device_map is not None and hf_quantizer is None: + if device_map is not None: #and hf_quantizer is None: expanded_device_map = expand_device_map(device_map, expected_keys) caching_allocator_warmup(model_to_load, expanded_device_map) @@ -5033,7 +5032,7 @@ def _load_pretrained_model( warner = logger.warning if model.__class__.__name__ in archs else logger.info warner( f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" - f" initializing {model.__class__.__name__}: {"\n\t".join(update_key_name(unexpected_keys))}\n- This IS expected if you are" + f" initializing {model.__class__.__name__}:\n\t{"\n\t".join(update_key_name(unexpected_keys))}\n- This IS expected if you are" f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" " with another architecture (e.g. initializing a BertForSequenceClassification model from a" " BertForPreTraining model).\n- This IS NOT expected if you are initializing" @@ -5045,7 +5044,7 @@ def _load_pretrained_model( if len(missing_keys) > 0: logger.warning( f"\nSome weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized: {"\n\t".join(update_key_name(missing_keys))}\nYou should probably" + f" {pretrained_model_name_or_path} and are newly initialized:\n\t{"\n\t".join(update_key_name(missing_keys))}\nYou should probably" " TRAIN this model on a down-stream task to be able to use it for predictions and inference." ) elif len(mismatched_keys) == 0: @@ -5887,8 +5886,6 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict): if _torch_distributed_available and torch.distributed.is_initialized() else None ) - print("Preallocating memory with cache allocator to improve loading perfs") - print(accelerator_device_map) total_byte_count = defaultdict(lambda: 0) for param_name, device in accelerator_device_map.items(): param = model.get_parameter_or_buffer(param_name) @@ -5908,7 +5905,7 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict): # Allow up to 95% of max device memory byte_count = min(byte_count, int(0.95 * device_memory)) # Allocate memory - _ = torch.empty(byte_count // 2, dtype=torch.float16, device=device, requires_grad=False) + _ = torch.empty(byte_count // 4, dtype=torch.float16, device=device, requires_grad=False) def get_disk_only_shard_files(device_map, weight_map): diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index ecedef5f5023..3231e2f542b9 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -142,14 +142,15 @@ class DeepseekV3Config(PretrainedConfig): "layers.*.mlp.experts.*.gate_proj": "local_colwise", "layers.*.mlp.experts.*.up_proj": "local_colwise", "layers.*.mlp.experts.*.down_proj": "local_rowwise", - "layers.*.mlp.experts": "gather", + "layers.*.mlp.experts.*": "local", # each expert is wrapped in a module list "layers.*.mlp.shared_experts.gate_proj": "local_colwise", "layers.*.mlp.shared_experts.up_proj": "local_colwise", "layers.*.mlp.shared_experts.down_proj": "local_rowwise", - "layers.*.mlp.shared_experts": "gather", # local's job + "layers.*.mlp.shared_experts": "local", # local's job "layers.*.mlp.gate_proj": "local_colwise", "layers.*.mlp.up_proj": "local_colwise", "layers.*.mlp.down_proj": "local_rowwise", + "layers.*.mlp": "gather" } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), From 7350a5d435228b22ffaaa1a9528564e3446de617 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 09:43:02 +0000 Subject: [PATCH 68/86] push most of the changes --- .../integrations/finegrained_fp8.py | 6 +- .../integrations/tensor_parallel.py | 4 +- src/transformers/modeling_utils.py | 11 +- .../deepseek_v3/configuration_deepseek_v3.py | 6 +- .../deepseek_v3/modeling_deepseek_v3.py | 35 ++++-- .../models/deepseek_v3/modular_deepseek_v3.py | 119 +++--------------- 6 files changed, 63 insertions(+), 118 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 7e3f1df057f1..19bd34cdbdbe 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -304,7 +304,7 @@ def __init__( device=None, activation_scheme="dynamic", ): - super().__init__(in_features,out_features) + super().__init__(in_features, out_features) self.in_features = in_features self.out_features = out_features @@ -313,7 +313,9 @@ def __init__( if self.weight.element_size() == 1: scale_out_features = (out_features + block_size[0] - 1) // block_size[0] scale_in_features = (in_features + block_size[1] - 1) // block_size[1] - self.weight_scale_inv = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float32, device=device)) + self.weight_scale_inv = nn.Parameter( + torch.empty(scale_out_features, scale_in_features, dtype=torch.float32, device=device) + ) else: self.register_parameter("weight_scale_inv", None) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 1a0c017f6f09..f8a04c037b93 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -487,7 +487,9 @@ def __init__(self): try: tp_layer.prepare_module_tp(module, device_mesh) except NotImplementedError as e: - print(f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP plan: {e}") + print( + f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP plan: {e}" + ) # 2. We add hooks to the parrent module if needed if "." in layer_name: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 5f2578371527..1e9fd2ea6b32 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -773,7 +773,7 @@ def _load_state_dict_into_meta_model( device_map_regex = "|".join([re.escape(k) for k in sorted(device_map.keys(), reverse=True)]) is_quantized = hf_quantizer is not None - is_meta_state_dict = shard_file.endswith(".safetensors")# and not is_quantized + is_meta_state_dict = shard_file.endswith(".safetensors") # and not is_quantized file_pointer = None if is_meta_state_dict: @@ -4134,7 +4134,9 @@ def from_pretrained( logger.warning("Tensor Parallel requires torch.distributed to be initialized first.") rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) - torch.distributed.init_process_group("nccl", rank=rank, world_size=world_size, init_method="env://") + torch.distributed.init_process_group( + "nccl", rank=rank, world_size=world_size, init_method="env://" + ) torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) except Exception as e: raise EnvironmentError( @@ -4144,11 +4146,10 @@ def from_pretrained( # Detect the accelerator on the machine. If no accelerator is available, it returns CPU. device_type = torch._C._get_accelerator().type - device_module = torch.get_device_module(device_type) - # Get device with index assuming equal number of devices per host tp_device = torch.device(device_type, torch.cuda.current_device()) if tp_device.index > 0: import sys + sys.stdout = open(os.devnull, "w") # This is the easiest way to dispatch to the current process device device_map = tp_device @@ -4892,7 +4893,7 @@ def _load_pretrained_model( expected_keys = hf_quantizer.update_expected_keys(model_to_load, expected_keys, checkpoint_keys) # Warmup cuda to load the weights much faster on devices - if device_map is not None: #and hf_quantizer is None: + if device_map is not None: # and hf_quantizer is None: expanded_device_map = expand_device_map(device_map, expected_keys) caching_allocator_warmup(model_to_load, expanded_device_map) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 3231e2f542b9..1c609052977c 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -142,15 +142,15 @@ class DeepseekV3Config(PretrainedConfig): "layers.*.mlp.experts.*.gate_proj": "local_colwise", "layers.*.mlp.experts.*.up_proj": "local_colwise", "layers.*.mlp.experts.*.down_proj": "local_rowwise", - "layers.*.mlp.experts.*": "local", # each expert is wrapped in a module list + "layers.*.mlp.experts.*": "local", # each expert is wrapped in a module list "layers.*.mlp.shared_experts.gate_proj": "local_colwise", "layers.*.mlp.shared_experts.up_proj": "local_colwise", "layers.*.mlp.shared_experts.down_proj": "local_rowwise", - "layers.*.mlp.shared_experts": "local", # local's job + "layers.*.mlp.shared_experts": "local", # local's job "layers.*.mlp.gate_proj": "local_colwise", "layers.*.mlp.up_proj": "local_colwise", "layers.*.mlp.down_proj": "local_rowwise", - "layers.*.mlp": "gather" + "layers.*.mlp": "gather", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 6ec116b8d724..6e0228097e0c 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -24,6 +24,7 @@ LossKwargs, add_start_docstrings, add_start_docstrings_to_model_forward, + is_torch_flex_attn_available, logging, replace_return_docstrings, ) @@ -31,6 +32,12 @@ from .configuration_deepseek_v3 import DeepseekV3Config +if is_torch_flex_attn_available(): + from torch.nn.attention.flex_attention import BlockMask + + from ...integrations.flex_attention import make_flex_block_causal_mask + + logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "DeepseekV3Config" @@ -98,13 +105,13 @@ def forward(self, x, position_ids): self._dynamic_frequency_update(position_ids, device=x.device) # Core RoPE block - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(position_ids.device) + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() # Force float32 (see https://github.com/huggingface/transformers/pull/29285) device_type = x.device.type device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" with torch.autocast(device_type=device_type, enabled=False): - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + freqs = (inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() @@ -208,6 +215,10 @@ def forward(self, hidden_states): return hidden_states def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + r""" + CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused + to not have to do a loop here (deepseek has 256 experts soooo yeah). + """ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) expert_mask = expert_mask.permute(2, 0, 1) @@ -303,7 +314,11 @@ def eager_attention_forward( def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. + r""" + TODO let's just use the original freqcis computation to not have the view + transpose + reshape! This is not optimized! + Applies Rotary Position Embedding to the query and key tensors. + Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. @@ -654,6 +669,8 @@ class DeepseekV3Model(DeepseekV3PreTrainedModel): config: DeepseekV3Config """ + _keys_to_ignore_on_load_unexpected = [r"model.layers.61.*"] + def __init__(self, config: DeepseekV3Config): super().__init__(config) self.padding_idx = config.pad_token_id @@ -661,7 +678,7 @@ def __init__(self, config: DeepseekV3Config): self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( - [DeepseekV3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers + 1)] + [DeepseekV3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = DeepseekV3RotaryEmbedding(config=config) @@ -789,12 +806,17 @@ def _update_causal_mask( input_tensor: torch.Tensor, cache_position: torch.Tensor, past_key_values: Cache, - output_attentions: bool, + output_attentions: bool = False, ): if self.config._attn_implementation == "flash_attention_2": if attention_mask is not None and (attention_mask == 0.0).any(): return attention_mask return None + if self.config._attn_implementation == "flex_attention": + if isinstance(attention_mask, torch.Tensor): + attention_mask = make_flex_block_causal_mask(attention_mask) + if isinstance(attention_mask, BlockMask): + return attention_mask # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail @@ -875,7 +897,7 @@ def _prepare_4d_causal_attention_mask_with_cache_position( dtype (`torch.dtype`): The dtype to use for the 4D attention mask. device (`torch.device`): - The device to plcae the 4D attention mask on. + The device to place the 4D attention mask on. cache_position (`torch.Tensor`): Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): @@ -962,7 +984,6 @@ def forward( **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: 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 diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index fe5ea154d8d0..1cd26c09878a 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -18,6 +18,7 @@ LlamaPreTrainedModel, LlamaRMSNorm, LlamaRotaryEmbedding, + LlamaDecoderLayer, apply_rotary_pos_emb, eager_attention_forward, rotate_half, @@ -37,7 +38,11 @@ class DeepseekV3RotaryEmbedding(LlamaRotaryEmbedding): def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. + r""" + TODO let's just use the original freqcis computation to not have the view + transpose + reshape! This is not optimized! + Applies Rotary Position Embedding to the query and key tensors. + Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. @@ -168,6 +173,10 @@ def forward(self, hidden_states): return hidden_states def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + r""" + CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused + to not have to do a loop here (deepseek has 256 experts soooo yeah). + """ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) expert_mask = expert_mask.permute(2, 0, 1) @@ -183,8 +192,11 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig expert_output = expert(expert_input) weighted_output = expert_output * expert_weights.unsqueeze(-1) final_hidden_states.index_add_(0, token_indices, weighted_output) - return final_hidden_states.type(hidden_states.dtype) + # in original deepseek, the output of the experts are gathered once we leave this module + # thus the moe module is itelsf an IsolatedParallel module + # and all expert are "local" meaning we shard but we don't gather + return final_hidden_states.type(hidden_states.dtype) class DeepseekV3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -260,7 +272,7 @@ def forward( k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings - if self.config.rope_interleave: + if self.config.rope_interleave: # support using interleaved weights for efficiency q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) else: q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) @@ -306,9 +318,9 @@ def forward( return attn_output, attn_weights -class DeepseekV3DecoderLayer(nn.Module): +class DeepseekV3DecoderLayer(LlamaDecoderLayer, nn.Module): def __init__(self, config: DeepseekV3Config, layer_idx: int): - super().__init__() + nn.Module().__init__() self.hidden_size = config.hidden_size self.self_attn = DeepseekV3Attention(config=config, layer_idx=layer_idx) @@ -321,106 +333,13 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Cache] = None, - output_attentions: Optional[bool] = False, - use_cache: Optional[bool] = False, - cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC - **kwargs: Unpack[FlashAttentionKwargs], - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - # Self Attention - hidden_states, self_attn_weights = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - **kwargs, - ) - hidden_states = residual + hidden_states - - # Fully Connected - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) - - return outputs - class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): pass class DeepseekV3Model(LlamaModel): - pass - # def __init__(self, config): - # super().__init__(config) - # self._register_load_state_dict_pre_hook(self.load_pre_hook) - # self._register_state_dict_hook(self.load_hook) - # self.post_init() - - # def load_pre_hook(self, state_dict, prefix, *args): - # """ - # Weights have to be permuted for correct rope formulation. We can't do this in the weights - # as every other framework already uses the `Llama` original function (which is copyrighted btw). - # And I am not even sure it's better.... anyways end of my rant - # """ - - # def permute_for_rope(input_tensor): - # """ - # When you go from the complex ROPE formulation to sin and cos one, you need - # to permute the query and key weights (to avoid doing it on the fly) - # """ - # n_heads, dim1, dim2 = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] - # input_tensor = input_tensor.reshape(n_heads * dim1, dim2) - # input_tensor = input_tensor.view(n_heads, dim1 // 2, 2, dim2) - # input_tensor = input_tensor.transpose(1, 2).reshape(n_heads, dim1, dim2) - # return input_tensor - - # def permute_layer_for_rope(key, num_heads, head_dim, rope_dim): - # weight = state_dict[key] - # weight = weight.view(num_heads, head_dim, -1) - # weight_rot = weight[:, -rope_dim:] - # weight_rot = permute_for_rope(weight_rot) - # weight[:, -rope_dim:] = weight_rot - # weight = weight.view(-1, weight.shape[-1]) - # state_dict[key] = weight - - # for k in state_dict: - # if "q_b_proj." in k: - # permute_layer_for_rope( - # k, - # num_heads=self.config.num_attention_heads, - # head_dim=self.config.qk_head_dim, - # rope_dim=self.config.qk_rope_head_dim, - # ) - # if "kv_a_proj_with_mqa." in k: - # permute_layer_for_rope( - # k, - # num_heads=1, - # head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, - # rope_dim=self.config.qk_rope_head_dim, - # ) - - # def load_hook(self, module, state_dict, prefix, *args): - # self.load_pre_hook(state_dict, prefix, *args) + _keys_to_ignore_on_load_unexpected = [r"model.layers.61.*"] class DeepseekV3ForCausalLM(LlamaForCausalLM): @@ -431,4 +350,4 @@ class DeepseekV3ForCausalLM(LlamaForCausalLM): "DeepseekV3PreTrainedModel", "DeepseekV3Model", "DeepseekV3ForCausalLM", -] +] \ No newline at end of file From a50c3512e1fba917586855a09f54cee3cbc70b3d Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 09:56:12 +0000 Subject: [PATCH 69/86] reorder functions + call for contributions! --- docs/source/en/model_doc/deepseek_v3.md | 9 ++++ .../deepseek_v3/modeling_deepseek_v3.py | 44 +++++++++---------- .../models/deepseek_v3/modular_deepseek_v3.py | 40 ++++++++--------- 3 files changed, 51 insertions(+), 42 deletions(-) diff --git a/docs/source/en/model_doc/deepseek_v3.md b/docs/source/en/model_doc/deepseek_v3.md index e44d78739069..6c9fc4b1de91 100644 --- a/docs/source/en/model_doc/deepseek_v3.md +++ b/docs/source/en/model_doc/deepseek_v3.md @@ -23,6 +23,15 @@ The DeepSeek-V3 model was proposed in [DeepSeek-V3 Technical Report](https://arx The abstract from the paper is the following: We present DeepSeek-V3, a strong Mixture-of-Experts (MoE) language model with 671B total parameters with 37B activated for each token. To achieve efficient inference and cost-effective training, DeepSeek-V3 adopts Multi-head Latent Attention (MLA) and DeepSeekMoE architectures, which were thoroughly validated in DeepSeek-V2. Furthermore, DeepSeek-V3 pioneers an auxiliary-loss-free strategy for load balancing and sets a multi-token prediction training objective for stronger performance. We pre-train DeepSeek-V3 on 14.8 trillion diverse and high-quality tokens, followed by Supervised Fine-Tuning and Reinforcement Learning stages to fully harness its capabilities. Comprehensive evaluations reveal that DeepSeek-V3 outperforms other open-source models and achieves performance comparable to leading closed-source models. Despite its excellent performance, DeepSeek-V3 requires only 2.788M H800 GPU hours for its full training. In addition, its training process is remarkably stable. Throughout the entire training process, we did not experience any irrecoverable loss spikes or perform any rollbacks. The model checkpoints are available at https://github.com/deepseek-ai/DeepSeek-V3. +## Limitations and call for contribution! + +We are super happy to make this code community-powered, and would love to see how you can best optimize the following: + +- current implementation uses the "naive" attention compution (so not really MLA) +- current implementation loops through the experts. This should be replaced. Pointers to use `get_packed_weights` from `intetrations/tensor_parallel`. +- current implementation uses the eleuther formula for ROPE, using the orginal one would be more efficient! (should still follow our API) +- static cache is not supported (this should be just a generation config issue / config shape issues) + ### Usage tips The model uses Multi-head Latent Attention (MLA) and DeepSeekMoE architectures for efficient inference and cost-effective training. It employs an auxiliary-loss-free strategy for load balancing and multi-token prediction training objective. The model can be used for various language tasks after being pre-trained on 14.8 trillion tokens and going through Supervised Fine-Tuning and Reinforcement Learning stages. diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 6e0228097e0c..c00a95ce3a47 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -154,18 +154,6 @@ def __init__(self, config): self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts))) - def forward(self, hidden_states): - hidden_states = hidden_states.view(-1, self.config.hidden_size) - router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) - scores = router_logits.sigmoid() - topk_indices = self.get_topk_indices(scores) - topk_weights = scores.gather(1, topk_indices) - if self.norm_topk_prob: - denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 - topk_weights /= denominator - topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights - def get_topk_indices(self, scores): scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( @@ -185,6 +173,18 @@ def get_topk_indices(self, scores): topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] return topk_indices + def forward(self, hidden_states): + hidden_states = hidden_states.view(-1, self.config.hidden_size) + router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + scores = router_logits.sigmoid() + topk_indices = self.get_topk_indices(scores) + topk_weights = scores.gather(1, topk_indices) + if self.norm_topk_prob: + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator + topk_weights = topk_weights * self.routed_scaling_factor + return topk_indices, topk_weights + class DeepseekV3MoE(nn.Module): """ @@ -205,15 +205,6 @@ def __init__(self, config): config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts ) - def forward(self, hidden_states): - residuals = hidden_states - orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) - hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states - def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): r""" CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused @@ -240,6 +231,15 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig # and all expert are "local" meaning we shard but we don't gather return final_hidden_states.type(hidden_states.dtype) + def forward(self, hidden_states): + residuals = hidden_states + orig_shape = hidden_states.shape + topk_indices, topk_weights = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = hidden_states + self.shared_experts(residuals) + return hidden_states + def rotate_half(x): """Rotates half the hidden dims of the input.""" @@ -431,7 +431,7 @@ def forward( k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings - if self.config.rope_interleave: + if self.config.rope_interleave: # support using interleaved weights for efficiency q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) else: q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 1cd26c09878a..9c27558f7995 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -112,18 +112,6 @@ def __init__(self, config): self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts))) - def forward(self, hidden_states): - hidden_states = hidden_states.view(-1, self.config.hidden_size) - router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) - scores = router_logits.sigmoid() - topk_indices = self.get_topk_indices(scores) - topk_weights = scores.gather(1, topk_indices) - if self.norm_topk_prob: - denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 - topk_weights /= denominator - topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights - def get_topk_indices(self, scores): scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( @@ -143,6 +131,17 @@ def get_topk_indices(self, scores): topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] return topk_indices + def forward(self, hidden_states): + hidden_states = hidden_states.view(-1, self.config.hidden_size) + router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + scores = router_logits.sigmoid() + topk_indices = self.get_topk_indices(scores) + topk_weights = scores.gather(1, topk_indices) + if self.norm_topk_prob: + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator + topk_weights = topk_weights * self.routed_scaling_factor + return topk_indices, topk_weights class DeepseekV3MoE(nn.Module): """ @@ -163,14 +162,6 @@ def __init__(self, config): config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts ) - def forward(self, hidden_states): - residuals = hidden_states - orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) - hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): r""" @@ -198,6 +189,15 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig # and all expert are "local" meaning we shard but we don't gather return final_hidden_states.type(hidden_states.dtype) + def forward(self, hidden_states): + residuals = hidden_states + orig_shape = hidden_states.shape + topk_indices, topk_weights = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = hidden_states + self.shared_experts(residuals) + return hidden_states + class DeepseekV3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" From 24557c37c3ed13fbfbd065ce111b64f638ee3de5 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 10:03:06 +0000 Subject: [PATCH 70/86] update readme --- docs/source/en/model_doc/deepseek_v3.md | 39 +++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/docs/source/en/model_doc/deepseek_v3.md b/docs/source/en/model_doc/deepseek_v3.md index 6c9fc4b1de91..1a300e2c00df 100644 --- a/docs/source/en/model_doc/deepseek_v3.md +++ b/docs/source/en/model_doc/deepseek_v3.md @@ -35,6 +35,45 @@ We are super happy to make this code community-powered, and would love to see ho ### Usage tips The model uses Multi-head Latent Attention (MLA) and DeepSeekMoE architectures for efficient inference and cost-effective training. It employs an auxiliary-loss-free strategy for load balancing and multi-token prediction training objective. The model can be used for various language tasks after being pre-trained on 14.8 trillion tokens and going through Supervised Fine-Tuning and Reinforcement Learning stages. +You can run the model in `FP8` automatically, using 2 nodes of 8 H100 should be more than enough! + +```python +# `run_deepseek_v1.py` +from transformers import AutoModelForCausalLM, AutoTokenizer +import torch +torch.manual_seed(30) + +tokenizer = AutoTokenizer.from_pretrained("deepseek-r1") + +chat = [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing great. How can I help you today?"}, + {"role": "user", "content": "I'd like to show off how chat templating works!"}, +] + + +model = AutoModelForCausalLM.from_pretrained("deepseek-r1", device_map="auto", torch_dtype=torch.bfloat16) +inputs = tokenizer.apply_chat_template(chat, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device) +import time +start = time.time() +outputs = model.generate(inputs, max_new_tokens=50) +print(tokenizer.batch_decode(outputs)) +print(time.time()-start) +``` +Use the following to run it +```bash +torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0|1 --rdzv-id an_id --rdzv-backend c10d --rdzv-endpoint master_addr:master_port run_deepseek_r1.py +``` + +If you have: +```bash +[rank0]: ncclInternalError: Internal check failed. +[rank0]: Last error: +[rank0]: Bootstrap : no socket interface found +``` +error, it means NCCL was probably not loaded. + + ## DeepseekV3Config [[autodoc]] DeepseekV3Config From d7da38bc3dd13e1aff2ba0cba2c5af9670146d9c Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 10:03:38 +0000 Subject: [PATCH 71/86] nits --- .../models/deepseek_v3/modular_deepseek_v3.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 9c27558f7995..c1bddffa633a 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -13,12 +13,12 @@ from ...processing_utils import Unpack from ...utils import logging from ..llama.modeling_llama import ( + LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, LlamaPreTrainedModel, LlamaRMSNorm, LlamaRotaryEmbedding, - LlamaDecoderLayer, apply_rotary_pos_emb, eager_attention_forward, rotate_half, @@ -39,10 +39,10 @@ class DeepseekV3RotaryEmbedding(LlamaRotaryEmbedding): def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): r""" - TODO let's just use the original freqcis computation to not have the view + TODO let's just use the original freqcis computation to not have the view transpose + reshape! This is not optimized! Applies Rotary Position Embedding to the query and key tensors. - + Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. @@ -143,6 +143,7 @@ def forward(self, hidden_states): topk_weights = topk_weights * self.routed_scaling_factor return topk_indices, topk_weights + class DeepseekV3MoE(nn.Module): """ A mixed expert module containing shared experts. @@ -162,11 +163,10 @@ def __init__(self, config): config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts ) - def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): r""" - CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused - to not have to do a loop here (deepseek has 256 experts soooo yeah). + CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused + to not have to do a loop here (deepseek has 256 experts soooo yeah). """ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) @@ -198,6 +198,7 @@ def forward(self, hidden_states): hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states + class DeepseekV3Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -272,7 +273,7 @@ def forward( k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) cos, sin = position_embeddings - if self.config.rope_interleave: # support using interleaved weights for efficiency + if self.config.rope_interleave: # support using interleaved weights for efficiency q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) else: q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) @@ -350,4 +351,4 @@ class DeepseekV3ForCausalLM(LlamaForCausalLM): "DeepseekV3PreTrainedModel", "DeepseekV3Model", "DeepseekV3ForCausalLM", -] \ No newline at end of file +] From 186e32b90235afc0c0b8eb1484f3f43b72e4adba Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 10:07:09 +0000 Subject: [PATCH 72/86] update --- src/transformers/models/deepseek_v3/modular_deepseek_v3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index c1bddffa633a..f621f03aefc4 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -340,7 +340,7 @@ class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): class DeepseekV3Model(LlamaModel): - _keys_to_ignore_on_load_unexpected = [r"model.layers.61.*"] + _keys_to_ignore_on_load_unexpected = [r"model\.layers\.61.*"] class DeepseekV3ForCausalLM(LlamaForCausalLM): From ee33cf7b7376ae1d19ac6ff3f129adcc571c66b6 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 10:13:14 +0000 Subject: [PATCH 73/86] ruff was updated on main --- src/transformers/modeling_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 7a49f9e250a3..151df32ea658 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -5050,7 +5050,7 @@ def _load_pretrained_model( warner = logger.warning if model.__class__.__name__ in archs else logger.info warner( f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" - f" initializing {model.__class__.__name__}:\n\t{"\n\t".join(update_key_name(unexpected_keys))}\n- This IS expected if you are" + f" initializing {model.__class__.__name__}:\n\t{'\n\t'.join(update_key_name(unexpected_keys))}\n- This IS expected if you are" f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" " with another architecture (e.g. initializing a BertForSequenceClassification model from a" " BertForPreTraining model).\n- This IS NOT expected if you are initializing" @@ -5062,7 +5062,7 @@ def _load_pretrained_model( if len(missing_keys) > 0: logger.warning( f"\nSome weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized:\n\t{"\n\t".join(update_key_name(missing_keys))}\nYou should probably" + f" {pretrained_model_name_or_path} and are newly initialized:\n\t{'\n\t'.join(update_key_name(missing_keys))}\nYou should probably" " TRAIN this model on a down-stream task to be able to use it for predictions and inference." ) elif len(mismatched_keys) == 0: From f2bb6f982f610e50991d94a7cb7b8d3738a029b2 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 10:17:06 +0000 Subject: [PATCH 74/86] merge with main and fix copies --- .../deepseek_v3/modeling_deepseek_v3.py | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index c00a95ce3a47..8001bcfff905 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -5,7 +5,7 @@ # modular_deepseek_v3.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 import math -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -617,20 +617,12 @@ def _init_weights(self, module): config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) - past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + past_key_values (`Cache`, *optional*): Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. - Two formats are allowed: - - a [`~cache_utils.Cache`] instance, see our - [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache); - - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of - shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy - cache format. - - The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the - legacy cache format will be returned. + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` @@ -669,7 +661,7 @@ class DeepseekV3Model(DeepseekV3PreTrainedModel): config: DeepseekV3Config """ - _keys_to_ignore_on_load_unexpected = [r"model.layers.61.*"] + _keys_to_ignore_on_load_unexpected = [r"model\.layers\.61.*"] def __init__(self, config: DeepseekV3Config): super().__init__(config) @@ -724,6 +716,10 @@ def forward( ) use_cache = False + # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache + if not isinstance(past_key_values, (type(None), Cache)): + raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.") + if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -972,7 +968,7 @@ def forward( input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, From 8cefd1c97622bb85267519cc8dbcd86c73536282 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 11:36:28 +0100 Subject: [PATCH 75/86] revert unrelated changes --- src/transformers/modeling_utils.py | 46 ++++-------------------------- tests/utils/test_modeling_utils.py | 40 -------------------------- 2 files changed, 5 insertions(+), 81 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 151df32ea658..8b6b2c831a4b 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -891,40 +891,6 @@ def _add_variant(weights_name: str, variant: Optional[str] = None) -> str: return weights_name -def update_key_name(keys): - """ - Updates a dictionary of keys to pack layers together as layer.{0, 1, 4} instead of layers.0, layers.1, layers.4. - """ - key_dict = defaultdict(list) - for key in keys: - all_digits = re.findall(r".(\d+).", key) - for i, k in enumerate(all_digits): - if len(key_dict[re.sub(r".(\d+).", ".*.", key)]) <= i: - key_dict[re.sub(r".(\d+).", ".*.", key)].append(set()) - key_dict[re.sub(r".(\d+).", ".*.", key)][i].add(int(k)) - - final_keys = set() - for key in keys: - text = re.sub(r".(\d+).", ".*.", key) - pattern = key_dict[text] - final_text = "" - for i, part in enumerate(text.split("*")): - if len(pattern) <= i: - final_text += part - else: - data = [str(i) for i in sorted(pattern[i])] - if len(data) > 10: - result = f"{data[0]}...{data[-1]}" - else: - result = ", ".join(data) # If there are only 1 or 2 elements, show them all - if len(data) > 1: - final_text += part + "{" + result + "}" - else: - final_text += part + data[0] - final_keys.add(final_text) - return list(final_keys) - - def _get_resolved_checkpoint_files( pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], subfolder: str, @@ -1368,10 +1334,8 @@ def _get_device_map( max_memory = hf_quantizer.adjust_max_memory(max_memory) device_map_kwargs["max_memory"] = max_memory - print(device_map) device_map = infer_auto_device_map(model, dtype=target_dtype, **device_map_kwargs) - print(device_map) - exit(0) + if hf_quantizer is not None: hf_quantizer.validate_environment(device_map=device_map) @@ -5050,7 +5014,7 @@ def _load_pretrained_model( warner = logger.warning if model.__class__.__name__ in archs else logger.info warner( f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" - f" initializing {model.__class__.__name__}:\n\t{'\n\t'.join(update_key_name(unexpected_keys))}\n- This IS expected if you are" + f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" " with another architecture (e.g. initializing a BertForSequenceClassification model from a" " BertForPreTraining model).\n- This IS NOT expected if you are initializing" @@ -5061,8 +5025,8 @@ def _load_pretrained_model( logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") if len(missing_keys) > 0: logger.warning( - f"\nSome weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized:\n\t{'\n\t'.join(update_key_name(missing_keys))}\nYou should probably" + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" " TRAIN this model on a down-stream task to be able to use it for predictions and inference." ) elif len(mismatched_keys) == 0: @@ -5080,7 +5044,7 @@ def _load_pretrained_model( ] ) logger.warning( - f"\nSome weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able" " to use it for predictions and inference." diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index d1e50df10f1e..212969ce150a 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -48,7 +48,6 @@ is_torch_available, logging, ) -from transformers.modeling_utils import update_key_name from transformers.testing_utils import ( TOKEN, CaptureLogger, @@ -1719,45 +1718,6 @@ def test_isin_mps_friendly(self): torch.equal(torch.isin(random_ids, random_test_tensor), isin_mps_friendly(random_ids, random_test_tensor)) ) - def test_update_key_name(self): - original_keys = [ - "model.language_model.0.self_attn.0.mlp0", - "model.language_model.0.self_attn.1.mlp.conv1.weight", - "model.language_model.1.self_attn.0.mlp.conv1.weight", - "model.language_model.1.self_attn.1.mlp.conv1.weight", - "model.language_model.2.self_attn.0.mlp.conv1.weight", - "model.language_model.2.self_attn.1.conv1.weight", - "model.language_model.3.self_attn.0.mlp.conv1.weight", - "model.language_model.3.self_attn.1.mlp.conv1.weight", - "model.language_model.0.self_attn.0.mlp.layer.conv1.weight", - "model.language_model.1.self_attn.0.mlp.layer.conv1.weight", - "model.language_model.2.self_attn.0.mlp.layer.conv1.weight", - "model.language_model.3.self_attn.0.mlp.layer.conv1.weight", - "model.language_model.0.self_attn.0.mlp.layer.1.weight", - "model.language_model.1.self_attn.0.mlp.layer.1.weight", - "model.language_model.2.self_attn.0.mlp.layer.1.weight", - "model.language_model.3.self_attn.0.mlp.layer.1.weight", - "model.language_model.4.self_attn.0.mlp.layer.1.weight", - "model.language_model.5.self_attn.0.mlp.layer.1.weight", - "model.language_model.6.self_attn.0.mlp.layer.1.weight", - "model.language_model.7.self_attn.0.mlp.layer.1.weight", - "model.language_model.8.self_attn.0.mlp.layer.1.weight", - "model.language_model.9.self_attn.0.mlp.layer.1.weight", - "model.language_model.10.self_attn.0.mlp.layer.1.weight", - "model.language_model.11.self_attn.0.mlp.layer.1.weight", - "model.language_model.12.self_attn.0.mlp.layer.1.weight", - "model.language_model.13.self_attn.0.mlp.layer.1.weigh,t", - ] - new_keys = update_key_name(original_keys) - expected_new_keys = [ - "model.language_model.{0, ..., 13}.self_attn.{0, 1}.mlp.layer.1.weight", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.conv1.weight", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp0", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.conv1.weight", - "model.language_model.{0, 1, 2, 3}.self_attn.{0, 1}.mlp.layer.conv1.weight", - ] - self.assertListEqual(list(new_keys), expected_new_keys) - def test_can_generate(self): """Tests the behavior of `PreTrainedModel.can_generate` method.""" logger = logging.get_logger("transformers.modeling_utils") From a8fff20e99e3a1350423581ce954028f13a128f8 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 13:37:25 +0100 Subject: [PATCH 76/86] route all tokens to all experts when testing to avoid no gradient iddues --- tests/models/deepseek_v3/test_modeling_deepseek_v3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index c5efee3f9a6b..75fb83ee9b77 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -74,7 +74,7 @@ def __init__( qk_nope_head_dim=32, n_group=2, topk_group=1, - num_experts_per_tok=2, + num_experts_per_tok=8, first_k_dense_replace=2, norm_topk_prob=True, aux_loss_alpha=0.001, From 13019a7fc90fb7cfdc7f328950ef3655fe9776f4 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 13:51:30 +0100 Subject: [PATCH 77/86] finish fixing all tests --- .../models/deepseek_v3/modeling_deepseek_v3.py | 5 +++++ .../models/deepseek_v3/modular_deepseek_v3.py | 16 +++++++++++++++- .../deepseek_v3/test_modeling_deepseek_v3.py | 2 -- 3 files changed, 20 insertions(+), 3 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index 8001bcfff905..cab1e41cd7cd 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -154,6 +154,7 @@ def __init__(self, config): self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts))) + @torch.no_grad() def get_topk_indices(self, scores): scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( @@ -580,6 +581,10 @@ def _init_weights(self, module): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() + elif isinstance(module, DeepseekV3TopkRouter): + module.weight.data.normal_(mean=0.0, std=std) + elif isinstance(module, nn.Parameter): + module.weight.data.normal_(mean=0.0, std=std) DEEPSEEK_V3_INPUTS_DOCSTRING = r""" diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index f621f03aefc4..7713eb3b2707 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -112,6 +112,7 @@ def __init__(self, config): self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts))) + @torch.no_grad() def get_topk_indices(self, scores): scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0) group_scores = ( @@ -336,7 +337,20 @@ def __init__(self, config: DeepseekV3Config, layer_idx: int): class DeepseekV3PreTrainedModel(LlamaPreTrainedModel): - pass + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + 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_() + elif isinstance(module, DeepseekV3TopkRouter): + module.weight.data.normal_(mean=0.0, std=std) + elif isinstance(module, nn.Parameter): + module.weight.data.normal_(mean=0.0, std=std) class DeepseekV3Model(LlamaModel): diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index 75fb83ee9b77..d68e62abc336 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -579,8 +579,6 @@ def test_eager_matches_sdpa_generate(self): if "SdpaAttention" in submodule.__class__.__name__: has_sdpa = True break - if not has_sdpa: - raise ValueError("The SDPA model should have SDPA attention layers") texts = [ "hi here's a longer context, getting longer and", From 9b310a16e258656678278b0f7a3147f89d445aac Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 13:54:12 +0100 Subject: [PATCH 78/86] fixup --- tests/models/deepseek_v3/test_modeling_deepseek_v3.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index d68e62abc336..f452b028a20b 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -570,16 +570,6 @@ def test_eager_matches_sdpa_generate(self): self.assertTrue(model_eager.config._attn_implementation == "eager") - for name, submodule in model_eager.named_modules(): - if "SdpaAttention" in submodule.__class__.__name__: - raise ValueError("The eager model should not have SDPA attention layers") - - has_sdpa = False - for name, submodule in model_sdpa.named_modules(): - if "SdpaAttention" in submodule.__class__.__name__: - has_sdpa = True - break - texts = [ "hi here's a longer context, getting longer and", "Hello this is a very long sentence my friend, very long for real", From e3628a3ebf9c57502887af71d0da82c57d1935cb Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 14:08:54 +0100 Subject: [PATCH 79/86] nit --- src/transformers/modeling_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 8b6b2c831a4b..dc4f9b5eac82 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -779,8 +779,7 @@ def _load_state_dict_into_meta_model( device_map_regex = "|".join([re.escape(k) for k in sorted(device_map.keys(), reverse=True)]) is_quantized = hf_quantizer is not None - is_meta_state_dict = shard_file.endswith(".safetensors") # and not is_quantized - + is_meta_state_dict = shard_file.endswith(".safetensors") file_pointer = None if is_meta_state_dict: file_pointer = safe_open(shard_file, framework="pt", device=tensor_device) @@ -4102,7 +4101,6 @@ def from_pretrained( raise EnvironmentError("tensor parallel is only supported for `torch>=2.5`.") if not torch.distributed.is_initialized(): try: - logger.warning("Tensor Parallel requires torch.distributed to be initialized first.") rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) torch.distributed.init_process_group( From 9eb38e6d8c9c8ef797254f58c5adb5e18e2fe698 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 14:47:49 +0100 Subject: [PATCH 80/86] clean config --- .../models/deepseek_v3/configuration_deepseek_v3.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 1c609052977c..8f04f9a8e9dd 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -132,13 +132,7 @@ class DeepseekV3Config(PretrainedConfig): model_type = "deepseek_v3" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `DeepseekV3Model` - base_model_tp_plan = { - # "layers.*.self_attn.q_b_proj": "colwise", - # "layers.*.self_attn.q_a_proj": "colwise", - # "layers.*.self_attn.kv_a_proj_with_mqa": "colwise", - # "layers.*.self_attn.kv_b_proj": "colwise", - # "layers.*.self_attn.o_proj": "rowwise", + base_model_tp_plan = { # TODO: only replicate attention layers when > first_k_dense_replace "layers.*.mlp.experts.*.gate_proj": "local_colwise", "layers.*.mlp.experts.*.up_proj": "local_colwise", "layers.*.mlp.experts.*.down_proj": "local_rowwise", @@ -146,11 +140,11 @@ class DeepseekV3Config(PretrainedConfig): "layers.*.mlp.shared_experts.gate_proj": "local_colwise", "layers.*.mlp.shared_experts.up_proj": "local_colwise", "layers.*.mlp.shared_experts.down_proj": "local_rowwise", - "layers.*.mlp.shared_experts": "local", # local's job + "layers.*.mlp.shared_experts": "local", "layers.*.mlp.gate_proj": "local_colwise", "layers.*.mlp.up_proj": "local_colwise", "layers.*.mlp.down_proj": "local_rowwise", - "layers.*.mlp": "gather", + "layers.*.mlp": "gather", # This is the only moment where results are gathered } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), From 8cb959b8f931573d2d3e050b6722d87ae689de3d Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 15:00:40 +0100 Subject: [PATCH 81/86] last readme changes --- docs/source/en/model_doc/deepseek_v3.md | 95 +++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/docs/source/en/model_doc/deepseek_v3.md b/docs/source/en/model_doc/deepseek_v3.md index 1a300e2c00df..c3322a102f6e 100644 --- a/docs/source/en/model_doc/deepseek_v3.md +++ b/docs/source/en/model_doc/deepseek_v3.md @@ -60,6 +60,101 @@ outputs = model.generate(inputs, max_new_tokens=50) print(tokenizer.batch_decode(outputs)) print(time.time()-start) ``` +This generated: + +`````` +<｜Assistant｜> +Okay, the user wants to demonstrate how chat templating works. Let me break down what that means. Chat templating is about structuring the conversation data, especially for models that need specific input formats. Maybe they're referring to something like how messages are formatted with roles (user, assistant, system) in APIs like OpenAI. + +First, I should explain what chat templating is. It's the process of formatting conversation data into a structured format that the model can understand. This usually includes roles and content. For example, user messages, assistant responses, and system messages each have their own role tags. + +They might want an example. Let me think of a simple conversation. The user says "Hello, how are you?" and the assistant responds "I'm doing great. How can I help you today?" Then the user follows up with wanting to show off chat templating. So the example should include the history and the new message. + +In some frameworks, like Hugging Face's Transformers, chat templates are applied using Jinja2 templates. The template might look something like combining system messages, then looping through user and assistant messages with appropriate tags. For instance, using {% for message in messages %} and assigning roles like <|user|>, <|assistant|>, etc. + +I should structure the example with the messages array, showing each role and content. Then apply a hypothetical template to convert that into a formatted string the model uses. Also, mention that different models have different templating requirements, like using special tokens or varying role labels. + +Wait, the user mentioned "chat templating" in the context of showing off. Maybe they want a practical example they can present. So providing a code snippet or a structured data example would be helpful. Let me outline a typical messages array and then the templated output. + +Also, it's important to note that proper templating ensures the model knows the conversation flow, which is crucial for generating coherent responses. Maybe include a note about why it's important, like maintaining context and role-specific processing. + +Let me check if there are any common mistakes or things to avoid. For example, not closing tags properly, or mismatching roles. But maybe that's too detailed unless the user asks. Focus on the positive example first. + +Putting it all together, the response should have an example messages array, the applied template, and the final formatted string. Maybe use angle brackets or special tokens as placeholders. Also, mention that this helps in training or fine-tuning models with structured data. + +I think that's a solid approach. Let me structure it step by step to make it clear. + + +Chat templating is a way to structure conversation data (e.g., user/assistant interactions) into a format that language models understand. This is especially important for models trained to handle multi-turn dialogues, where the input must explicitly separate roles (user, assistant, system, etc.) and messages. Let’s break this down with an example! + +--- + +### **Step 1: Raw Conversation History** +Suppose we have this conversation: +- **User**: "Hello, how are you?" +- **Assistant**: "I'm doing great. How can I help you today?" +- **User**: "I'd like to show off how chat templating works!" + +--- + +### **Step 2: Structured Messages** +In frameworks like Hugging Face Transformers or OpenAI, conversations are often formatted as a list of dictionaries with `role` and `content`: +```python +messages = [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing great. How can I help you today?"}, + {"role": "user", "content": "I'd like to show off how chat templating works!"}, +] +``` + +--- + +### **Step 3: Apply a Chat Template** +A **chat template** converts this structured data into a single string formatted for the model. For example, using a Jinja-style template (common in Hugging Face): + +```jinja +{% for message in messages %} + {% if message['role'] == 'user' %} + <|user|>{{ message['content'] }}<|end|> + {% elif message['role'] == 'assistant' %} + <|assistant|>{{ message['content'] }}<|end|> + {% endif %} +{% endfor %} +<|assistant|> +``` + +--- + +### **Step 4: Final Templated Output** +Applying the template to our `messages` list would produce: +```text +<|user|>Hello, how are you?<|end|> +<|assistant|>I'm doing great. How can I help you today?<|end|> +<|user|>I'd like to show off how chat templating works!<|end|> +<|assistant|> +``` + +This tells the model: +1. The conversation history (user/assistant turns). +2. The model’s turn to generate a response (`<|assistant|>` at the end). + +--- + +### **Key Notes**: +- **Role Separation**: Tags like `<|user|>` and `<|assistant|>` help the model distinguish speakers. +- **Special Tokens**: Models often use unique tokens (e.g., `<|end|>`) to mark message boundaries. +- **Flexibility**: Templates vary by model (e.g., OpenAI uses `{"role": "user", "content": "..."}` instead of tags). + +--- + +### **Why This Matters**: +- **Consistency**: Ensures the model understands dialogue structure. +- **Context Preservation**: Maintains the flow of multi-turn conversations. +- **Alignment**: Matches the format the model was trained on for better performance. + +Want to dive deeper or see a specific framework’s implementation (e.g., OpenAI, Llama, Mistral)? Let me know! 😊<｜end▁of▁sentence｜> +`````` + Use the following to run it ```bash torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0|1 --rdzv-id an_id --rdzv-backend c10d --rdzv-endpoint master_addr:master_port run_deepseek_r1.py From a55630b6608f5d3f0f08ca6319f0574038a5e653 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 15:01:26 +0100 Subject: [PATCH 82/86] nit --- src/transformers/integrations/finegrained_fp8.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 19bd34cdbdbe..5583a88455ba 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -413,7 +413,6 @@ def replace_with_fp8_linear( if quantization_config.modules_to_not_convert is not None: modules_to_not_convert.extend(quantization_config.modules_to_not_convert) modules_to_not_convert = list(set(modules_to_not_convert)) - print(modules_to_not_convert) model, has_been_replaced = _replace_with_fp8_linear( model, tp_plan=model._tp_plan, From bce2073866b562c1f8d6a085c54069f3f7f19979 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 15:08:18 +0100 Subject: [PATCH 83/86] do cnit --- src/transformers/integrations/finegrained_fp8.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 5583a88455ba..342b7e108d73 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -382,10 +382,7 @@ def _replace_with_fp8_linear( block_size=quantization_config.weight_block_size, ) has_been_replaced = True - # import re - # TODO local should be determined here - # tp_plan[re.sub("\d+", "*", current_key_name_str)] = "local_" + tp_plan[re.sub("\d+", "*", current_key_name_str)] - # tp_plan[re.sub("\d+", "*", current_key_name_str).rsplit(".")[0]] = "gather" + # when changing a layer the TP PLAN for that layer should be updated. TODO if len(list(module.children())) > 0: _, has_been_replaced = _replace_with_fp8_linear( From a1f1f3fcb75d4780101e8569cde2a733a7ded8ca Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 15:12:00 +0100 Subject: [PATCH 84/86] typo --- src/transformers/integrations/finegrained_fp8.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 342b7e108d73..d4e472a990dc 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -382,7 +382,7 @@ def _replace_with_fp8_linear( block_size=quantization_config.weight_block_size, ) has_been_replaced = True - # when changing a layer the TP PLAN for that layer should be updated. TODO + # when changing a layer the TP PLAN for that layer should be updated. TODO if len(list(module.children())) > 0: _, has_been_replaced = _replace_with_fp8_linear( From d2ae07204a147f4d8dd8bd05145193797a155527 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 15:39:37 +0100 Subject: [PATCH 85/86] last nit --- src/transformers/modeling_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index dc4f9b5eac82..137bd01c01de 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4874,7 +4874,7 @@ def _load_pretrained_model( # Warmup cuda to load the weights much faster on devices if device_map is not None: # and hf_quantizer is None: expanded_device_map = expand_device_map(device_map, expected_keys) - caching_allocator_warmup(model_to_load, expanded_device_map) + caching_allocator_warmup(model_to_load, expanded_device_map, factor=2 if hf_quantizer is None else 4) error_msgs = [] mismatched_keys = [] @@ -5835,7 +5835,7 @@ def expand_device_map(device_map, param_names): return new_device_map -def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict): +def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict, factor=2): """This function warm-ups the caching allocator based on the size of the model tensors that will reside on each device. It allows to have one large call to Malloc, instead of recursively calling it later when loading the model, which is actually the loading speed botteneck. @@ -5886,7 +5886,7 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict): # Allow up to 95% of max device memory byte_count = min(byte_count, int(0.95 * device_memory)) # Allocate memory - _ = torch.empty(byte_count // 4, dtype=torch.float16, device=device, requires_grad=False) + _ = torch.empty(byte_count // factor, dtype=torch.float16, device=device, requires_grad=False) def get_disk_only_shard_files(device_map, weight_map): From 372efd65a39bd7304e2dd462c8dba4b63825e2d6 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Fri, 28 Mar 2025 15:46:33 +0100 Subject: [PATCH 86/86] one more one more --- src/transformers/integrations/fbgemm_fp8.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/integrations/fbgemm_fp8.py b/src/transformers/integrations/fbgemm_fp8.py index 98bc919caa15..71c2b570cc0a 100644 --- a/src/transformers/integrations/fbgemm_fp8.py +++ b/src/transformers/integrations/fbgemm_fp8.py @@ -28,7 +28,7 @@ logger = logging.get_logger(__name__) -class FbgemmFp8Linear(torch.nn.Linear): +class FbgemmFp8Linear(torch.nn.Module): def __init__(self, in_features, out_features, bias, weight_dtype=torch.float32): super().__init__() self.in_features = in_features