DEEPSEEK_R1_COMPLETE.txt · Lines 1501–3000

1501 1502 class DeepseekV3ForSequenceClassification(DeepseekV3PreTrainedModel) L1734 1503 __init__(config) L1735 1504 get_input_embeddings() / set_input_embeddings(value) L1744 L1747 1505 forward(...) -> SequenceClassifierOutputWithPast L1751 1506 1507-------------------------------------------------------------------------------- 150812.3 DeepseekV3Config PARAMETERS [SPEC] 1509-------------------------------------------------------------------------------- 1510 Constructor defaults, from configuration_deepseek.py L112. Values differing in 1511 this checkpoint's config.json are marked. 1512 1513 vocab_size 129280 1514 hidden_size 7168 1515 intermediate_size 18432 1516 moe_intermediate_size 2048 1517 num_hidden_layers 61 1518 num_nextn_predict_layers 1 1519 num_attention_heads 128 1520 num_key_value_heads 128 1521 n_shared_experts 1 1522 n_routed_experts 256 1523 ep_size 1 1524 routed_scaling_factor 2.5 1525 kv_lora_rank 512 1526 q_lora_rank 1536 1527 qk_rope_head_dim 64 1528 v_head_dim 128 1529 qk_nope_head_dim 128 1530 topk_method "noaux_tc" 1531 n_group 8 1532 topk_group 4 1533 num_experts_per_tok 8 1534 moe_layer_freq 1 1535 first_k_dense_replace 3 1536 norm_topk_prob True 1537 scoring_func "sigmoid" 1538 aux_loss_alpha 0.001 (training only) 1539 seq_aux True (training only) 1540 hidden_act "silu" 1541 max_position_embeddings 163840 (default in the class is 2048) 1542 initializer_range 0.02 1543 rms_norm_eps 1e-6 1544 use_cache True 1545 pad_token_id None 1546 bos_token_id 0 1547 eos_token_id 1 1548 pretraining_tp 1 1549 tie_word_embeddings False 1550 rope_theta 10000.0 1551 rope_scaling None (set to the yarn dict in config.json) 1552 attention_bias False 1553 attention_dropout 0.0 1554 1555-------------------------------------------------------------------------------- 155612.4 forward() INPUT / OUTPUT CONTRACT [SPEC] 1557-------------------------------------------------------------------------------- 1558 DeepseekV3ForCausalLM.forward( 1559 input_ids: LongTensor (batch, seq) or None 1560 attention_mask: Tensor (batch, seq) 1=keep 0=pad or None 1561 position_ids: LongTensor (batch, seq) or None 1562 past_key_values: Cache or legacy tuple or None 1563 inputs_embeds: FloatTensor (batch, seq, 7168) or None 1564 labels: LongTensor (batch, seq) or None 1565 use_cache: bool or None 1566 output_attentions: bool or None 1567 output_hidden_states: bool or None 1568 return_dict: bool or None 1569 ) -> CausalLMOutputWithPast( 1570 loss: FloatTensor scalar, only if labels given 1571 logits: FloatTensor (batch, seq, 129280), cast to float32 1572 past_key_values: Cache 1573 hidden_states: tuple of (batch, seq, 7168), only if requested 1574 attentions: tuple, only if requested 1575 ) 1576 1577 Exactly one of input_ids / inputs_embeds must be supplied. 1578 Labels are shifted internally: logits[..., :-1] predict labels[..., 1:]. 1579 Loss is CrossEntropyLoss over the flattened vocab. 1580 1581 prepare_inputs_for_generation handles cache trimming and builds position_ids 1582 from the attention mask via cumsum - which matters, because with left-padding 1583 the position ids are NOT simply arange. 1584 1585-------------------------------------------------------------------------------- 158612.5 OFFICIAL HOSTED API ENDPOINTS [CARD] + [EXTERNAL] 1587-------------------------------------------------------------------------------- 1588 Stated by the model card (section 5, reproduced verbatim in 15.6): 1589 1590 Chat UI https://chat.deepseek.com - enable the "DeepThink" toggle 1591 API platform https://platform.deepseek.com - "OpenAI-Compatible API" 1592 1593 The platform exposes an OpenAI-shaped surface. The canonical usage is to point 1594 an OpenAI client at DeepSeek's base URL and select the reasoning model: 1595 1596 from openai import OpenAI 1597 client = OpenAI(api_key="<key>", base_url="https://api.deepseek.com") 1598 resp = client.chat.completions.create( 1599 model="deepseek-reasoner", # R1 1600 messages=[{"role":"user","content":"..."}]) 1601 1602 Typical routes on that surface: 1603 POST /chat/completions chat, streaming or not 1604 POST /completions legacy text completion (FIM) 1605 GET /models model list 1606 GET /user/balance account balance 1607 1608 [EXTERNAL] Model identifiers, route paths, pricing and parameter support are 1609 DeepSeek's to change and are not verifiable from this machine. Treat the two 1610 URLs above as the authoritative pointers and confirm specifics against the live 1611 documentation at platform.deepseek.com before writing an integration. 1612 1613-------------------------------------------------------------------------------- 161412.6 SELF-HOSTED SERVING ENDPOINTS [CARD] + [EXTERNAL] 1615-------------------------------------------------------------------------------- 1616 The model card directs full-R1 users to the DeepSeek-V3 repository: 1617 1618 https://github.com/deepseek-ai/DeepSeek-V3 1619 1620 and gives concrete launch lines for the DISTILLED models (which are ordinary 1621 Qwen/Llama architectures and therefore already supported): 1622 1623 vllm serve deepseek-ai/DeepSeek-R1-Distill-Qwen-32B \ 1624 --tensor-parallel-size 2 --max-model-len 32768 --enforce-eager 1625 1626 python3 -m sglang.launch_server \ 1627 --model deepseek-ai/DeepSeek-R1-Distill-Qwen-32B \ 1628 --trust-remote-code --tp 2 1629 1630 Both servers present an OpenAI-compatible HTTP surface, conventionally: 1631 1632 POST /v1/chat/completions chat, SSE streaming with stream=true 1633 POST /v1/completions raw completion 1634 GET /v1/models model list 1635 GET /health liveness 1636 1637 vLLM additionally serves /metrics (Prometheus); SGLang additionally serves 1638 /generate (its native, non-OpenAI route). 1639 1640 [EXTERNAL] Flags and routes belong to those projects and change between 1641 releases; verify against the version you install. 1642 1643 ON THIS MACHINE: none of these are installed, and the 643 GB fp8 checkpoint 1644 needs a runtime that implements 128x128 block dequant. See PART 13 for why 1645 the memory envelope is the binding constraint here regardless. 1646 1647-------------------------------------------------------------------------------- 164812.7 THE REASONING-CONTENT CONTRACT 1649-------------------------------------------------------------------------------- 1650 R1 emits a reasoning trace before its answer, delimited by <think> (128798) and 1651 </think> (128799). Two consequences that any integration must handle: 1652 1653 1. SEPARATE THE CHANNELS. The trace is not the answer. On the hosted API this 1654 arrives as a distinct field alongside the normal content; when self-hosting 1655 you must split on </think> yourself. 1656 1657 2. DO NOT FEED TRACES BACK. The chat template explicitly strips everything up 1658 to and including </think> from prior assistant turns (see 11.3). Echoing old 1659 reasoning into the next request diverges from how the model was trained. 1660 1661 USAGE RECOMMENDATIONS FROM THE MODEL CARD [CARD]: 1662 - temperature 0.5-0.7, 0.6 recommended, to avoid endless repetition 1663 - AVOID a system prompt; put all instructions in the user message 1664 - for math, ask for the final answer inside \boxed{} 1665 - average over several runs when benchmarking 1666 - force the response to begin with "<think>\n"; the model otherwise 1667 sometimes emits an empty "<think>\n\n</think>" and skips reasoning 1668 - evaluation used max generation length 32,768 tokens, temperature 0.6, 1669 top_p 0.95, 64 samples per query for pass@1 1670 1671 1672================================================================================ 1673PART 13 FOOTPRINT: MEMORY, DISK I/O, COMPUTE [COMPUTED] 1674================================================================================ 1675Derived from the measured tensor census. Batch 1, decode step, layers 0..60 1676(the MTP layer excluded). 1677 1678WEIGHTS THAT MUST BE RESIDENT (everything except the routed experts) 1679 --------------------------------------------------------------------------- 1680 component GiB basis 1681 --------------------------------------------------------------------------- 1682 attention, 5 projections x 61 layers 10.630 187,105,280 B per layer 1683 shared experts x 58 layers 2.379 3 x 2048 x 7168 fp8 1684 dense MLP, layers 0-2 1.107 3 x 3 x 18432 x 7168 fp8 1685 MoE router weights x 58 0.198 256 x 7168 BF16 1686 e_score_correction_bias + all norms 0.002 F32 / BF16 1687 embed_tokens 1.726 129280 x 7168 BF16 1688 lm_head 1.726 129280 x 7168 BF16 1689 --------------------------------------------------------------------------- 1690 TOTAL RESIDENT, native dtypes 17.768 1691 with embed_tokens + lm_head at int8 16.042 1692 --------------------------------------------------------------------------- 1693 1694ROUTED EXPERT TRAFFIC (streamed, per token) 1695 8 experts x 58 MoE layers = 464 expert reads per token 1696 464 x 44,050,944 B = 20,439,638,016 B = 19.04 GiB per token 1697 1698 Against the measured drive rates from PART 1: 1699 --------------------------------------------------------------------------- 1700 at 507 MB/s single stream 40.3 s per token 1701 at 1,220 MB/s 4 parallel streams 16.8 s per token 1702 at 326 MB/s scattered small reads 62.7 s per token 1703 --------------------------------------------------------------------------- 1704 Bandwidth floors only: no syscall overhead, no dequantization, no compute 1705 overlap, no cache reuse. Nothing has been run end to end. 1706 1707KV CACHE - see 10.3 for the table and the kv_b_proj coupling. 1708 1709COMPUTE PER TOKEN 1710 --------------------------------------------------------------------------- 1711 attention (5 projections x 61) 11.41 GMAC 31.3% 1712 routed experts (8 x 58 x 3 matrices) 20.43 GMAC 56.0% 1713 shared experts (58 x 3 matrices) 2.55 GMAC 7.0% 1714 dense MLP (3 layers) 1.19 GMAC 3.3% 1715 lm_head 0.93 GMAC 2.5% 1716 --------------------------------------------------------------------------- 1717 TOTAL 36.52 GMAC = 73.0 GFLOP per token 1718 --------------------------------------------------------------------------- 1719 On 12 threads of AVX2 without VNNI, roughly 150 GOPS sustained int8 is 1720 realistic for this CPU: about 0.24 s per token of arithmetic. 1721 1722THE SHAPE OF THE PROBLEM ON THIS MACHINE 1723 disk I/O per token 17 - 40 s 1724 arithmetic per token ~0.24 s 1725 ratio 70x to 170x 1726 I/O-bound by two orders of magnitude. Attention is 31% of the arithmetic but 1727 none of the streamed I/O, since it is resident. 1728 1729WHAT FITS, AND WHAT DOES NOT 1730 RAM available to WSL2 25 GiB 1731 free at scan time 9.6 GiB 1732 resident weights needed 16.0 GiB 1733 + KV at 16K context, bf16 1.1 GiB 1734 + activations, LUTs, OS reserve ~1.5 GiB 1735 -------- 1736 minimum working set ~18.6 GiB 1737 1738 The resident weight set alone exceeds currently free memory. The 163,840-token 1739 context is unreachable here: 10.72 GiB of KV on top of 16.0 GiB of weights 1740 exceeds the 25 GiB the VM has. 16K-32K is the honest ceiling for this machine. 1741 1742 Disk: 619.65 GiB of the checkpoint is routed experts against 228 GiB free on /. 1743 There is no room to write a second, reorganized copy of the expert weights. 1744 The contiguous on-disk layout in 2.6 means none is needed. 1745 1746 1747================================================================================ 1748PART 14 DISCREPANCIES, TRAPS AND GOTCHAS 1749================================================================================ 1750 1751THREE-WAY CONTEXT-LENGTH DISAGREEMENT 1752 config.json max_position_embeddings 163840 1753 README.md "Context Length" 128K [CARD] 1754 tokenizer_config.json model_max_length 16384 1755 Three different numbers in three files of the same repository. The architecture 1756 supports 163,840 positions; the card advertises 128K; the tokenizer will refuse 1757 or warn past 16,384. Decide explicitly which governs your deployment; do not 1758 let a library pick for you. 1759 1760INDEX total_size IS WRONG BY 1.99x 1761 1,369,062,772,000 claimed vs 688,574,839,360 actual. FP8 elements were costed 1762 at two bytes. Headers are authoritative. 1763 1764e_score_correction_bias IS F32 1765 Shape (256,), one per MoE layer, sitting among BF16 neighbours in the same 1766 layer namespace. Read it as bf16 and routing degrades in every MoE layer with 1767 no crash and no error - only worse output. 1768 1769kv_a_proj_with_mqa HAS THE ONLY PADDED BLOCK GRID 1770 Shape (576, 7168) against a scale grid of (5, 56). 576/128 = 4.5, so the fifth 1771 block row covers only 64 of 128 rows. Every other tensor among all 91,991 1772 divides by 128 exactly. A dequant loop assuming exact divisibility reads 64 1773 rows of garbage per layer across 62 layers, silently. 1774 1775ROUTER GATHERS PRE-BIAS SCORES 1776 Selection uses scores + bias; weighting uses scores without bias. Both, in the 1777 same function, deliberately. See PART 7 step 7. 1778 1779softmax_scale IS NOT 1/sqrt(d) 1780 0.135237, not 0.0721688, because mscale_all_dim=1.0 activates the mscale^2 1781 correction. A 1.874x factor. 1782 1783routed_scaling_factor IS APPLIED AFTER RENORMALIZATION 1784 Multiply by 2.5 after dividing by the sum, not before. Easy to drop entirely. 1785 1786THE SHARED EXPERT CONSUMES THE PRE-MoE RESIDUAL 1787 shared_experts(identity), not shared_experts(routed_output), and it is added 1788 to the routed sum afterwards. 1789 1790HALF OF ALL TENSORS ARE TINY F32 SCALES 1791 45,867 of 91,991. Never size work, progress, or allocation by tensor count. 1792 1793FP8 CANNOT BE INTEGER-ACCUMULATED 1794 It has an exponent field. Raw fp8 bytes are not fixed-point; any integer MAC 1795 scheme has to dequantize or work per-128x128-tile with a shared exponent. 1796 1797EXPERTS STRADDLE SHARD BOUNDARIES 1798 Contiguity holds WITHIN a shard. About 162 of 14,848 experts span two files. 1799 Handle it, or corrupt one expert per boundary. 1800 1801lm_head IS IN SHARD 160, NOT 163 1802 Shards 161-163 are entirely MTP layer 61. A loader that assumes "the output 1803 head is in the last shard" will not find it. 1804 1805465 UNUSED VOCAB SLOTS 1806 Token ids run 0..128,814 but vocab_size is 129,280. Rows 128,815..129,279 of 1807 embed_tokens and lm_head are never indexed by the tokenizer, but must still be 1808 allocated. 1809 1810NO FP8 REFERENCE IMPLEMENTATION IN THE REPOSITORY 1811 The shipped modeling code assumes already-dequantized bf16. Anyone writing a 1812 loader is writing the 128x128 block dequant without a local reference to check 1813 against. See PART 6. 1814 1815MACHINE-SPECIFIC (megamind-agi) 1816 * No AVX-512, no VNNI, no AMX. int8 dot = vpmaddubsw + vpmaddwd on 256-bit. 1817 * Warm re-reads are SLOWER than cold (408 vs 507 MB/s). No page-cache benefit. 1818 * 4-way parallel reads give 2.4x - fan out, do not serialize. 1819 * Python 3.14 has no torch wheels. 1820 * WSL2 kills backgrounded jobs when a one-shot SSH command returns; use systemd. 1821 * cmd.exe caps command lines at 8191 chars; base64 anything long. 1822 1823 1824================================================================================ 1825PART 15 SOURCE CODE 1826================================================================================ 1827Everything below is reproduced byte-for-byte from the files in 1828/home/joseph/DeepSeek-R1 as of 2026-08-05. 1829 1830 1831================================================================================ 183215.1 config.json 1833================================================================================ 1834{ 1835 "architectures": [ 1836 "DeepseekV3ForCausalLM" 1837 ], 1838 "attention_bias": false, 1839 "attention_dropout": 0.0, 1840 "auto_map": { 1841 "AutoConfig": "configuration_deepseek.DeepseekV3Config", 1842 "AutoModel": "modeling_deepseek.DeepseekV3Model", 1843 "AutoModelForCausalLM": "modeling_deepseek.DeepseekV3ForCausalLM" 1844 }, 1845 "bos_token_id": 0, 1846 "eos_token_id": 1, 1847 "ep_size": 1, 1848 "first_k_dense_replace": 3, 1849 "hidden_act": "silu", 1850 "hidden_size": 7168, 1851 "initializer_range": 0.02, 1852 "intermediate_size": 18432, 1853 "kv_lora_rank": 512, 1854 "max_position_embeddings": 163840, 1855 "model_type": "deepseek_v3", 1856 "moe_intermediate_size": 2048, 1857 "moe_layer_freq": 1, 1858 "n_group": 8, 1859 "n_routed_experts": 256, 1860 "n_shared_experts": 1, 1861 "norm_topk_prob": true, 1862 "num_attention_heads": 128, 1863 "num_experts_per_tok": 8, 1864 "num_hidden_layers": 61, 1865 "num_key_value_heads": 128, 1866 "num_nextn_predict_layers": 1, 1867 "q_lora_rank": 1536, 1868 "qk_nope_head_dim": 128, 1869 "qk_rope_head_dim": 64, 1870 "quantization_config": { 1871 "activation_scheme": "dynamic", 1872 "fmt": "e4m3", 1873 "quant_method": "fp8", 1874 "weight_block_size": [ 1875 128, 1876 128 1877 ] 1878 }, 1879 "rms_norm_eps": 1e-06, 1880 "rope_scaling": { 1881 "beta_fast": 32, 1882 "beta_slow": 1, 1883 "factor": 40, 1884 "mscale": 1.0, 1885 "mscale_all_dim": 1.0, 1886 "original_max_position_embeddings": 4096, 1887 "type": "yarn" 1888 }, 1889 "rope_theta": 10000, 1890 "routed_scaling_factor": 2.5, 1891 "scoring_func": "sigmoid", 1892 "tie_word_embeddings": false, 1893 "topk_group": 4, 1894 "topk_method": "noaux_tc", 1895 "torch_dtype": "bfloat16", 1896 "transformers_version": "4.46.3", 1897 "use_cache": true, 1898 "v_head_dim": 128, 1899 "vocab_size": 129280 1900} 1901 1902 1903================================================================================ 190415.2 configuration_deepseek.py (199 lines, 9,897 bytes) 1905================================================================================ 1906from transformers.configuration_utils import PretrainedConfig 1907from transformers.utils import logging 1908 1909logger = logging.get_logger(__name__) 1910 1911DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} 1912class DeepseekV3Config(PretrainedConfig): 1913 r""" 1914 This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek 1915 model according to the specified arguments, defining the model architecture. Instantiating a configuration with the 1916 defaults will yield a similar configuration to that of the DeepSeek-V3. 1917 1918 Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the 1919 documentation from [`PretrainedConfig`] for more information. 1920 1921 1922 Args: 1923 vocab_size (`int`, *optional*, defaults to 129280): 1924 Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the 1925 `inputs_ids` passed when calling [`DeepseekV3Model`] 1926 hidden_size (`int`, *optional*, defaults to 4096): 1927 Dimension of the hidden representations. 1928 intermediate_size (`int`, *optional*, defaults to 11008): 1929 Dimension of the MLP representations. 1930 moe_intermediate_size (`int`, *optional*, defaults to 1407): 1931 Dimension of the MoE representations. 1932 num_hidden_layers (`int`, *optional*, defaults to 32): 1933 Number of hidden layers in the Transformer decoder. 1934 num_nextn_predict_layers (`int`, *optional*, defaults to 1): 1935 Number of nextn predict layers in the DeepSeekV3 Model. 1936 num_attention_heads (`int`, *optional*, defaults to 32): 1937 Number of attention heads for each attention layer in the Transformer decoder. 1938 n_shared_experts (`int`, *optional*, defaults to None): 1939 Number of shared experts, None means dense model. 1940 n_routed_experts (`int`, *optional*, defaults to None): 1941 Number of routed experts, None means dense model. 1942 routed_scaling_factor (`float`, *optional*, defaults to 1.0): 1943 Scaling factor or routed experts. 1944 topk_method (`str`, *optional*, defaults to `gready`): 1945 Topk method used in routed gate. 1946 n_group (`int`, *optional*, defaults to None): 1947 Number of groups for routed experts. 1948 topk_group (`int`, *optional*, defaults to None): 1949 Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). 1950 num_experts_per_tok (`int`, *optional*, defaults to None): 1951 Number of selected experts, None means dense model. 1952 moe_layer_freq (`int`, *optional*, defaults to 1): 1953 The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. 1954 first_k_dense_replace (`int`, *optional*, defaults to 0): 1955 Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). 1956 \--k dense layers--/ 1957 norm_topk_prob (`bool`, *optional*, defaults to False): 1958 Whether to normalize the weights of the routed experts. 1959 scoring_func (`str`, *optional*, defaults to 'softmax'): 1960 Method of computing expert weights. 1961 aux_loss_alpha (`float`, *optional*, defaults to 0.001): 1962 Auxiliary loss weight coefficient. 1963 seq_aux = (`bool`, *optional*, defaults to True): 1964 Whether to compute the auxiliary loss for each individual sample. 1965 num_key_value_heads (`int`, *optional*): 1966 This is the number of key_value heads that should be used to implement Grouped Query Attention. If 1967 `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if 1968 `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When 1969 converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed 1970 by meanpooling all the original heads within that group. For more details checkout [this 1971 paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to 1972 `num_attention_heads`. 1973 hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): 1974 The non-linear activation function (function or string) in the decoder. 1975 max_position_embeddings (`int`, *optional*, defaults to 2048): 1976 The maximum sequence length that this model might ever be used with. 1977 initializer_range (`float`, *optional*, defaults to 0.02): 1978 The standard deviation of the truncated_normal_initializer for initializing all weight matrices. 1979 rms_norm_eps (`float`, *optional*, defaults to 1e-06): 1980 The epsilon used by the rms normalization layers. 1981 use_cache (`bool`, *optional*, defaults to `True`): 1982 Whether or not the model should return the last key/values attentions (not used by all models). Only 1983 relevant if `config.is_decoder=True`. 1984 pad_token_id (`int`, *optional*): 1985 Padding token id. 1986 bos_token_id (`int`, *optional*, defaults to 1): 1987 Beginning of stream token id. 1988 eos_token_id (`int`, *optional*, defaults to 2): 1989 End of stream token id. 1990 tie_word_embeddings (`bool`, *optional*, defaults to `False`): 1991 Whether to tie weight embeddings 1992 rope_theta (`float`, *optional*, defaults to 10000.0): 1993 The base period of the RoPE embeddings. 1994 rope_scaling (`Dict`, *optional*): 1995 Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling 1996 strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is 1997 `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update 1998 `max_position_embeddings` to the expected new maximum. 1999 attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): 2000 Whether to use a bias in the query, key, value and output projection layers during self-attention. 2001 attention_dropout (`float`, *optional*, defaults to 0.0): 2002 The dropout ratio for the attention probabilities. 2003 2004 ```python 2005 >>> from transformers import DeepseekV3Model, DeepseekV3Config 2006 2007 >>> # Initializing a Deepseek-V3 style configuration 2008 >>> configuration = DeepseekV3Config() 2009 2010 >>> # Accessing the model configuration 2011 >>> configuration = model.config 2012 ```""" 2013 2014 model_type = "deepseek_v3" 2015 keys_to_ignore_at_inference = ["past_key_values"] 2016 2017 def __init__( 2018 self, 2019 vocab_size=129280, 2020 hidden_size=7168, 2021 intermediate_size=18432, 2022 moe_intermediate_size = 2048, 2023 num_hidden_layers=61, 2024 num_nextn_predict_layers=1, 2025 num_attention_heads=128, 2026 num_key_value_heads=128, 2027 n_shared_experts = 1, 2028 n_routed_experts = 256, 2029 ep_size = 1, 2030 routed_scaling_factor = 2.5, 2031 kv_lora_rank = 512, 2032 q_lora_rank = 1536, 2033 qk_rope_head_dim = 64, 2034 v_head_dim = 128, 2035 qk_nope_head_dim = 128, 2036 topk_method = 'noaux_tc', 2037 n_group = 8, 2038 topk_group = 4, 2039 num_experts_per_tok = 8, 2040 moe_layer_freq = 1, 2041 first_k_dense_replace = 3, 2042 norm_topk_prob = True, 2043 scoring_func = 'sigmoid', 2044 hidden_act="silu", 2045 max_position_embeddings=4096, 2046 initializer_range=0.02, 2047 rms_norm_eps=1e-6, 2048 use_cache=True, 2049 pad_token_id=None, 2050 bos_token_id=0, 2051 eos_token_id=1, 2052 tie_word_embeddings=False, 2053 rope_theta=10000.0, 2054 rope_scaling=None, 2055 attention_bias=False, 2056 attention_dropout=0.0, 2057 **kwargs, 2058 ): 2059 self.vocab_size = vocab_size 2060 self.max_position_embeddings = max_position_embeddings 2061 self.hidden_size = hidden_size 2062 self.intermediate_size = intermediate_size 2063 self.moe_intermediate_size = moe_intermediate_size 2064 self.num_hidden_layers = num_hidden_layers 2065 self.num_nextn_predict_layers = num_nextn_predict_layers 2066 self.num_attention_heads = num_attention_heads 2067 self.n_shared_experts = n_shared_experts 2068 self.n_routed_experts = n_routed_experts 2069 self.ep_size = ep_size 2070 self.routed_scaling_factor = routed_scaling_factor 2071 self.kv_lora_rank = kv_lora_rank 2072 self.q_lora_rank = q_lora_rank 2073 self.qk_rope_head_dim = qk_rope_head_dim 2074 self.v_head_dim = v_head_dim 2075 self.qk_nope_head_dim = qk_nope_head_dim 2076 self.topk_method = topk_method 2077 self.n_group = n_group 2078 self.topk_group = topk_group 2079 self.num_experts_per_tok = num_experts_per_tok 2080 self.moe_layer_freq = moe_layer_freq 2081 self.first_k_dense_replace = first_k_dense_replace 2082 self.norm_topk_prob = norm_topk_prob 2083 self.scoring_func = scoring_func 2084 # for backward compatibility 2085 if num_key_value_heads is None: 2086 num_key_value_heads = num_attention_heads 2087 2088 self.num_key_value_heads = num_key_value_heads 2089 self.hidden_act = hidden_act 2090 self.initializer_range = initializer_range 2091 self.rms_norm_eps = rms_norm_eps 2092 self.use_cache = use_cache 2093 self.rope_theta = rope_theta 2094 self.rope_scaling = rope_scaling 2095 self.attention_bias = attention_bias 2096 self.attention_dropout = attention_dropout 2097 2098 super().__init__( 2099 pad_token_id=pad_token_id, 2100 bos_token_id=bos_token_id, 2101 eos_token_id=eos_token_id, 2102 tie_word_embeddings=tie_word_embeddings, 2103 **kwargs, 2104 ) 2105 2106================================================================================ 210715.3 modeling_deepseek.py (1,848 lines, 75,741 bytes) 2108 The reference forward pass. This is the authoritative description of 2109 how DeepSeek-R1 computes. 2110================================================================================ 2111# coding=utf-8 2112# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved. 2113# 2114# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX 2115# and OPT implementations in this library. It has been modified from its 2116# original forms to accommodate minor architectural differences compared 2117# to GPT-NeoX and OPT used by the Meta AI team that trained the model. 2118# 2119# Licensed under the Apache License, Version 2.0 (the "License"); 2120# you may not use this file except in compliance with the License. 2121# You may obtain a copy of the License at 2122# 2123# http://www.apache.org/licenses/LICENSE-2.0 2124# 2125# Unless required by applicable law or agreed to in writing, software 2126# distributed under the License is distributed on an "AS IS" BASIS, 2127# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 2128# See the License for the specific language governing permissions and 2129# limitations under the License. 2130""" PyTorch DeepSeek model.""" 2131import math 2132import warnings 2133from typing import List, Optional, Tuple, Union 2134 2135import torch 2136import torch.nn.functional as F 2137import torch.utils.checkpoint 2138from torch import nn 2139from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss 2140 2141from transformers.activations import ACT2FN 2142from transformers.cache_utils import Cache, DynamicCache 2143from transformers.modeling_attn_mask_utils import ( 2144 AttentionMaskConverter, 2145 _prepare_4d_attention_mask, 2146 _prepare_4d_causal_attention_mask, 2147) 2148from transformers.modeling_outputs import ( 2149 BaseModelOutputWithPast, 2150 CausalLMOutputWithPast, 2151 SequenceClassifierOutputWithPast, 2152) 2153from transformers.modeling_utils import PreTrainedModel 2154from transformers.pytorch_utils import ( 2155 ALL_LAYERNORM_LAYERS, 2156 is_torch_greater_or_equal_than_1_13, 2157) 2158from transformers.utils import ( 2159 add_start_docstrings, 2160 add_start_docstrings_to_model_forward, 2161 is_flash_attn_2_available, 2162 is_flash_attn_greater_or_equal_2_10, 2163 logging, 2164 replace_return_docstrings, 2165) 2166from transformers.utils.import_utils import is_torch_fx_available 2167from .configuration_deepseek import DeepseekV3Config 2168import torch.distributed as dist 2169import numpy as np 2170 2171if is_flash_attn_2_available(): 2172 from flash_attn import flash_attn_func, flash_attn_varlen_func 2173 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa 2174 2175 2176# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph. 2177# It means that the function will not be traced through and simply appear as a node in the graph. 2178if is_torch_fx_available(): 2179 if not is_torch_greater_or_equal_than_1_13: 2180 import torch.fx 2181 2182 _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask) 2183 2184 2185logger = logging.get_logger(__name__) 2186 2187_CONFIG_FOR_DOC = "DeepseekV3Config" 2188 2189 2190def _get_unpad_data(attention_mask): 2191 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) 2192 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() 2193 max_seqlen_in_batch = seqlens_in_batch.max().item() 2194 cu_seqlens = F.pad( 2195 torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0) 2196 ) 2197 return ( 2198 indices, 2199 cu_seqlens, 2200 max_seqlen_in_batch, 2201 ) 2202 2203 2204class DeepseekV3RMSNorm(nn.Module): 2205 def __init__(self, hidden_size, eps=1e-6): 2206 """ 2207 DeepseekV3RMSNorm is equivalent to T5LayerNorm 2208 """ 2209 super().__init__() 2210 self.weight = nn.Parameter(torch.ones(hidden_size)) 2211 self.variance_epsilon = eps 2212 2213 def forward(self, hidden_states): 2214 input_dtype = hidden_states.dtype 2215 hidden_states = hidden_states.to(torch.float32) 2216 variance = hidden_states.pow(2).mean(-1, keepdim=True) 2217 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) 2218 return self.weight * hidden_states.to(input_dtype) 2219 2220 2221ALL_LAYERNORM_LAYERS.append(DeepseekV3RMSNorm) 2222 2223 2224class DeepseekV3RotaryEmbedding(nn.Module): 2225 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): 2226 super().__init__() 2227 2228 self.dim = dim 2229 self.max_position_embeddings = max_position_embeddings 2230 self.base = base 2231 inv_freq = 1.0 / ( 2232 self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim) 2233 ) 2234 self.register_buffer("inv_freq", inv_freq, persistent=False) 2235 2236 # Build here to make `torch.jit.trace` work. 2237 self._set_cos_sin_cache( 2238 seq_len=max_position_embeddings, 2239 device=self.inv_freq.device, 2240 dtype=torch.get_default_dtype(), 2241 ) 2242 self.max_seq_len_cached = None 2243 2244 def _set_cos_sin_cache(self, seq_len, device, dtype): 2245 self.max_seq_len_cached = seq_len 2246 t = torch.arange( 2247 self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype 2248 ) 2249 2250 freqs = torch.outer(t, self.inv_freq.to(t.device)) 2251 # Different from paper, but it uses a different permutation in order to obtain the same calculation 2252 emb = torch.cat((freqs, freqs), dim=-1) 2253 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) 2254 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) 2255 2256 def forward(self, x, seq_len=None): 2257 # x: [bs, num_attention_heads, seq_len, head_size] 2258 if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached: 2259 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) 2260 2261 return ( 2262 self.cos_cached[:seq_len].to(dtype=x.dtype), 2263 self.sin_cached[:seq_len].to(dtype=x.dtype), 2264 ) 2265 2266 2267# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->DeepseekV3 2268class DeepseekV3LinearScalingRotaryEmbedding(DeepseekV3RotaryEmbedding): 2269 """DeepseekV3RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" 2270 2271 def __init__( 2272 self, 2273 dim, 2274 max_position_embeddings=2048, 2275 base=10000, 2276 device=None, 2277 scaling_factor=1.0, 2278 ): 2279 self.scaling_factor = scaling_factor 2280 super().__init__(dim, max_position_embeddings, base, device) 2281 2282 def _set_cos_sin_cache(self, seq_len, device, dtype): 2283 self.max_seq_len_cached = seq_len 2284 t = torch.arange( 2285 self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype 2286 ) 2287 t = t / self.scaling_factor 2288 2289 freqs = torch.outer(t, self.inv_freq) 2290 # Different from paper, but it uses a different permutation in order to obtain the same calculation 2291 emb = torch.cat((freqs, freqs), dim=-1) 2292 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) 2293 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) 2294 2295 2296# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->DeepseekV3 2297class DeepseekV3DynamicNTKScalingRotaryEmbedding(DeepseekV3RotaryEmbedding): 2298 """DeepseekV3RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" 2299 2300 def __init__( 2301 self, 2302 dim, 2303 max_position_embeddings=2048, 2304 base=10000, 2305 device=None, 2306 scaling_factor=1.0, 2307 ): 2308 self.scaling_factor = scaling_factor 2309 super().__init__(dim, max_position_embeddings, base, device) 2310 2311 def _set_cos_sin_cache(self, seq_len, device, dtype): 2312 self.max_seq_len_cached = seq_len 2313 2314 if seq_len > self.max_position_embeddings: 2315 base = self.base * ( 2316 (self.scaling_factor * seq_len / self.max_position_embeddings) 2317 - (self.scaling_factor - 1) 2318 ) ** (self.dim / (self.dim - 2)) 2319 inv_freq = 1.0 / ( 2320 base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim) 2321 ) 2322 self.register_buffer("inv_freq", inv_freq, persistent=False) 2323 2324 t = torch.arange( 2325 self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype 2326 ) 2327 2328 freqs = torch.outer(t, self.inv_freq) 2329 # Different from paper, but it uses a different permutation in order to obtain the same calculation 2330 emb = torch.cat((freqs, freqs), dim=-1) 2331 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) 2332 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) 2333 2334 2335# Inverse dim formula to find dim based on number of rotations 2336def yarn_find_correction_dim( 2337 num_rotations, dim, base=10000, max_position_embeddings=2048 2338): 2339 return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( 2340 2 * math.log(base) 2341 ) 2342 2343 2344# Find dim range bounds based on rotations 2345def yarn_find_correction_range( 2346 low_rot, high_rot, dim, base=10000, max_position_embeddings=2048 2347): 2348 low = math.floor( 2349 yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) 2350 ) 2351 high = math.ceil( 2352 yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) 2353 ) 2354 return max(low, 0), min(high, dim - 1) # Clamp values just in case 2355 2356 2357def yarn_get_mscale(scale=1, mscale=1): 2358 if scale <= 1: 2359 return 1.0 2360 return 0.1 * mscale * math.log(scale) + 1.0 2361 2362 2363def yarn_linear_ramp_mask(min, max, dim): 2364 if min == max: 2365 max += 0.001 # Prevent singularity 2366 2367 linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) 2368 ramp_func = torch.clamp(linear_func, 0, 1) 2369 return ramp_func 2370 2371 2372class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding): 2373 2374 def __init__( 2375 self, 2376 dim, 2377 max_position_embeddings=2048, 2378 base=10000, 2379 device=None, 2380 scaling_factor=1.0, 2381 original_max_position_embeddings=4096, 2382 beta_fast=32, 2383 beta_slow=1, 2384 mscale=1, 2385 mscale_all_dim=0, 2386 ): 2387 self.scaling_factor = scaling_factor 2388 self.original_max_position_embeddings = original_max_position_embeddings 2389 self.beta_fast = beta_fast 2390 self.beta_slow = beta_slow 2391 self.mscale = mscale 2392 self.mscale_all_dim = mscale_all_dim 2393 super().__init__(dim, max_position_embeddings, base, device) 2394 2395 def _set_cos_sin_cache(self, seq_len, device, dtype): 2396 self.max_seq_len_cached = seq_len 2397 dim = self.dim 2398 2399 freq_extra = 1.0 / ( 2400 self.base 2401 ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) 2402 ) 2403 freq_inter = 1.0 / ( 2404 self.scaling_factor 2405 * self.base 2406 ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) 2407 ) 2408 2409 low, high = yarn_find_correction_range( 2410 self.beta_fast, 2411 self.beta_slow, 2412 dim, 2413 self.base, 2414 self.original_max_position_embeddings, 2415 ) 2416 inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to( 2417 device=device, dtype=torch.float32 2418 ) 2419 inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask 2420 self.register_buffer("inv_freq", inv_freq, persistent=False) 2421 2422 t = torch.arange(seq_len, device=device, dtype=torch.float32) 2423 2424 freqs = torch.outer(t, inv_freq) 2425 2426 _mscale = float( 2427 yarn_get_mscale(self.scaling_factor, self.mscale) 2428 / yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) 2429 ) 2430 2431 emb = torch.cat((freqs, freqs), dim=-1) 2432 self.register_buffer( 2433 "cos_cached", (emb.cos() * _mscale).to(dtype), persistent=False 2434 ) 2435 self.register_buffer( 2436 "sin_cached", (emb.sin() * _mscale).to(dtype), persistent=False 2437 ) 2438 2439 2440# Copied from transformers.models.llama.modeling_llama.rotate_half 2441def rotate_half(x): 2442 """Rotates half the hidden dims of the input.""" 2443 x1 = x[..., : x.shape[-1] // 2] 2444 x2 = x[..., x.shape[-1] // 2 :] 2445 return torch.cat((-x2, x1), dim=-1) 2446 2447 2448# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb 2449def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): 2450 """Applies Rotary Position Embedding to the query and key tensors. 2451 2452 Args: 2453 q (`torch.Tensor`): The query tensor. 2454 k (`torch.Tensor`): The key tensor. 2455 cos (`torch.Tensor`): The cosine part of the rotary embedding. 2456 sin (`torch.Tensor`): The sine part of the rotary embedding. 2457 position_ids (`torch.Tensor`): 2458 The position indices of the tokens corresponding to the query and key tensors. For example, this can be 2459 used to pass offsetted position ids when working with a KV-cache. 2460 unsqueeze_dim (`int`, *optional*, defaults to 1): 2461 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and 2462 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note 2463 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and 2464 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes 2465 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have 2466 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. 2467 Returns: 2468 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. 2469 """ 2470 cos = cos[position_ids].unsqueeze(unsqueeze_dim) 2471 sin = sin[position_ids].unsqueeze(unsqueeze_dim) 2472 2473 b, h, s, d = q.shape 2474 q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) 2475 2476 b, h, s, d = k.shape 2477 k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) 2478 2479 q_embed = (q * cos) + (rotate_half(q) * sin) 2480 k_embed = (k * cos) + (rotate_half(k) * sin) 2481 return q_embed, k_embed 2482 2483 2484class DeepseekV3MLP(nn.Module): 2485 def __init__(self, config, hidden_size=None, intermediate_size=None): 2486 super().__init__() 2487 self.config = config 2488 self.hidden_size = config.hidden_size if hidden_size is None else hidden_size 2489 self.intermediate_size = ( 2490 config.intermediate_size if intermediate_size is None else intermediate_size 2491 ) 2492 2493 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) 2494 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) 2495 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) 2496 self.act_fn = ACT2FN[config.hidden_act] 2497 2498 def forward(self, x): 2499 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) 2500 return down_proj 2501 2502 2503class MoEGate(nn.Module): 2504 def __init__(self, config): 2505 super().__init__() 2506 self.config = config 2507 self.top_k = config.num_experts_per_tok 2508 self.n_routed_experts = config.n_routed_experts 2509 self.routed_scaling_factor = config.routed_scaling_factor 2510 self.scoring_func = config.scoring_func 2511 self.topk_method = config.topk_method 2512 self.n_group = config.n_group 2513 self.topk_group = config.topk_group 2514 2515 # topk selection algorithm 2516 self.norm_topk_prob = config.norm_topk_prob 2517 self.gating_dim = config.hidden_size 2518 self.weight = nn.Parameter( 2519 torch.empty((self.n_routed_experts, self.gating_dim)) 2520 ) 2521 if self.topk_method == "noaux_tc": 2522 self.e_score_correction_bias = nn.Parameter( 2523 torch.empty((self.n_routed_experts)) 2524 ) 2525 self.reset_parameters() 2526 2527 def reset_parameters(self) -> None: 2528 import torch.nn.init as init 2529 2530 init.kaiming_uniform_(self.weight, a=math.sqrt(5)) 2531 2532 def forward(self, hidden_states): 2533 bsz, seq_len, h = hidden_states.shape 2534 ### compute gating score 2535 hidden_states = hidden_states.view(-1, h) 2536 logits = F.linear( 2537 hidden_states.type(torch.float32), self.weight.type(torch.float32), None 2538 ) 2539 if self.scoring_func == "sigmoid": 2540 scores = logits.sigmoid() 2541 else: 2542 raise NotImplementedError( 2543 f"insupportable scoring function for MoE gating: {self.scoring_func}" 2544 ) 2545 2546 ### select top-k experts 2547 if self.topk_method == "noaux_tc": 2548 assert not self.training 2549 scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0) 2550 group_scores = ( 2551 scores_for_choice.view(bsz * seq_len, self.n_group, -1).topk(2, dim=-1)[0].sum(dim = -1) 2552 ) # [n, n_group] 2553 group_idx = torch.topk( 2554 group_scores, k=self.topk_group, dim=-1, sorted=False 2555 )[ 2556 1 2557 ] # [n, top_k_group] 2558 group_mask = torch.zeros_like(group_scores) # [n, n_group] 2559 group_mask.scatter_(1, group_idx, 1) # [n, n_group] 2560 score_mask = ( 2561 group_mask.unsqueeze(-1) 2562 .expand( 2563 bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group 2564 ) 2565 .reshape(bsz * seq_len, -1) 2566 ) # [n, e] 2567 tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) # [n, e] 2568 _, topk_idx = torch.topk( 2569 tmp_scores, k=self.top_k, dim=-1, sorted=False 2570 ) 2571 topk_weight = scores.gather(1, topk_idx) 2572 else: 2573 raise NotImplementedError( 2574 f"insupportable TopK function for MoE gating: {self.topk_method}" 2575 ) 2576 2577 ### norm gate to sum 1 2578 if self.top_k > 1 and self.norm_topk_prob: 2579 denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 2580 topk_weight = topk_weight / denominator 2581 topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor 2582 2583 return topk_idx, topk_weight 2584 2585class DeepseekV3MoE(nn.Module): 2586 """ 2587 A mixed expert module containing shared experts. 2588 """ 2589 2590 def __init__(self, config): 2591 super().__init__() 2592 self.config = config 2593 self.num_experts_per_tok = config.num_experts_per_tok 2594 2595 if hasattr(config, "ep_size") and config.ep_size > 1: 2596 assert config.ep_size == dist.get_world_size() 2597 self.ep_size = config.ep_size 2598 self.experts_per_rank = config.n_routed_experts // config.ep_size 2599 self.ep_rank = dist.get_rank() 2600 self.experts = nn.ModuleList( 2601 [ 2602 ( 2603 DeepseekV3MLP( 2604 config, intermediate_size=config.moe_intermediate_size 2605 ) 2606 if i >= self.ep_rank * self.experts_per_rank 2607 and i < (self.ep_rank + 1) * self.experts_per_rank 2608 else None 2609 ) 2610 for i in range(config.n_routed_experts) 2611 ] 2612 ) 2613 else: 2614 self.ep_size = 1 2615 self.experts_per_rank = config.n_routed_experts 2616 self.ep_rank = 0 2617 self.experts = nn.ModuleList( 2618 [ 2619 DeepseekV3MLP( 2620 config, intermediate_size=config.moe_intermediate_size 2621 ) 2622 for i in range(config.n_routed_experts) 2623 ] 2624 ) 2625 self.gate = MoEGate(config) 2626 if config.n_shared_experts is not None: 2627 intermediate_size = config.moe_intermediate_size * config.n_shared_experts 2628 self.shared_experts = DeepseekV3MLP( 2629 config=config, intermediate_size=intermediate_size 2630 ) 2631 2632 def forward(self, hidden_states): 2633 identity = hidden_states 2634 orig_shape = hidden_states.shape 2635 topk_idx, topk_weight = self.gate(hidden_states) 2636 hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) 2637 flat_topk_idx = topk_idx.view(-1) 2638 if not self.training: 2639 y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape) 2640 if self.config.n_shared_experts is not None: 2641 y = y + self.shared_experts(identity) 2642 return y 2643 2644 @torch.no_grad() 2645 def moe_infer(self, x, topk_ids, topk_weight): 2646 cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) 2647 cnts.scatter_(1, topk_ids, 1) 2648 tokens_per_expert = cnts.sum(dim=0) 2649 idxs = topk_ids.view(-1).argsort() 2650 sorted_tokens = x[idxs // topk_ids.shape[1]] 2651 sorted_tokens_shape = sorted_tokens.shape 2652 if self.ep_size > 1: 2653 tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1) 2654 tokens_per_expert_group = tokens_per_expert.new_empty( 2655 tokens_per_expert.shape[0] 2656 ) 2657 dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert) 2658 output_splits = ( 2659 tokens_per_expert_group.view(self.ep_size, -1) 2660 .sum(1) 2661 .cpu() 2662 .numpy() 2663 .tolist() 2664 ) 2665 gathered_tokens = sorted_tokens.new_empty( 2666 tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1] 2667 ) 2668 input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist() 2669 dist.all_to_all( 2670 list(gathered_tokens.split(output_splits)), 2671 list(sorted_tokens.split(input_split_sizes)), 2672 ) 2673 tokens_per_expert_post_gather = tokens_per_expert_group.view( 2674 self.ep_size, self.experts_per_rank 2675 ).sum(dim=0) 2676 gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32) 2677 s = 0 2678 for i, k in enumerate(tokens_per_expert_group.cpu().numpy()): 2679 gatherd_idxs[s : s + k] = i % self.experts_per_rank 2680 s += k 2681 gatherd_idxs = gatherd_idxs.argsort() 2682 sorted_tokens = gathered_tokens[gatherd_idxs] 2683 tokens_per_expert = tokens_per_expert_post_gather 2684 tokens_per_expert = tokens_per_expert.cpu().numpy() 2685 2686 outputs = [] 2687 start_idx = 0 2688 for i, num_tokens in enumerate(tokens_per_expert): 2689 end_idx = start_idx + num_tokens 2690 if num_tokens == 0: 2691 continue 2692 expert = self.experts[i + self.ep_rank * self.experts_per_rank] 2693 tokens_for_this_expert = sorted_tokens[start_idx:end_idx] 2694 expert_out = expert(tokens_for_this_expert) 2695 outputs.append(expert_out) 2696 start_idx = end_idx 2697 2698 outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) 2699 if self.ep_size > 1: 2700 new_x = torch.empty_like(outs) 2701 new_x[gatherd_idxs] = outs 2702 gathered_tokens = new_x.new_empty(*sorted_tokens_shape) 2703 dist.all_to_all( 2704 list(gathered_tokens.split(input_split_sizes)), 2705 list(new_x.split(output_splits)), 2706 ) 2707 outs = gathered_tokens 2708 2709 new_x = torch.empty_like(outs) 2710 new_x[idxs] = outs 2711 final_out = ( 2712 new_x.view(*topk_ids.shape, -1) 2713 .type(topk_weight.dtype) 2714 .mul_(topk_weight.unsqueeze(dim=-1)) 2715 .sum(dim=1) 2716 .type(new_x.dtype) 2717 ) 2718 return final_out 2719 2720 2721# Copied from transformers.models.llama.modeling_llama.repeat_kv 2722def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: 2723 """ 2724 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, 2725 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) 2726 """ 2727 batch, num_key_value_heads, slen, head_dim = hidden_states.shape 2728 if n_rep == 1: 2729 return hidden_states 2730 hidden_states = hidden_states[:, :, None, :, :].expand( 2731 batch, num_key_value_heads, n_rep, slen, head_dim 2732 ) 2733 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) 2734 2735 2736# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV3 2737class DeepseekV3Attention(nn.Module): 2738 """Multi-headed attention from 'Attention Is All You Need' paper""" 2739 2740 def __init__(self, config: DeepseekV3Config, layer_idx: Optional[int] = None): 2741 super().__init__() 2742 self.config = config 2743 self.layer_idx = layer_idx 2744 if layer_idx is None: 2745 logger.warning_once( 2746 f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " 2747 "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " 2748 "when creating this class." 2749 ) 2750 2751 self.attention_dropout = config.attention_dropout 2752 self.hidden_size = config.hidden_size 2753 self.num_heads = config.num_attention_heads 2754 2755 self.max_position_embeddings = config.max_position_embeddings 2756 self.rope_theta = config.rope_theta 2757 self.q_lora_rank = config.q_lora_rank 2758 self.qk_rope_head_dim = config.qk_rope_head_dim 2759 self.kv_lora_rank = config.kv_lora_rank 2760 self.v_head_dim = config.v_head_dim 2761 self.qk_nope_head_dim = config.qk_nope_head_dim 2762 self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim 2763 2764 self.is_causal = True 2765 2766 if self.q_lora_rank is None: 2767 self.q_proj = nn.Linear( 2768 self.hidden_size, self.num_heads * self.q_head_dim, bias=False 2769 ) 2770 else: 2771 self.q_a_proj = nn.Linear( 2772 self.hidden_size, config.q_lora_rank, bias=config.attention_bias 2773 ) 2774 self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank) 2775 self.q_b_proj = nn.Linear( 2776 config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False 2777 ) 2778 2779 self.kv_a_proj_with_mqa = nn.Linear( 2780 self.hidden_size, 2781 config.kv_lora_rank + config.qk_rope_head_dim, 2782 bias=config.attention_bias, 2783 ) 2784 self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank) 2785 self.kv_b_proj = nn.Linear( 2786 config.kv_lora_rank, 2787 self.num_heads 2788 * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), 2789 bias=False, 2790 ) 2791 2792 self.o_proj = nn.Linear( 2793 self.num_heads * self.v_head_dim, 2794 self.hidden_size, 2795 bias=config.attention_bias, 2796 ) 2797 self._init_rope() 2798 2799 self.softmax_scale = self.q_head_dim ** (-0.5) 2800 if self.config.rope_scaling is not None: 2801 mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) 2802 scaling_factor = self.config.rope_scaling["factor"] 2803 if mscale_all_dim: 2804 mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) 2805 self.softmax_scale = self.softmax_scale * mscale * mscale 2806 2807 def _init_rope(self): 2808 if self.config.rope_scaling is None: 2809 self.rotary_emb = DeepseekV3RotaryEmbedding( 2810 self.qk_rope_head_dim, 2811 max_position_embeddings=self.max_position_embeddings, 2812 base=self.rope_theta, 2813 ) 2814 else: 2815 scaling_type = self.config.rope_scaling["type"] 2816 scaling_factor = self.config.rope_scaling["factor"] 2817 if scaling_type == "linear": 2818 self.rotary_emb = DeepseekV3LinearScalingRotaryEmbedding( 2819 self.qk_rope_head_dim, 2820 max_position_embeddings=self.max_position_embeddings, 2821 scaling_factor=scaling_factor, 2822 base=self.rope_theta, 2823 ) 2824 elif scaling_type == "dynamic": 2825 self.rotary_emb = DeepseekV3DynamicNTKScalingRotaryEmbedding( 2826 self.qk_rope_head_dim, 2827 max_position_embeddings=self.max_position_embeddings, 2828 scaling_factor=scaling_factor, 2829 base=self.rope_theta, 2830 ) 2831 elif scaling_type == "yarn": 2832 kwargs = { 2833 key: self.config.rope_scaling[key] 2834 for key in [ 2835 "original_max_position_embeddings", 2836 "beta_fast", 2837 "beta_slow", 2838 "mscale", 2839 "mscale_all_dim", 2840 ] 2841 if key in self.config.rope_scaling 2842 } 2843 self.rotary_emb = DeepseekV3YarnRotaryEmbedding( 2844 self.qk_rope_head_dim, 2845 max_position_embeddings=self.max_position_embeddings, 2846 scaling_factor=scaling_factor, 2847 base=self.rope_theta, 2848 **kwargs, 2849 ) 2850 else: 2851 raise ValueError(f"Unknown RoPE scaling type {scaling_type}") 2852 2853 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): 2854 return ( 2855 tensor.view(bsz, seq_len, self.num_heads, self.v_head_dim) 2856 .transpose(1, 2) 2857 .contiguous() 2858 ) 2859 2860 def forward( 2861 self, 2862 hidden_states: torch.Tensor, 2863 attention_mask: Optional[torch.Tensor] = None, 2864 position_ids: Optional[torch.LongTensor] = None, 2865 past_key_value: Optional[Cache] = None, 2866 output_attentions: bool = False, 2867 use_cache: bool = False, 2868 **kwargs, 2869 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: 2870 if "padding_mask" in kwargs: 2871 warnings.warn( 2872 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" 2873 ) 2874 bsz, q_len, _ = hidden_states.size() 2875 2876 if self.q_lora_rank is None: 2877 q = self.q_proj(hidden_states) 2878 else: 2879 q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) 2880 q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) 2881 q_nope, q_pe = torch.split( 2882 q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 2883 ) 2884 2885 compressed_kv = self.kv_a_proj_with_mqa(hidden_states) 2886 compressed_kv, k_pe = torch.split( 2887 compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 2888 ) 2889 k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) 2890 kv = ( 2891 self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) 2892 .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) 2893 .transpose(1, 2) 2894 ) 2895 2896 k_nope, value_states = torch.split( 2897 kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 2898 ) 2899 kv_seq_len = value_states.shape[-2] 2900 if past_key_value is not None: 2901 if self.layer_idx is None: 2902 raise ValueError( 2903 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " 2904 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " 2905 "with a layer index." 2906 ) 2907 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) 2908 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) 2909 2910 q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) 2911 2912 query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) 2913 query_states[:, :, :, : self.qk_nope_head_dim] = q_nope 2914 query_states[:, :, :, self.qk_nope_head_dim :] = q_pe 2915 2916 key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) 2917 key_states[:, :, :, : self.qk_nope_head_dim] = k_nope 2918 key_states[:, :, :, self.qk_nope_head_dim :] = k_pe 2919 if past_key_value is not None: 2920 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models 2921 key_states, value_states = past_key_value.update( 2922 key_states, value_states, self.layer_idx, cache_kwargs 2923 ) 2924 2925 attn_weights = ( 2926 torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale 2927 ) 2928 2929 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): 2930 raise ValueError( 2931 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" 2932 f" {attn_weights.size()}" 2933 ) 2934 assert attention_mask is not None 2935 if attention_mask is not None: 2936 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): 2937 raise ValueError( 2938 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" 2939 ) 2940 attn_weights = attn_weights + attention_mask 2941 2942 # upcast attention to fp32 2943 attn_weights = nn.functional.softmax( 2944 attn_weights, dim=-1, dtype=torch.float32 2945 ).to(query_states.dtype) 2946 attn_weights = nn.functional.dropout( 2947 attn_weights, p=self.attention_dropout, training=self.training 2948 ) 2949 attn_output = torch.matmul(attn_weights, value_states) 2950 2951 if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim): 2952 raise ValueError( 2953 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is" 2954 f" {attn_output.size()}" 2955 ) 2956 2957 attn_output = attn_output.transpose(1, 2).contiguous() 2958 2959 attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) 2960 2961 attn_output = self.o_proj(attn_output) 2962 2963 if not output_attentions: 2964 attn_weights = None 2965 2966 return attn_output, attn_weights, past_key_value 2967 2968 2969# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->DeepseekV3 2970class DeepseekV3FlashAttention2(DeepseekV3Attention): 2971 """ 2972 DeepseekV3 flash attention module. This module inherits from `DeepseekV3Attention` as the weights of the module stays 2973 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of 2974 flash attention and deal with padding tokens in case the input contains any of them. 2975 """ 2976 2977 def __init__(self, *args, **kwargs): 2978 super().__init__(*args, **kwargs) 2979 2980 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. 2981 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. 2982 # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). 2983 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() 2984 2985 def forward( 2986 self, 2987 hidden_states: torch.Tensor, 2988 attention_mask: Optional[torch.LongTensor] = None, 2989 position_ids: Optional[torch.LongTensor] = None, 2990 past_key_value: Optional[Cache] = None, 2991 output_attentions: bool = False, 2992 use_cache: bool = False, 2993 **kwargs, 2994 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: 2995 # DeepseekV3FlashAttention2 attention does not support output_attentions 2996 if "padding_mask" in kwargs: 2997 warnings.warn( 2998 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" 2999 ) 3000