Upload 3 files

config.json

@@ -0,0 +1,40 @@
+{
+  "architectures": [
+    "Mymodel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 2,
+  "embedding_dim": 4096,
+  "embedding_method": "pma",
+  "encoder_mode": "post_normal",
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "inf_seq_length": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 13440,
+  "keep_max_layer": 32,
+  "max_position_embeddings": 65536,
+  "max_window_layers": 28,
+  "model_type": "qwen2",
+  "num_attention_heads": 32,
+  "num_encoder_layers": 0,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 4,
+  "padding_side": "right",
+  "pma_ln": true,
+  "pma_norm": false,
+  "pma_norm_mode": "post_normal",
+  "pma_num_heads": 32,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 1000000,
+  "rotary_emb_base": 1000000,
+  "seq_length": 65536,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.39.2",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 92416
+}

configuration_d2coder.py

@@ -0,0 +1,149 @@
+from transformers import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class D2CoderConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`D2LLM`]. It is used to instantiate a
+    Qwen2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Qwen2-7B-beta [Qwen/Qwen2-7B-beta](https://huggingface.co/Qwen/Qwen2-7B-beta).
+
+    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 151936):
+            Vocabulary size of the Qwen2 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`D2LLM`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 22016):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 32):
+            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 `32`.
+        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 32768):
+            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`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 28):
+            The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import Qwen2Model, Qwen2Config
+
+    >>> # Initializing a Qwen2 style configuration
+    >>> configuration = Qwen2Config()
+
+    >>> # Initializing a model from the Qwen2-7B style configuration
+    >>> model = Qwen2Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "qwen2"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=4096,
+        intermediate_size=22016,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=32,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        attention_dropout=0.0,
+
+        embedding_method="pma",
+        inf_seq_length=1024,
+        encoder_mode ="post_normal",
+        num_encoder_layers =0,
+        padding_side ="right",
+
+        keep_max_layer=32,
+        pma_num_heads=32,
+        pma_ln=True,
+        pma_norm=False,
+        pma_norm_mode="post_normal",
+
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        # 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.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+
+        self.embedding_method = config.embedding_method
+        self.inf_seq_length = config.inf_seq_length
+        self.encoder_mode = config.encoder_mode
+        self.num_encoder_layers = config.num_encoder_layers
+        self.padding_side = config.padding_side
+
+        self.keep_max_layer = config.keep_max_layer
+        self.pma_num_heads = config.pma_num_heads
+        self.pma_ln = config.pma_ln
+        self.pma_norm = config.pma_norm
+        self.pma_norm_mode = config.pma_norm_mode
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

modeling_d2coder.py

@@ -0,0 +1,638 @@
+from transformers import Qwen2Config
+import inspect
+import math
+import os
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers import PretrainedConfig
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa, _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+import numpy as np
+from transformers import Qwen2Config
+from transformers import Qwen2ForCausalLM
+import inspect
+import math
+import os
+import warnings
+from typing import List, Optional, Tuple, Union
+from tqdm import tqdm, trange
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa, _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+import numpy as np
+import torch
+import os
+import argparse
+import json
+from tqdm import tqdm
+from typing import cast, List, Union, Tuple
+from transformers import AutoTokenizer, AutoModel  # pylint: disable=C0413
+from peft import LoraConfig, get_peft_model, TaskType
+import time
+import torch.nn.functional as F
+import sys
+import time
+import torch 
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from tqdm import tqdm, trange
+from collections import defaultdict
+from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM, AutoConfig
+import torch.distributed as dist
+from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import re
+
+
+# PMA部分 post_normal
+class MAB_POST(nn.Module):
+    def __init__(self, dim_Q, dim_K, dim_V, num_heads, ln=False):
+        super(MAB_POST, self).__init__()
+        self.dim_V = dim_V
+        self.num_heads = num_heads
+        self.fc_q = nn.Linear(dim_Q, dim_V)
+        self.fc_k = nn.Linear(dim_K, dim_V)
+        self.fc_v = nn.Linear(dim_K, dim_V)
+
+        if ln:
+            self.ln0 = nn.LayerNorm(dim_V)
+            self.ln1 = nn.LayerNorm(dim_V)
+        self.fc_o = nn.Linear(dim_V, dim_V)
+        nn.init.xavier_uniform_(self.fc_q.weight)
+        nn.init.xavier_uniform_(self.fc_k.weight)
+        nn.init.xavier_uniform_(self.fc_v.weight)
+        nn.init.xavier_uniform_(self.fc_o.weight)
+
+
+
+    # Q(bs, 1, emb), pad_mask (bs, seq) Post-LN
+    def forward(self, Q, K, pad_mask=None):
+
+        Q_ = self.fc_q(Q)
+        K_, V_ = self.fc_k(K), self.fc_v(K)
+
+        dim_split = self.dim_V // self.num_heads
+        Q_ = torch.cat(Q_.split(dim_split, 2), 0) # (bs* num_head, 1, emb)
+        K_ = torch.cat(K_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        V_ = torch.cat(V_.split(dim_split, 2), 0)
+
+        pad_mask = pad_mask.unsqueeze(1).repeat(self.num_heads, Q.size(1), 1) # (bs*num_head, 1, seq)
+        score = Q_.bmm(K_.transpose(1,2))/math.sqrt(self.dim_V)
+        score = score.masked_fill(pad_mask == 0, -1e12)
+        A = torch.softmax(score, 2)  # (bs*num_head, 1, seq)
+        A = A * pad_mask
+        O = torch.cat(A.bmm(V_).split(Q.size(0), 0), 2) # (bs, 1, emb)
+        O = Q + O
+        # O = torch.cat((Q_ + A.bmm(V_)).split(Q.size(0), 0), 2)
+        O = O if getattr(self, 'ln0', None) is None else self.ln0(O)
+        O = O + F.relu(self.fc_o(O))
+        O = O if getattr(self, 'ln1', None) is None else self.ln1(O)
+        return O
+
+
+# PMA部分 pre_normal
+class MAB_PRE_NORMAL(nn.Module):
+    def __init__(self, dim_Q, dim_K, dim_V, num_heads, ln=False):
+        super(MAB_PRE_NORMAL, self).__init__()
+        self.dim_V = dim_V
+        self.num_heads = num_heads
+        self.fc_q = nn.Linear(dim_Q, dim_V)
+        self.fc_k = nn.Linear(dim_K, dim_V)
+        self.fc_v = nn.Linear(dim_K, dim_V)
+
+        if ln:
+            self.ln_q = nn.LayerNorm(dim_V)
+            self.ln_kv = nn.LayerNorm(dim_V)
+            self.ln_o = nn.LayerNorm(dim_V)
+            self.ln_final = nn.LayerNorm(dim_V)
+        
+        self.fc_o = nn.Linear(dim_V, dim_V)
+        nn.init.xavier_uniform_(self.fc_q.weight)
+        nn.init.xavier_uniform_(self.fc_k.weight)
+        nn.init.xavier_uniform_(self.fc_v.weight)
+        nn.init.xavier_uniform_(self.fc_o.weight)
+
+
+
+    
+    # pad_mask (bs, seq) Pre-LN 正常架构
+    def forward(self, Q, K, pad_mask=None):
+
+        Q_ = Q if getattr(self, 'ln_q', None) is None else self.ln_q(Q)
+        K_ = K if getattr(self, 'ln_kv', None) is None else self.ln_kv(K)
+
+        Q_ = self.fc_q(Q_) 
+        K_, V_ = self.fc_k(K_), self.fc_v(K_)
+
+        dim_split = self.dim_V // self.num_heads
+
+
+        Q_ = torch.cat(Q_.split(dim_split, 2), 0) # (bs* num_head, 1, emb)
+        K_ = torch.cat(K_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        V_ = torch.cat(V_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        pad_mask = pad_mask.unsqueeze(1).repeat(self.num_heads, Q.size(1), 1) # (bs*num_head, 1, seq)
+        score = Q_.bmm(K_.transpose(1,2))/math.sqrt(self.dim_V)
+        score = score.masked_fill(pad_mask == 0, -1e12)
+        A = torch.softmax(score, 2)  # (bs*num_head, 1, seq)
+        A = A * pad_mask
+
+
+        O = torch.cat(A.bmm(V_).split(Q.size(0), 0), 2) 
+        O = Q + O
+
+        O_ = O if getattr(self, 'ln_o', None) is None else self.ln_o(O) # O的layernorm分支
+        O_ = O + F.relu(self.fc_o(O_))
+        return O_ if getattr(self, 'ln_final', None) is None else self.ln_final(O_)
+
+# PMA部分 pre_gptj
+class MAB_PRE_GPTJ(nn.Module):
+    def __init__(self, dim_Q, dim_K, dim_V, num_heads, ln=False):
+        super(MAB_PRE_GPTJ, self).__init__()
+        self.dim_V = dim_V
+        self.num_heads = num_heads
+        self.fc_q = nn.Linear(dim_Q, dim_V)
+        self.fc_k = nn.Linear(dim_K, dim_V)
+        self.fc_v = nn.Linear(dim_K, dim_V)
+        self.fc_o = nn.Linear(dim_V, dim_V)
+        
+        nn.init.xavier_uniform_(self.fc_q.weight)
+        nn.init.xavier_uniform_(self.fc_k.weight)
+        nn.init.xavier_uniform_(self.fc_v.weight)
+        nn.init.xavier_uniform_(self.fc_o.weight)
+        if ln:
+            self.ln_q = nn.LayerNorm(dim_V)
+            self.ln_kv = nn.LayerNorm(dim_V)
+            self.ln_final = nn.LayerNorm(dim_V)
+
+    # pad_mask (bs, seq) 
+    def forward(self, Q, K, pad_mask=None):
+        
+        # layernorm
+        Q_ = Q if getattr(self, 'ln_q', None) is None else self.ln_q(Q)
+        K_ = K if getattr(self, 'ln_kv', None) is None else self.ln_kv(K)
+
+
+        Q1 = self.fc_q(Q_) 
+        K1, V1 = self.fc_k(K_), self.fc_v(K_)
+        dim_split = self.dim_V // self.num_heads
+
+
+        Q1 = torch.cat(Q1.split(dim_split, 2), 0) # (bs* num_head, 1, emb)
+        K1 = torch.cat(K1.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        V1 = torch.cat(V1.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+
+
+        pad_mask = pad_mask.unsqueeze(1).repeat(self.num_heads, Q.size(1), 1) # (bs*num_head, 1, seq)
+        score = Q1.bmm(K1.transpose(1,2))/math.sqrt(self.dim_V)
+        score = score.masked_fill(pad_mask == 0, -1e12)
+        A = torch.softmax(score, 2)  # (bs*num_head, 1, seq)
+        A = A * pad_mask
+        O1 = torch.cat(A.bmm(V1).split(Q.size(0), 0), 2)  # (bs, 1, emb)
+
+        O2 = F.relu(self.fc_o(Q_))  # (bs, 1, emb)
+
+        O_final = Q + O1 + O2
+
+        return O_final if getattr(self, 'ln_final', None) is None else self.ln_final(O_final)
+
+
+
+
+class PMA(nn.Module):
+    def __init__(self, dim, num_heads, num_seeds, ln=False, pma_mode=None):
+        super(PMA, self).__init__()
+        self.S = nn.Parameter(torch.Tensor(1, num_seeds, dim))
+        nn.init.xavier_uniform_(self.S)
+        if pma_mode == 'post_normal':
+            self.mab = MAB_POST(dim, dim, dim, num_heads, ln=ln)
+        elif pma_mode == 'pre_normal':
+            self.mab = MAB_PRE_NORMAL(dim, dim, dim, num_heads, ln=ln)
+        elif pma_mode == 'pre_gptj':
+            self.mab = MAB_PRE_GPTJ(dim, dim, dim, num_heads, ln=ln)
+        else:
+            raise ValueError(f"Error, the pma_mode {pma_mode} is not implemented !")
+    # X: (bs, seq, emb), pad_mask: (bs, seq)
+    def forward(self, X, pad_mask):
+        if self.S.dtype != torch.bfloat16:
+            X = X.float()
+        return self.mab(self.S.repeat(X.size(0), 1, 1), X, pad_mask)
+
+
+# 普通双向transformer encoder, post_normal
+class EncoderLayer_POST(nn.Module):
+    def __init__(self, dim_V, num_heads, ln=False):
+        super(EncoderLayer_POST, self).__init__()
+        self.dim_V = dim_V
+        self.num_heads = num_heads
+        self.fc_q = nn.Linear(dim_V, dim_V)
+        self.fc_k = nn.Linear(dim_V, dim_V)
+        self.fc_v = nn.Linear(dim_V, dim_V)
+        self.fc_o = nn.Linear(dim_V, dim_V)
+        
+        
+        nn.init.xavier_uniform_(self.fc_q.weight)
+        nn.init.xavier_uniform_(self.fc_k.weight)
+        nn.init.xavier_uniform_(self.fc_v.weight)
+        nn.init.xavier_uniform_(self.fc_o.weight)
+
+        if ln:
+            self.ln0 = nn.LayerNorm(dim_V)
+            self.ln1 = nn.LayerNorm(dim_V)
+
+    # Q:(bs, seq, emb), pad_mask:(bs, seq)
+    def forward(self, Q, pad_mask=None):
+
+        Q_, K_, V_ = self.fc_q(Q), self.fc_k(Q), self.fc_v(Q) 
+
+        dim_split = self.dim_V // self.num_heads
+        Q_ = torch.cat(Q_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        K_ = torch.cat(K_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        V_ = torch.cat(V_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+
+        pad_mask = pad_mask.unsqueeze(1).repeat(self.num_heads, Q.size(1), 1)  # (bs*num_head, seq, seq)
+
+        score = Q_.bmm(K_.transpose(1,2))/math.sqrt(self.dim_V)
+        score = score.masked_fill(pad_mask == 0, -1e12)
+        A = torch.softmax(score, 2)  # (bs*num_head, seq, seq)
+        A = A * pad_mask  # (bs*num_head, seq, seq)
+
+        O = torch.cat(A.bmm(V_).split(Q.size(0), 0), 2) # (bs, seq, emb)
+        O = Q + O
+
+        O = O if getattr(self, 'ln0', None) is None else self.ln0(O)  
+        O = O + F.relu(self.fc_o(O))
+        O = O if getattr(self, 'ln1', None) is None else self.ln1(O)
+        return O
+
+
+# 普通双向transformer encoder, pre LN norm
+class EncoderLayer_PRE_NORMAL(nn.Module):
+    def __init__(self, dim_V, num_heads, ln=False):
+        super(EncoderLayer_PRE_NORMAL, self).__init__()
+        self.dim_V = dim_V
+        self.num_heads = num_heads
+        self.fc_q = nn.Linear(dim_V, dim_V)
+        self.fc_k = nn.Linear(dim_V, dim_V)
+        self.fc_v = nn.Linear(dim_V, dim_V)
+        self.fc_o = nn.Linear(dim_V, dim_V)
+        
+        
+        nn.init.xavier_uniform_(self.fc_q.weight)
+        nn.init.xavier_uniform_(self.fc_k.weight)
+        nn.init.xavier_uniform_(self.fc_v.weight)
+        nn.init.xavier_uniform_(self.fc_o.weight)
+
+        if ln:
+            self.ln_qkv = nn.LayerNorm(dim_V)
+            self.ln_o = nn.LayerNorm(dim_V)
+
+    # Q:(bs, seq, emb), pad_mask:(bs, seq)
+    def forward(self, Q, pad_mask=None):
+
+        Q_ = Q if getattr(self, 'ln_qkv', None) is None else self.ln_qkv(Q) # layernorm
+
+        Q_, K_, V_ = self.fc_q(Q_), self.fc_k(Q_), self.fc_v(Q_)
+        dim_split = self.dim_V // self.num_heads
+        Q_ = torch.cat(Q_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        K_ = torch.cat(K_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        V_ = torch.cat(V_.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        pad_mask = pad_mask.unsqueeze(1).repeat(self.num_heads, Q.size(1), 1)  # (bs*num_head, seq, seq)
+        score = Q_.bmm(K_.transpose(1,2))/math.sqrt(self.dim_V)
+        score = score.masked_fill(pad_mask == 0, -1e12)
+        A = torch.softmax(score, 2)  # (bs*num_head, seq, seq)
+        A = A * pad_mask
+
+        O = torch.cat(A.bmm(V_).split(Q.size(0), 0), 2)
+        O = Q + O
+
+        O_ = O if getattr(self, 'ln_o', None) is None else self.ln_o(O) # O的layernorm分支
+
+        O_ = O + F.relu(self.fc_o(O_))
+
+        return O_ 
+
+# 普通双向transformer encoder, pre LN gptj
+class EncoderLayer_PRE_GPTJ(nn.Module):
+    def __init__(self, dim_V, num_heads, ln=False):
+        super(EncoderLayer_PRE_GPTJ, self).__init__()
+        self.dim_V = dim_V
+        self.num_heads = num_heads
+        self.fc_q = nn.Linear(dim_V, dim_V)
+        self.fc_k = nn.Linear(dim_V, dim_V)
+        self.fc_v = nn.Linear(dim_V, dim_V)
+        self.fc_o = nn.Linear(dim_V, dim_V)
+        
+        
+        nn.init.xavier_uniform_(self.fc_q.weight)
+        nn.init.xavier_uniform_(self.fc_k.weight)
+        nn.init.xavier_uniform_(self.fc_v.weight)
+        nn.init.xavier_uniform_(self.fc_o.weight)
+
+        if ln:
+            self.ln_qkv = nn.LayerNorm(dim_V)
+
+    # Q:(bs, seq, emb), pad_mask:(bs, seq)
+    def forward(self, Q, pad_mask=None):
+
+        Q_ = Q if getattr(self, 'ln_qkv', None) is None else self.ln_qkv(Q) # layernorm
+
+
+        Q1, K1, V1 = self.fc_q(Q_), self.fc_k(Q_), self.fc_v(Q_)
+        dim_split = self.dim_V // self.num_heads
+        Q1 = torch.cat(Q1.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        K1 = torch.cat(K1.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        V1 = torch.cat(V1.split(dim_split, 2), 0) # (bs* num_head, seq, emb)
+        pad_mask = pad_mask.unsqueeze(1).repeat(self.num_heads, Q.size(1), 1)  # (bs*num_head, seq, seq)
+        score = Q1.bmm(K1.transpose(1,2))/math.sqrt(self.dim_V)
+        score = score.masked_fill(pad_mask == 0, -1e12)
+        A = torch.softmax(score, 2)  # (bs*num_head, seq, seq)
+        A = A * pad_mask
+        O1 = torch.cat(A.bmm(V1).split(Q.size(0), 0), 2) # （bs, seq, emb）
+
+        O2 = F.relu(self.fc_o(Q_)) 
+
+        O_final = Q + O1 + O2
+
+        return O_final
+
+
+class Encoder(nn.Module):
+    def __init__(self, emb_dim, num_heads, ln, encoder_mode, num_encoder_layers):
+        super(Encoder, self).__init__()
+        self.num_encoder_layers = num_encoder_layers
+        if encoder_mode == 'post_normal':
+            self.layers = nn.ModuleList([EncoderLayer_POST(dim_V=emb_dim, num_heads=num_heads, ln=ln)
+                                        for _ in range(num_encoder_layers)])
+        elif encoder_mode == 'pre_normal':
+            self.layers = nn.ModuleList([EncoderLayer_PRE_NORMAL(dim_V=emb_dim, num_heads=num_heads, ln=ln)
+                                        for _ in range(num_encoder_layers)])
+        elif encoder_mode == 'pre_gptj':
+            self.layers = nn.ModuleList([EncoderLayer_PRE_GPTJ(dim_V=emb_dim, num_heads=num_heads, ln=ln)
+                                        for _ in range(num_encoder_layers)])
+        else:
+            raise ValueError(f"Error, the encoder_mode {encoder_mode} is not implemented !")
+    
+    # X:(bs, seq, emb), mask: (bs, seq)
+    def forward(self, X, mask):
+        if self.num_encoder_layers == 0:
+            return X
+        if self.layers[0].fc_q.weight.dtype != torch.bfloat16:
+            X = X.float()
+        for layer in self.layers:
+            X = layer(X, mask)
+
+        return X
+
+class D2LLMConfig(PretrainedConfig):
+    model_type = "qwen2"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=4096,
+        intermediate_size=22016,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=32,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        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.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        # 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.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+class D2Coder(PreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.plm_model = Qwen2ForCausalLM(config)
+
+        self.embedding_method = config.embedding_method
+        self.inf_seq_length = config.inf_seq_length
+        self.encoder_mode = config.encoder_mode
+        self.num_encoder_layers = config.num_encoder_layers
+        self.padding_side = config.padding_side
+
+        self.keep_max_layer = config.keep_max_layer
+        self.emb_dim = self.plm_model.model.embed_tokens.weight.size(1)
+        self.num_heads = config.pma_num_heads
+        self.ln = config.pma_ln
+        self.norm = config.pma_norm
+        self.pma_mode = config.pma_norm_mode
+        self.encoder = Encoder(self.emb_dim, self.num_heads, self.ln, self.encoder_mode, self.num_encoder_layers)
+        self.mha_pma = PMA(self.emb_dim, self.num_heads, 1, ln=self.ln, pma_mode=self.pma_mode)
+
+    def forward(self, inputs_all, mode, args):
+        # output_embeddings_a = self.get_sentence_embedding(self.embedding_method, **inputs_a)
+        
+        # output_embeddings_b = self.get_sentence_embedding(self.embedding_method, **inputs_b) # (bs, emb_size)
+        bs = self.args.batch_size
+        if mode == 'train':
+            output_embeddings_all = self.get_sentence_embedding(self.embedding_method, **inputs_all).reshape(2+self.args.neg_K, bs, -1) # (2+K, bs, emb_size)
+            # if self.to_compress:
+            #     output_embeddings_all = self.projector(output_embeddings_all)
+
+            output_embeddings_hardneg = output_embeddings_all[2:] # (neg_K, bs, emb)
+            hn_norm = torch.nn.functional.normalize(output_embeddings_hardneg, p=2, dim=-1)
+        elif mode == 'eval':
+            output_embeddings_all = self.get_sentence_embedding(self.embedding_method, **inputs_all).reshape(2, bs, -1) # (2, bs, emb_size)
+            # if self.to_compress:
+            #     output_embeddings_all = self.projector(output_embeddings_all)
+        else:
+            raise ValueError('Error of mode value')
+
+        output_embeddings_a = output_embeddings_all[0] # (bs, emb)
+        output_embeddings_b = output_embeddings_all[1] # (bs, emb)
+        a_norm = torch.nn.functional.normalize(output_embeddings_a, p=2, dim=-1)
+        b_norm = torch.nn.functional.normalize(output_embeddings_b, p=2, dim=-1)
+        
+        
+
+        b_cross_gpus = gather_across_devices(output_embeddings_b, args.global_rank, self.world_size) 
+        b_norm_cross_gpus = torch.nn.functional.normalize(b_cross_gpus, p=2, dim=-1) # ()
+
+
+        assert a_norm.size(0) == b_norm.size(0)
+        bs = output_embeddings_a.size(0)
+        # in-batch计算部分
+        output_in_batch_local_gpu = torch.matmul(a_norm, b_norm.t())
+        output_in_batch_global_gpu = torch.matmul(a_norm, b_norm_cross_gpus.t())
+
+        if mode == 'train':
+            # hard neg计算部分
+            pos_neg_emb = torch.cat([b_norm.unsqueeze(0), hn_norm], dim=0) # (1+neg_K, bs, emb)
+            output_hardneg_specific_task = torch.matmul(a_norm.unsqueeze(1), pos_neg_emb.permute(1,2,0)).squeeze() #  (bs, 1+neg_K)
+            # output_pos_hardneg_rep_specific_task = torch.cat([output_embeddings_a.unsqueeze(0).expand(pos_neg_emb.size(0),-1,-1), pos_neg_emb],dim=-1)
+
+        elif mode == 'eval':
+            output_hardneg_specific_task = None
+            output_pos_hardneg_rep_specific_task = None
+        
+        return output_in_batch_local_gpu, output_in_batch_global_gpu, output_hardneg_specific_task  # (bs, bs) (bs, world_size*bs), (bs, 1+neg_K)  
+        # return output_in_batch_specific_task, output_hardneg_specific_task, output_pos_hardneg_rep_specific_task 
+
+    def last_embedding(self, A, index):
+        bs, seq, emb = A.size()
+        res = A[torch.arange(bs), index, :]
+        return res
+
+    def mean_embedding(self, A, mask):
+        bs, seq, emb = A.size()
+        res = (A * (mask.unsqueeze(-1))).sum(1) / (mask.sum(1).unsqueeze(-1))
+        return res
+    
+    # A (bs, seq, emb_size), mask (bs, 1, seq)
+    def weighted_embedding(self, A, mask):
+        weights = (torch.arange(start=1, end=A.size(1) + 1).unsqueeze(0).unsqueeze(-1).expand(A.size()).float()).to(A.device)
+        input_mask_expanded = (mask.squeeze(1).unsqueeze(-1).expand(A.size()).float()).to(A.device)
+        sum_embedding = torch.sum(A * input_mask_expanded * weights, dim=1)
+        sum_mask = torch.sum(input_mask_expanded * weights, dim=1)
+        weighted_embedding = sum_embedding / sum_mask
+        
+        return weighted_embedding
+
+    def pma_embedding(self, A, mask):
+        res = self.mha_pma(A, mask).squeeze(1)
+        return res
+
+
+    def get_sentence_embedding(self, embedding_method, **inputs):
+        outputs = self.plm_model(inputs['input_ids'], inputs['attention_mask'], output_hidden_states=True)
+        if embedding_method == 'last':
+            embedding = outputs.hidden_states[self.keep_max_layer]
+            index = inputs['attention_mask'].sum(-1).long() - 1
+            res_embedding = self.last_embedding(embedding, index)
+        elif embedding_method == 'mean':
+            embedding = outputs.hidden_states[self.keep_max_layer]
+            res_embedding = self.mean_embedding(embedding, inputs['attention_mask'])
+        elif embedding_method == 'weighted':
+            embedding = outputs.hidden_states[self.keep_max_layer]
+            res_embedding = self.weighted_embedding(embedding, inputs['attention_mask'])
+        elif embedding_method == 'pma':
+            embedding = outputs.hidden_states[self.keep_max_layer] # Qwen.hidden_state: (33, bs, seq, emb)
+            attention_mask = inputs['attention_mask']
+            embedding = self.encoder(embedding, attention_mask)
+            res_embedding = self.pma_embedding(embedding, attention_mask) # embedding: (bs, seq, emb), inputs['attention_mask']: (bs, seq)
+        else:
+            logger.debug('Error, no {} way to obtain embbedings'.format(embedding_method))
+        
+        if not self.norm:
+            res_embedding = torch.nn.functional.normalize(res_embedding, p=2.0, dim=-1, eps=1e-12, out=None)
+        return res_embedding
+
+
+
+    def encode(self, tokenizer, sentences, batch_size=32, convert_to_numpy=True,
+            convert_to_tensor=False, show_progress_bar=True, max_seq_length=None, **kwargs):
+
+        if max_seq_length is None:
+            max_seq_length = self.inf_seq_length
+
+        input_is_string = False        
+        if isinstance(sentences, str) or not hasattr(sentences, "__len__"):
+            sentences = [sentences]
+            input_is_string = True
+
+
+        all_embeddings = []
+        length_sorted_idx = np.argsort([-len(s) for s in sentences])
+        sentences_sorted = [sentences[idx] for idx in length_sorted_idx] # 大到小重排
+        with torch.no_grad():
+            for start_index in trange(0, len(sentences), batch_size, desc="Batches", disable=not show_progress_bar):
+                sentences_batch = sentences_sorted[start_index: start_index + batch_size]
+                # Compute sentences embeddingsz
+                with torch.no_grad():
+                    inputs = tokenizer(sentences_batch, padding=True, truncation=True, max_length=max_seq_length, add_special_tokens=False, return_tensors='pt').to(self.plm_model.device)
+                    embeddings = self.get_sentence_embedding(self.embedding_method, **inputs)
+                    # if self.to_compress:
+                    #     embeddings = self.projector(embeddings)
+                embeddings = embeddings.detach()
+                if convert_to_numpy:
+                    if embeddings.dtype == torch.bfloat16:
+                        embeddings = embeddings.cpu().to(torch.float32)
+                    else:
+                        embeddings = embeddings.cpu()
+                all_embeddings.extend(embeddings)
+        all_embeddings = [all_embeddings[idx] for idx in np.argsort(length_sorted_idx)]
+        if convert_to_tensor:
+            all_embeddings = torch.stack(all_embeddings)
+        elif convert_to_numpy:
+            all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
+
+        if input_is_string:
+            all_embeddings = all_embeddings[0]
+        return all_embeddings