[2025.06.10] adjust directory structure

src_dev/src_inference/dolly_inference.py

@@ -0,0 +1,39 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+model_name = ""
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
+
+
+prompt = "Give me a short introduction to large language model."
+messages = [
+    {"role": "user", "content": prompt}
+]
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+    enable_thinking=True 
+)
+model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+generated_ids = model.generate(
+    **model_inputs,
+    max_new_tokens=32768
+)
+output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
+
+
+try:
+    index = len(output_ids) - output_ids[::-1].index(151668)
+except ValueError:
+    index = 0
+
+thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
+content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
+
+print("thinking content:", thinking_content)
+print("content:", content)

src_dev/src_llm/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_dolly import DollyConfig
+from .modeling_dolly import DollyModel, DollyPreTrainedModel

src_dev/src_llm/configuration_dolly.py

@@ -0,0 +1,87 @@
+from transformers import PretrainedConfig
+from transformers import logging
+
+logger = logging.get_logger(__name__)
+
+class DollyConfig(PretrainedConfig):
+    model_type = "dolly"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # 添加张量并行配置
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",  # self-attention 的查询投影使用列并行
+        "layers.*.self_attn.k_proj": "colwise",  # self-attention 的键投影使用列并行
+        "layers.*.self_attn.v_proj": "colwise",  # self-attention 的值投影使用列并行
+        "layers.*.self_attn.o_proj": "rowwise",  # self-attention 的输出投影使用行并行
+        "layers.*.mlp.gate_proj": "colwise",    # MLP 的门投影使用列并行
+        "layers.*.mlp.up_proj": "colwise",      # MLP 的上升投影使用列并行
+        "layers.*.mlp.down_proj": "rowwise",    # MLP 的下降投影使用行并行
+    }
+
+    # 添加模型并行配置
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),  # 输入 ID 使用 embed_tokens 层生成嵌入
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),  # layers 层的输入输出
+        "norm": (["hidden_states"], ["hidden_states"]),  # 标准化层的输入输出
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=1024,
+        intermediate_size=3072,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=8,
+        head_dim=128,
+        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=1000000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        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  # we check `use_sliding_window` in the modeling code
+        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.head_dim = head_dim
+        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.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["DollyConfig"]

src_dev/src_llm/modeling_dolly.py

@@ -0,0 +1,916 @@
+from .configuration_dolly import DollyConfig
+from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS
+from functools import partial
+from typing import Callable, Optional, Tuple, Union
+from transformers.activations import ACT2FN
+from transformers.modeling_rope_utils import dynamic_rope_update, ROPE_INIT_FUNCTIONS
+import torch
+from torch import nn
+from transformers.utils import can_return_tuple
+from transformers.generation import GenerationMixin
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils import logging,LossKwargs,add_code_sample_docstrings,add_start_docstrings, add_start_docstrings_to_model_forward
+from transformers.processing_utils import Unpack
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast,QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast, TokenClassifierOutput
+from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
+logger = logging.get_logger(__name__)
+class DollyRMSNorm(nn.Module):
+    '''
+    输入的形状为batch_size*heads*seq_len*heads_dim,
+    仅在最后一个维度进行归一化，
+    参数量为heads_dim
+    '''
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        #添加一个非零值，防止除零
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        #x=(x^2)/n
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        #归一化并缩放
+        hidden_states = (hidden_states * torch.rsqrt(variance + self.variance_epsilon)) * self.weight
+        return hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+    
+class DollyMLP(nn.Module):
+    '''
+    DollyMLP层的实现, code from Qwen3
+    可训练参数量为: hidden_size * intermediate_size + hidden_size * intermediate_size + intermediate_size * hidden_size
+    '''
+    #只传入需要的三个参数
+    def __init__(self, hidden_size,intermediate_size,hidden_act):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        #可训练参数量 hidden_size * intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        #可训练参数量 hidden_size * intermediate_size
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        #可训练参数量 intermediate_size * hidden_size
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        #默认激活函数SiLU，=Sigmoid*x,输出范围(0, ∞)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        #通过激活函数后的向量[batch_size*seq_len*intermediate_size]和扩展后的向量进行元素乘
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+class DollyRotaryEmbedding(nn.Module):
+    def __init__(self, config: DollyConfig, device=None):
+        super().__init__()
+
+        #获取"rope_type"的值，如果没有使用"default"
+        self.config = config
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        # self.rope_type = getattr(config, "rope_type", "default")
+        self.original_max_seq_len = config.max_position_embeddings
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        #得到频率倒数,基础版本公式如下：
+        #inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2) / dim))
+        #得到一个长度为attention_dim/2的小数序列，如attention_dim=128,inv_freq=64
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        #注册缓存区
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        
+    @torch.no_grad()
+    @dynamic_rope_update
+    def forward(self, x, position_ids):
+        #将inv_freq进行维度扩展首先扩展为1*64*1，然后扩展为batch_size*64*1
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        #position_ids进行扩展batch_size* 1*seq_len
+        position_ids_expanded = position_ids[:, None, :].float()
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        #强制float32，禁止混合精度
+        with torch.autocast(device_type=device_type, enabled=False):  
+            #矩阵乘以后，形状变为batch_size*64*seq_len后，再转为batch_size*seq_len*64
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            #沿64的维度进行拼接batch_size*seq_len*128
+            emb = torch.cat((freqs, freqs), dim=-1)
+            #emb的cos()部分乘缩放系数:batch_size*seq_len*128
+            cos = emb.cos() * self.attention_scaling
+            #emb的sin()部分乘缩放系数:batch_size*seq_len*128
+            sin = emb.sin() * self.attention_scaling
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    #x分为两段，各取batch_size*heads_num*seq_len*64
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    #下部分*-1,按下+上重新组合x
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    '''
+    对q和k分别使用RoPE
+    '''
+    #由batch_size*seq_len*128增加一个维度变为batch_size*1*seq_len*128
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    #q:batch_size*heads_num*seq_len*128和sin及con进行元素乘batch_size*1*seq_len*128
+    #cos部分直接乘，sin部分进行位置变化后再乘，最后相加完成位置编码的嵌入
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    '''
+    在GQA中, k和v的heads数量小于q的heads数量, 
+    所以在进行attention计算前需要进行repeat_kv操作, 
+    保持第2维度的数量相等,n=config.num_attention_heads // config.num_key_value_heads
+    '''
+    #GQA中num_key_value_heads的数量小于attention heads数量
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    #在第3维度[index=2]进行扩充
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    #repeat kv，确保形状和query一致
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+    #计算attention score
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    #casual attention
+    if attention_mask is not None:
+        #causal_mask应当为下三角为0，上三角为负无穷
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+    # sofmax计算
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    # 训练时给attn_weights的随机位置置零
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    #乘value过程
+    attn_output = torch.matmul(attn_weights, value_states)
+    #最后转置，形状变为batch_size*seq_len*num_heads*head_dim
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class DollyAttention(nn.Module):
+    """
+    Attention in transformer
+    """
+    def __init__(self, config: DollyConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        #num_key_value_groups即repeat kv的n
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        #attention socre的缩放系数
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        #Q权重矩阵
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        #K权重矩阵
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        #V权重矩阵
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        #Attention的最后一步，线性连接层
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        #给输入q和k进行归一化操作，使用RMSNom，注意仅在最后一个维度进行归一化，参数量为head_dim
+        self.q_norm = DollyRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = DollyRMSNorm(self.head_dim, eps=config.rms_norm_eps) 
+        #支持滑动窗口
+        self.sliding_window = config.sliding_window
+        if not (
+            self.config.use_sliding_window
+            and getattr(self.config, "sliding_window", None) is not None
+            and self.layer_idx >= self.config.max_window_layers
+        ):
+            self.sliding_window = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        #input_shape: batch_size* seq_len
+        input_shape = hidden_states.shape[:-1]
+        #hidden_shape: batch_size* seq_len*-1* head_dim
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        #将输入经过Q矩阵变为query，然后调整形状为batch_size* seq_len*-1* head_dim，然后进行归一化，再调整形状batch_size* -1*seq_len* head_dim
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        #将输入经过K矩阵变为key，然后调整形状为batch_size* seq_len*-1* head_dim，然后进行归一化，再调整形状batch_size* -1*seq_len* head_dim
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        #将输入经过V矩阵变为value，然后调整形状为batch_size* seq_len*-1* head_dim，再调整形状batch_size* -1*seq_len* head_dim
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        #得到嵌入旋转位置编码信息后的query和key
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        
+        if past_key_value is not None:
+            # KVCache的cache更新位置信息和kv值
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
+                logger.warning_once(
+                    "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+                    'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+                )
+            else:
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        #attention的具体计算，返回结果和中间的attention score
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,  
+            **kwargs,
+        )
+        #将batch_size*seq_len*num_heads*head_dim 变为 batch_size*seq_len*hiddent_size
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        #attention的线性层
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+class DollyDecoderLayer(nn.Module):
+    def __init__(self, config: DollyConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = DollyAttention(config=config, layer_idx=layer_idx)
+        self.mlp = DollyMLP(config.hidden_size, config.intermediate_size, config.hidden_act)
+        self.input_layernorm = DollyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = DollyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        #仅flash_attention_2支持滑动窗口
+        if config.sliding_window and config._attn_implementation != "flash_attention_2":
+           config.sliding_window = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, 
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        #保留最开始的值，后面做残差网络
+        residual = hidden_states
+        #step1:归一化
+        hidden_states = self.input_layernorm(hidden_states)
+
+        #step2: Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        #step3:残差
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        #step4:归一化
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        #step5:MLP
+        hidden_states = self.mlp(hidden_states)
+        #step6:残差
+        hidden_states = residual + hidden_states
+        outputs = (hidden_states,)
+        #是否输出attention weights
+        if output_attentions:
+            outputs += (self_attn_weights,)
+        return outputs
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+class DollyPreTrainedModel(PreTrainedModel):
+    config_class = DollyConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["DollyDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            #初始线性层化权重参数
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            #初始化embedding层参数
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+class DollyModel(DollyPreTrainedModel):
+    def __init__(self, config: DollyConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        #embedding层
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        #transformer层
+        self.layers = nn.ModuleList(
+            [DollyDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        #RMSNorm层
+        self.norm = DollyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        #RoPE层
+        self.rotary_emb = DollyRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        #进行模型初始化
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutputWithPast:
+        #是否输出attention的计算结果
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        #是否输出隐藏层状态
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        #是否使用cache
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if not isinstance(past_key_values, (type(None), Cache)):
+            raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        )
+
+        hidden_states = inputs_embeds
+
+        #共享RoPE位置编码
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # 否是输出hidden state
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        #transformer的计算：
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            #如果启用了梯度检查点，并且处于训练模式下，模型将调用 self._gradient_checkpointing_func，
+            # 该函数使用 partial 函数来计算并重新计算部分前向过程，从而节省内存。
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    partial(decoder_layer.__call__, **flash_attn_kwargs),
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                    position_embeddings,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    position_embeddings=position_embeddings,
+                    **flash_attn_kwargs,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+        #归一化
+        hidden_states = self.norm(hidden_states)
+
+        # 添加最后一层hidden_states
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        #将结果打包为BaseModelOutputWithPast数据结构
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    def _update_causal_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_tensor: torch.Tensor,
+        cache_position: torch.Tensor,
+        past_key_values: Cache,
+        output_attentions: bool = False,
+    ):
+        #处理batch生成时padding问题
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and past_key_values is not None:
+                is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0]
+                if is_padding_right:
+                    raise ValueError(
+                        "You are attempting to perform batched generation with padding_side='right'"
+                        " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
+                    )
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+
+        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
+        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
+        # to infer the attention mask.
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        using_static_cache = isinstance(past_key_values, StaticCache)
+        using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)
+
+        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+        if (
+            self.config._attn_implementation == "sdpa"
+            and not (using_static_cache or using_sliding_window_cache)
+            and not output_attentions
+        ):
+            if AttentionMaskConverter._ignore_causal_mask_sdpa(
+                attention_mask,
+                inputs_embeds=input_tensor,
+                past_key_values_length=past_seen_tokens,
+                sliding_window=self.config.sliding_window,
+                is_training=self.training,
+            ):
+                return None
+
+        dtype, device = input_tensor.dtype, input_tensor.device
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        # SlidingWindowCache or StaticCache
+        if using_sliding_window_cache or using_static_cache:
+            target_length = past_key_values.get_max_cache_shape()
+        # DynamicCache or no cache
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else past_seen_tokens + sequence_length + 1
+            )
+
+        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+            attention_mask,
+            sequence_length=sequence_length,
+            target_length=target_length,
+            dtype=dtype,
+            device=device,
+            cache_position=cache_position,
+            batch_size=input_tensor.shape[0],
+            config=self.config,
+            past_key_values=past_key_values,
+        )
+
+        if (
+            self.config._attn_implementation == "sdpa"
+            and attention_mask is not None
+            and attention_mask.device.type in ["cuda", "xpu"]
+            and not output_attentions
+        ):
+            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+        return causal_mask
+
+    @staticmethod
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        config: DollyConfig,
+        past_key_values: Cache,
+    ):
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            causal_mask = attention_mask
+        else:
+            min_dtype = torch.finfo(dtype).min
+            causal_mask = torch.full(
+                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
+            )
+            diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+            if config.sliding_window is not None:
+                # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
+                # the check is needed to verify is current checkpoint was trained with sliding window or not
+                if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:
+                    sliding_attend_mask = torch.arange(target_length, device=device) <= (
+                        cache_position.reshape(-1, 1) - config.sliding_window
+                    )
+                    diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
+            causal_mask *= diagonal_attend_mask
+            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+            if attention_mask is not None:
+                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+                if attention_mask.shape[-1] > target_length:
+                    attention_mask = attention_mask[:, :target_length]
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+                    causal_mask.device
+                )
+                padding_mask = padding_mask == 0
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    padding_mask, min_dtype
+                )
+        return causal_mask
+    
+
+class DollyForCausalLM(DollyPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = DollyModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @can_return_tuple
+    @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[KwargsForCausalLM],
+    ) -> CausalLMOutputWithPast:
+       
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class DollyForSequenceClassification(DollyPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = DollyModel(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+    ) -> SequenceClassifierOutputWithPast:
+
+        transformer_outputs: BaseModelOutputWithPast = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        hidden_states = transformer_outputs.last_hidden_state
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            last_non_pad_token = -1
+        elif input_ids is not None:
+            # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+            token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+        else:
+            last_non_pad_token = -1
+            logger.warning_once(
+                f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+                "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+            )
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
+class DollyForTokenClassification(DollyPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = DollyModel(config)
+        if getattr(config, "classifier_dropout", None) is not None:
+            classifier_dropout = config.classifier_dropout
+        elif getattr(config, "hidden_dropout", None) is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.score = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+    ) -> TokenClassifierOutput:
+       
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        sequence_output = outputs.last_hidden_state
+        sequence_output = self.dropout(sequence_output)
+        logits = self.score(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.config)
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class DollyForQuestionAnswering(DollyPreTrainedModel):
+    base_model_prefix = "transformer"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = DollyModel(config)
+        self.qa_outputs = nn.Linear(config.hidden_size, 2)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.transformer.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.transformer.embed_tokens = value
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        start_positions: Optional[torch.LongTensor] = None,
+        end_positions: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> QuestionAnsweringModelOutput:
+
+        outputs: BaseModelOutputWithPast = self.transformer(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+        sequence_output = outputs.last_hidden_state
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        loss = None
+        if start_positions is not None and end_positions is not None:
+            loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs)
+
+        return QuestionAnsweringModelOutput(
+            loss=loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "DollyForCausalLM",
+    "DollyForQuestionAnswering",
+    "DollyModel",
+    "DollyPreTrainedModel",
+    "DollyForSequenceClassification",
+    "DollyForTokenClassification",
+]

src_dev/src_llm_moe/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_dollyMoE import DollyMoEConfig
+from .modeling_dollyMoE import DollyMoEModel, DollyMoEPreTrainedModel

src_dev/src_llm_moe/configuration_dollyMoE.py

@@ -0,0 +1,95 @@
+from transformers import PretrainedConfig
+from transformers import logging
+from transformers.modeling_rope_utils import rope_config_validation
+logger = logging.get_logger(__name__)
+
+
+class DollyMoEConfig(PretrainedConfig):
+    model_type = "dolly_moe"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.experts.*.gate_proj": "colwise",
+        "layers.*.mlp.experts.*.up_proj": "colwise",
+        "layers.*.mlp.experts.*.down_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=1024,
+        intermediate_size=3072,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=8,
+        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=1000000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        attention_dropout=0.0,
+        decoder_sparse_step=1,
+        moe_intermediate_size=768,
+        num_experts_per_tok=4,
+        num_experts=8,
+        norm_topk_prob=False,
+        output_router_logits=False,
+        router_aux_loss_coef=0.001,
+        mlp_only_layers=None,
+        **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
+
+        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.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+        self.decoder_sparse_step = decoder_sparse_step
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_experts = num_experts
+        self.norm_topk_prob = norm_topk_prob
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+__all__ = ["DollyMoEConfig"]

src_dev/src_pretrain/pretrain_deepspeed.py

@@ -0,0 +1,120 @@
+import torch
+from src_llm.modeling_dolly import DollyForCausalLM
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+    AutoTokenizer,
+    set_seed
+)
+from datasets import load_dataset
+from accelerate import Accelerator
+import os
+import logging
+import glob
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+accelerator = Accelerator()
+set_seed(42)
+
+model_name = "/mnt/han.luzhi/dolly_llm/weights_llm"  
+resume_option = None                     
+# resume_option = True                    
+# resume_option = "/mnt/han.luzhi/dolly_llm/checkpoints_ds/checkpoint-500"  
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+if isinstance(resume_option, str) and os.path.exists(resume_option):
+    logger.info(f"从指定检查点 {resume_option} 加载模型...")
+    model = DollyForCausalLM.from_pretrained(resume_option)
+else:
+    logger.info("未找到检查点，从预训练模型加载...")
+    logger.info(f"从预训练模型 {model_name} 加载...")
+    model = DollyForCausalLM.from_pretrained(model_name)
+model.gradient_checkpointing_enable()
+
+data_files = glob.glob("/mnt/han.luzhi/dolly_llm/corpus/*.jsonl")
+logger.info(f"找到训练语料: {data_files}")
+
+dataset = load_dataset('json', data_files=data_files, split='train')
+# dataset = load_dataset('text', data_files={'train': '/mnt/han.luzhi/dolly_llm/corpus/wikipedia.txt'})
+
+def preprocess_function(examples):
+    return tokenizer(
+        examples['text'],
+        truncation=True,
+        max_length=512,
+        padding='max_length',
+        return_tensors='pt'
+    )
+
+tokenized_datasets = dataset.map(
+    preprocess_function,
+    batched=True,
+    # remove_columns=['text'], 
+    remove_columns=dataset.column_names,
+    num_proc=16  
+)
+
+data_collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer,
+    mlm=False 
+)
+
+training_args = TrainingArguments(
+    output_dir="/mnt/han.luzhi/dolly_llm/checkpoints_ds",
+    deepspeed="/mnt/han.luzhi/dolly_llm/deepspeed_config/ds_stage2.json",  
+    do_eval=False,
+    save_strategy="steps",
+    save_steps=500,
+    learning_rate=1e-4,
+    per_device_train_batch_size=36,
+    num_train_epochs=2,
+    logging_dir="/mnt/han.luzhi/dolly_llm/logs",
+    logging_steps=20,
+    save_total_limit=2,
+    fp16=True,  
+    dataloader_num_workers=3,
+    gradient_accumulation_steps=2,
+    report_to="tensorboard",
+    optim="adamw_torch",
+    lr_scheduler_type="cosine",
+    warmup_steps=500,
+    weight_decay=0.01,
+    local_rank=int(os.environ.get("LOCAL_RANK", 0)),
+    disable_tqdm=False,
+    logging_first_step=True,
+    log_level="info",
+    ddp_find_unused_parameters=False,
+    overwrite_output_dir=True,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets,
+    processing_class=tokenizer, 
+    data_collator=data_collator,
+)
+
+logger.info("***** 开始训练 *****")
+logger.info(f" 进程数 = {accelerator.num_processes}")
+logger.info(f" 设备 = {accelerator.device}")
+logger.info(f" 总批大小 = {training_args.per_device_train_batch_size * accelerator.num_processes * training_args.gradient_accumulation_steps}")
+
+if resume_option !=None:
+    trainer.train(resume_from_checkpoint=resume_option)
+else:
+    trainer.train()
+
+
+if accelerator.is_main_process: 
+    trainer.save_model("/mnt/han.luzhi/dolly_llm/result/final")
+    tokenizer.save_pretrained("/mnt/han.luzhi/dolly_llm/result/final")
+    logger.info("***** 训练完成，模型已保存 *****")
+
+##  deepspeed --num_gpus=4 --master_port=12345 pretrain_transformers.py
+#   deepspeed --include="localhost:1,2,3" --master_port=12345 pretrain_transformers.py

src_dev/src_pretrain/pretrain_deepspeed_MoE.py

@@ -0,0 +1,120 @@
+import torch
+from src_llm_moe.modeling_dollyMoE import DollyMoEForCausalLM
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+    AutoTokenizer,
+    set_seed
+)
+from datasets import load_dataset
+from accelerate import Accelerator
+import os
+import logging
+import glob
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+accelerator = Accelerator()
+set_seed(42)
+
+model_name = "/mnt/han.luzhi/dolly/weights_MoE"  
+resume_option = None                     
+# resume_option = True                    
+# resume_option = "/mnt/han.luzhi/dolly_llm/checkpoints_ds/checkpoint-500"  
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+if isinstance(resume_option, str) and os.path.exists(resume_option):
+    logger.info(f"从指定检查点 {resume_option} 加载模型...")
+    model = DollyMoEForCausalLM.from_pretrained(resume_option)
+else:
+    logger.info("未找到检查点，从预训练模型加载...")
+    logger.info(f"从预训练模型 {model_name} 加载...")
+    model = DollyMoEForCausalLM.from_pretrained(model_name)
+model.gradient_checkpointing_enable()
+
+data_files = glob.glob("/mnt/han.luzhi/dolly/corpus/*.jsonl")
+logger.info(f"找到训练语料: {data_files}")
+
+dataset = load_dataset('json', data_files=data_files, split='train')
+# dataset = load_dataset('text', data_files={'train': '/mnt/han.luzhi/dolly_llm/corpus/wikipedia.txt'})
+
+def preprocess_function(examples):
+    return tokenizer(
+        examples['text'],
+        truncation=True,
+        max_length=512,
+        padding='max_length',
+        return_tensors='pt'
+    )
+
+tokenized_datasets = dataset.map(
+    preprocess_function,
+    batched=True,
+    # remove_columns=['text'], 
+    remove_columns=dataset.column_names,
+    num_proc=16  
+)
+
+data_collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer,
+    mlm=False 
+)
+
+training_args = TrainingArguments(
+    output_dir="/mnt/han.luzhi/dolly/checkpoints_MoE",
+    deepspeed="/mnt/han.luzhi/dolly/deepspeed_config/ds_stage2.json",  
+    do_eval=False,
+    save_strategy="steps",
+    save_steps=500,
+    learning_rate=1e-4,
+    per_device_train_batch_size=16,
+    num_train_epochs=2,
+    logging_dir="/mnt/han.luzhi/dolly/logs",
+    logging_steps=20,
+    save_total_limit=2,
+    fp16=True,  
+    dataloader_num_workers=3,
+    gradient_accumulation_steps=2,
+    report_to="tensorboard",
+    optim="adamw_torch",
+    lr_scheduler_type="cosine",
+    warmup_steps=500,
+    weight_decay=0.01,
+    local_rank=int(os.environ.get("LOCAL_RANK", 0)),
+    disable_tqdm=False,
+    logging_first_step=True,
+    log_level="info",
+    ddp_find_unused_parameters=False,
+    overwrite_output_dir=True,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets,
+    processing_class=tokenizer, 
+    data_collator=data_collator,
+)
+
+logger.info("***** 开始训练 *****")
+logger.info(f" 进程数 = {accelerator.num_processes}")
+logger.info(f" 设备 = {accelerator.device}")
+logger.info(f" 总批大小 = {training_args.per_device_train_batch_size * accelerator.num_processes * training_args.gradient_accumulation_steps}")
+
+if resume_option !=None:
+    trainer.train(resume_from_checkpoint=resume_option)
+else:
+    trainer.train()
+
+
+if accelerator.is_main_process: 
+    trainer.save_model("/mnt/han.luzhi/dolly/chechpoints_MoE/final")
+    tokenizer.save_pretrained("/mnt/han.luzhi/dolly/chechpoints_MoE/final")
+    logger.info("***** 训练完成，模型已保存 *****")
+
+## deepspeed --num_gpus=4 --master_port=12345 pretrain_transformers.py
+#  CUDA_VISIBLE_DEVICES=1,2,3 deepspeed --include="localhost:1,2,3" --master_port=12345 pretrain_transformers.py

src_dev/src_pretrain/pretrain_transformers.py

@@ -0,0 +1,109 @@
+import torch
+from src_llm.modeling_dolly import DollyForCausalLM
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+    AutoTokenizer,
+    set_seed
+)
+from datasets import load_dataset
+from accelerate import Accelerator
+import os
+import logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+accelerator = Accelerator()
+set_seed(42)
+
+model_name = "/mnt/han.luzhi/dolly_llm/weights_llm"  
+# resume_option = None                     
+# resume_option = True                    
+resume_option = "/mnt/han.luzhi/dolly_llm/result/checkpoint-59000"  
+# resume_option = "/mnt/han.luzhi/dolly_llm/result/final" 
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+if isinstance(resume_option, str) and os.path.exists(resume_option):
+    logger.info(f"从指定检查点 {resume_option} 加载模型...")
+    model = DollyForCausalLM.from_pretrained(resume_option)
+else:
+    logger.info("未找到检查点，从预训练模型加载...")
+    logger.info(f"从预训练模型 {model_name} 加载...")
+    model = DollyForCausalLM.from_pretrained(model_name)
+
+model.gradient_checkpointing_enable()
+
+dataset = load_dataset('text', data_files={'train': '/mnt/han.luzhi/dolly_llm/corpus/wikipedia.txt'})
+
+def preprocess_function(examples):
+    return tokenizer(
+        examples['text'],
+        truncation=True,
+        max_length=512,
+        padding='max_length',
+        return_tensors='pt'
+    )
+
+tokenized_datasets = dataset.map(
+    preprocess_function,
+    batched=True,
+    remove_columns=['text'], 
+    num_proc=16  
+)
+
+data_collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer,
+    mlm=False 
+)
+
+training_args = TrainingArguments(
+    output_dir="/mnt/han.luzhi/dolly_llm/result",
+    do_eval=False,  
+    save_strategy="steps", 
+    save_steps=500,
+    learning_rate=5e-5,
+    per_device_train_batch_size=26,
+    num_train_epochs=5,
+    logging_dir="/mnt/han.luzhi/dolly_llm/logs",
+    logging_steps=10,
+    save_total_limit=2,
+    fp16=True, 
+    dataloader_num_workers=3,
+    gradient_accumulation_steps=2, 
+    report_to="tensorboard",  
+    optim="adamw_torch",  
+    lr_scheduler_type="cosine", 
+    warmup_steps=500, 
+    weight_decay=0.01,
+    local_rank=int(os.environ.get("LOCAL_RANK", 0)),  
+    disable_tqdm=False, 
+    logging_first_step=True,  
+    log_level="info",  
+    log_level_replica="info",  
+    ddp_find_unused_parameters=False, 
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets['train'],
+    processing_class=tokenizer, 
+    data_collator=data_collator,
+)
+
+logger.info("***** 开始训练 *****")
+logger.info(f" 进程数 = {accelerator.num_processes}")
+logger.info(f" 设备 = {accelerator.device}")
+logger.info(f" 总批大小 = {training_args.per_device_train_batch_size * accelerator.num_processes * training_args.gradient_accumulation_steps}")
+
+trainer.train(resume_from_checkpoint=resume_option)
+
+if accelerator.is_main_process: 
+    trainer.save_model("/mnt/han.luzhi/dolly_llm/result/final")
+    tokenizer.save_pretrained("/mnt/han.luzhi/dolly_llm/result/final")
+    logger.info("***** 训练完成，模型已保存 *****")
+
+##  python3 -m torch.distributed.launch --nproc_per_node=4 --master_port=12345 pretrain_transformers.py

src_dev/src_tokenizer/README.md

@@ -0,0 +1,3 @@
+# USAGE
+使用sentencepiece和huggingface transformers分别进行tokenizer的训练
+support BPE and BBPE

src_dev/src_tokenizer/hf_tokenizer_code/step1.1_hf_train_from_old_largefile.py

@@ -0,0 +1,23 @@
+from transformers import AutoTokenizer
+
+def get_training_corpus(file_path, batch_size=1000):
+    with open(file_path, "r", encoding="utf-8") as f:
+        batch = []
+        for line in f:
+            batch.append(line.strip())  
+            if len(batch) >= batch_size:
+                yield batch
+                batch = []
+        if batch:  
+            yield batch
+
+old_tokenizer = AutoTokenizer.from_pretrained("tokenizer_file/dolly_tokenizer_hf2")  
+training_corpus = get_training_corpus("corpus/wikipedia.txt")  
+
+tokenizer = old_tokenizer.train_new_from_iterator(
+    training_corpus,
+    vocab_size=34000
+)
+
+
+tokenizer.save_pretrained("tokenizer_file/new_tokenizer2")

src_dev/src_tokenizer/hf_tokenizer_code/step1_hf_train_from_old.py

@@ -0,0 +1,16 @@
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+dataset = load_dataset("text", data_files="corpus/wikipedia.txt", streaming=True)
+
+def batch_iterator(batch_size=1000):
+    for batch in dataset["train"].iter(batch_size):
+        yield batch["text"]
+
+old_tokenizer = AutoTokenizer.from_pretrained("tokenizer_file/dolly_tokenizer_hf2")
+tokenizer = old_tokenizer.train_new_from_iterator(
+    batch_iterator(),
+    vocab_size=33000,
+)
+
+tokenizer.save_pretrained("new_tokenizer")

src_dev/src_tokenizer/hf_tokenizer_code/step2_hf_train_from_scratch_BPEandBBPE.py

@@ -0,0 +1,149 @@
+import os
+import json
+from pathlib import Path
+from tokenizers import (
+    decoders,
+    models,
+    normalizers,
+    pre_tokenizers,
+    processors,
+    trainers,
+    Tokenizer,
+)
+from transformers import PreTrainedTokenizerFast
+import logging
+from tqdm import tqdm
+log_dir = "/mnt/hanluzhi/dolly/logs"
+os.makedirs(log_dir, exist_ok=True)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(), 
+        logging.FileHandler(os.path.join(log_dir, "training.log"))  
+    ]
+)
+
+def get_training_corpus():
+    for file_path in tqdm(jsonl_files, desc="processing files", unit="file"):
+        with open(file_path, "r", encoding="utf-8") as f:
+            for line in f:
+                data = json.loads(line)
+                if "text" in data:
+                    yield data["text"]
+
+
+corpus_dir = Path("./corpus/skypile_2020_head")
+jsonl_files = list(corpus_dir.glob("**/*.jsonl"))
+if not jsonl_files:
+    raise ValueError("No training files found")
+
+
+special_tokens = [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+]
+
+
+tokenizer = Tokenizer(models.BPE())
+
+
+tokenizer.normalizer = normalizers.Sequence([
+    normalizers.NFKC(),
+    normalizers.StripAccents(),
+])
+
+
+tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
+    [
+        pre_tokenizers.Split(r"[\r\n]", "isolated"),   
+        pre_tokenizers.Split(r"\s?[A-Za-zµÀ-ÖØ-öø-ƺƼ-ƿDŽ-ʓʕ-ʯͰ-ͳͶͷͻ-ͽͿΆΈ-ΊΌΎ-ΡΣ-ϵϷ-ҁҊ-ԯԱ-ՖႠ-ჅᎠ-Ᏽᏸ-ᏽᲐ-ᲺᲽ-Ჿᴀ-ᴫᵫ-ᵷᵹ-ᶚḀ-ἕἘ-Ἕἠ-ὅὈ-Ὅὐ-ὗὙὛὝὟ-ώᾀ-ᾴᾶ-ᾼιῂ-ῄῆ-ῌῐ-ΐῖ-Ίῠ-Ῥῲ-ῴῶ-ῼℂℇℊ-ℓℕℙ-ℝℤΩℨK-ℭℯ-ℴℹℼ-ℿⅅ-ⅉⅎↃↄⰀ-ⱻⱾ-ⳤⳫ-ⳮⳲⳳꙀ-ꙭꚀ-ꚛꜢ-ꝯꝱ-ꞇꞋ-ꞎꭰ-ꮿff-stﬓ-ﬗＡ-Ｚａ-ｚ𐐀-𐑏𐒰-𐓓𐓘-𐓻𐲀-𐲲𐳀-𐳲𑢠-𑣟𞤀-𞥃]+","isolated"), 
+        pre_tokenizers.Split(r"\s?[!-/:-~！-／：-～‘-‟　-。]+", "isolated"),   
+        pre_tokenizers.Split(r"\s+$", "isolated"), 
+        pre_tokenizers.Split("[一-龥ࠀ-一가-퟿]+","isolated"),
+        pre_tokenizers.Digits(individual_digits=True),
+        pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
+    ])
+
+# tokenizer.post_processor = None
+tokenizer.post_processor = processors.ByteLevel(use_regex=True, trim_offsets=False, add_prefix_space=True)
+
+tokenizer.decoder = decoders.ByteLevel(
+    use_regex=True,
+    trim_offsets=True,
+    add_prefix_space=False  
+)
+
+
+trainer = trainers.BpeTrainer(
+    vocab_size=12800,
+    special_tokens=special_tokens, 
+    show_progress=True  
+)
+
+logging.info("Training tokenizer...")
+tokenizer.train_from_iterator(
+    get_training_corpus(),
+    trainer=trainer,
+    length=sum(1 for _ in get_training_corpus())  
+)
+
+wrapped_tokenizer = PreTrainedTokenizerFast(
+    tokenizer_object=tokenizer,
+    bos_token=None,  
+    eos_token="<|endoftext|>",
+    pad_token="<|endoftext|>",
+    unk_token=None,  
+    additional_special_tokens=["<|im_start|>", "<|im_end|>"],
+    model_max_length=32768,  
+    clean_up_tokenization_spaces=False,  
+    errors="replace", 
+    split_special_tokens=False,  
+)
+logging.info("train done ")
+output_dir = Path("./llm_tokenizer")
+output_dir.mkdir(exist_ok=True, parents=True)
+
+wrapped_tokenizer.save_pretrained(output_dir)
+logging.info(f"Tokenizer saved to {output_dir}")
+
+QWEN2_CHAT_TEMPLATE = (
+    "{% for message in messages %}"
+    "{% if loop.first and messages[0]['role'] != 'system' %}"
+    "{{ '<|im_start|>system\nYou are a helpful assistant<|im_end|>\n' }}"
+    "{% endif %}"
+    "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}"
+    "{{ '<|im_start|>assistant\n' }}"
+    "{% endif %}"
+)
+
+wrapped_tokenizer.chat_template = QWEN2_CHAT_TEMPLATE
+
+wrapped_tokenizer.save_pretrained(output_dir)
+
+messages = [
+    {"role": "user", "content": "你好！"},
+    {"role": "assistant", "content": "你好，有什么可以帮您？"},
+    {"role": "user", "content": "Python怎么用？"}
+]
+
+formatted = wrapped_tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+logging.info("\n格式化后的对话:")
+logging.info(formatted)
+
+test_text = "深度学习是人工智能的一个重要分支。Deep learning is a subset of machine learning."
+logging.info("\nTest tokenization:")
+encoded = wrapped_tokenizer.tokenize(test_text)
+logging.info(encoded)
+logging.info(encoded)
+decoded_text = wrapped_tokenizer.decode(encoded) 
+logging.info(decoded_text)
+logging.info(decoded_text)

src_dev/src_tokenizer/hf_tokenizer_code/step3_hf_add_new_token.py

@@ -0,0 +1,10 @@
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("tokenizer_file/hf_tokenizer")
+
+new_tokens = ["<ACTION_1>", "<ACTION_2>"]
+tokenizer.add_tokens(new_tokens)
+
+text = "<ACTION_1>"
+print(tokenizer.tokenize(text))  
+tokenizer.save_pretrained("updated_tokenizer")

src_dev/src_tokenizer/spm_tokenizer_code/step1_spm_train_tokenizer.py

@@ -0,0 +1,118 @@
+import os
+import json
+from pathlib import Path
+import sentencepiece as spm
+import logging
+import tempfile
+import unicodedata
+from transformers import PreTrainedTokenizerFast
+from tqdm import tqdm
+logging.basicConfig(level=logging.INFO, format='%(message)s')
+
+def normalize_text(text):
+    text = text.replace('\x00', '')
+    text = unicodedata.normalize('NFKC', text)
+    text = ''.join(c for c in unicodedata.normalize('NFD', text)
+                  if unicodedata.category(c) != 'Mn')
+    return text
+
+corpus_dir = Path("./corpus/skypile_2020_head")
+jsonl_files = list(corpus_dir.glob("**/*.jsonl"))
+if not jsonl_files:
+    raise ValueError("No training files found")
+
+special_tokens = ["<|endoftext|>", "<|im_start|>", "<|im_end|>"]
+logging.info("Staring processing corpus")
+with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as temp_corpus:
+    temp_corpus_path = temp_corpus.name
+    file_progress = tqdm(jsonl_files, desc="处理文件中", unit="file")
+    for file_path in file_progress:
+        file_progress.set_description(f"处理: {file_path.name[:40]}...")
+        with open(file_path, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    data = json.loads(line)
+                    if "text" in data:
+                        normalized_text = normalize_text(data["text"])
+                        temp_corpus.write(normalized_text + "\n")
+                except json.JSONDecodeError as e:
+                    logging.warning(f"JSON 解析错误在 {file_path}: {e}")
+                except UnicodeDecodeError as e:
+                    logging.warning(f"编码错误在 {file_path}: {e}")
+logging.info("corpus processing completed")
+logging.info("Starting SentencePiece training")
+
+vocab_size = 32000  
+
+# character_coverage = 0.9995  
+
+spm.SentencePieceTrainer.train(
+    input=temp_corpus_path,
+    model_prefix='tokenizer_model',
+    vocab_size=vocab_size,  
+    model_type='bpe',
+    character_coverage=1.0,  
+    pad_id=-1,
+    bos_id=-1,
+    eos_id=-1,
+    unk_id=0,
+    user_defined_symbols=special_tokens,
+    split_digits=True,
+    byte_fallback=True,
+    remove_extra_whitespaces=False,
+    add_dummy_prefix=False,
+    normalization_rule_name='identity',
+    num_threads=os.cpu_count(),
+    max_sentencepiece_length=32,  
+    shuffle_input_sentence=True, 
+)
+
+os.unlink(temp_corpus_path)
+logging.info("SentencePiece training completed")
+
+tokenizer = PreTrainedTokenizerFast(
+    tokenizer_file='tokenizer_model.model',
+    bos_token=None,
+    eos_token="<|endoftext|>",
+    pad_token="<|endoftext|>",
+    unk_token=None,
+    additional_special_tokens=["<|im_start|>", "<|im_end|>"],
+    model_max_length=32768,
+    clean_up_tokenization_spaces=False,
+    split_special_tokens=False,
+    errors="replace"
+)
+
+output_dir = Path("./llm_tokenizer")
+output_dir.mkdir(exist_ok=True, parents=True)
+tokenizer.save_pretrained(output_dir)
+logging.info(f"Tokenizer saved to {output_dir}")
+
+QWEN2_CHAT_TEMPLATE = (
+    "{% for message in messages %}"
+    "{% if loop.first and messages[0]['role'] != 'system' %}"
+    "{{ '<|im_start|>system\nYou are a helpful assistant<|im_end|>\n' }}"
+    "{% endif %}"
+    "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}"
+    "{{ '<|im_start|>assistant\n' }}"
+    "{% endif %}"
+)
+
+tokenizer.chat_template = QWEN2_CHAT_TEMPLATE
+tokenizer.save_pretrained(output_dir)
+
+messages = [
+    {"role": "user", "content": "你好！"},
+    {"role": "assistant", "content": "你好，有什么可以帮您？"},
+    {"role": "user", "content": "Python怎么用？"}
+]
+
+formatted = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+print("\n格式化后的对话:")
+print(formatted)

src_dev/src_tokenizer/spm_tokenizer_code/step2_spm_merge_tokenizer.py

@@ -0,0 +1,33 @@
+import sentencepiece as spm
+from sentencepiece import sentencepiece_model_pb2 as sp_model_pb2
+
+model_file = "tokenizer_part/dolly_tokenizer_spm/tokenizer_zh.model"
+en_model_file = "tokenizer_part/dolly_tokenizer_spm/tokenizer_en.model"
+
+model_file_new = "tokenizer_part/dolly_tokenizer_spm/tokenizer.model"
+
+chinese_sp = spm.SentencePieceProcessor(model_file=model_file)
+en_sp = spm.SentencePieceProcessor(model_file=en_model_file)
+
+print(chinese_sp.vocab_size())#32000
+
+chinese_spm = sp_model_pb2.ModelProto()
+chinese_spm.ParseFromString(chinese_sp.serialized_model_proto())
+
+en_spm = sp_model_pb2.ModelProto()
+en_spm.ParseFromString(en_sp.serialized_model_proto())
+
+chinese_sp_set = set(p.piece for p in chinese_spm.pieces)
+
+for p in en_spm.pieces:
+    piece = p.piece
+    if piece not in chinese_sp_set:
+        chinese_sp_set.add(piece)
+        new_piece = sp_model_pb2.ModelProto().SentencePiece()
+        new_piece.piece = piece
+        new_piece.score = 0
+        chinese_spm.pieces.append(new_piece)
+
+print(f"after expand vocab={len(chinese_spm.pieces)}") # 32361,增加361个
+with open(model_file_new, 'wb') as file:
+    file.write(chinese_spm.SerializeToString())

src_dev/src_tokenizer/spm_tokenizer_code/step3_spm_to_hf.py

@@ -0,0 +1,18 @@
+from transformers import AutoTokenizer,XLMRobertaTokenizer
+#由spm转为hf_tokenizer
+tokenizer = XLMRobertaTokenizer.from_pretrained("tokenizer_file/dolly_tokenizer_spm/tokenizer.model")
+text = "A Frenchman from Paris"
+encoded_input = tokenizer.encode(text)
+print(encoded_input)
+decoded_output = tokenizer.decode(encoded_input)
+print(decoded_output)
+tokenizer.save_pretrained('tokenizer_file/dolly_tokenizer_hf')
+
+#转换为常见的tokenzier.json格式
+tokenizer = AutoTokenizer.from_pretrained("tokenizer_file/dolly_tokenizer_hf")
+text = "A Frenchman from Paris"
+encoded_input = tokenizer.encode(text)
+print(encoded_input)
+decoded_output = tokenizer.decode(encoded_input)
+print(decoded_output)
+tokenizer.save_pretrained('tokenizer_file/dolly_tokenizer_hf2')

src_dev/src_tokenizer/spm_tokenizer_code/step4_spm_expand_vocab.py

@@ -0,0 +1,21 @@
+import sentencepiece as spm
+from sentencepiece import sentencepiece_model_pb2 as sp_model_pb2
+
+sp = spm.SentencePieceProcessor(model_file='tokenizer_file/dolly_tokenizer_spm/tokenizer.model')
+print(sp.vocab_size())  # 32461
+
+old_model = sp_model_pb2.ModelProto()
+old_model.ParseFromString(sp.serialized_model_proto())
+pieces = ['拉布拉多', '货拉拉', '拉不拉']
+existed_pieces = set(p.piece for p in old_model.pieces)
+for piece in pieces:
+    if piece not in existed_pieces:
+        existed_pieces.add(piece)
+        new_piece = sp_model_pb2.ModelProto().SentencePiece()
+        new_piece.piece = piece
+        new_piece.score = 0
+        old_model.pieces.append(new_piece)
+    
+print(len(old_model.pieces))  # 32461
+with open('tokenizer_file/dolly_tokenizer_spm/expand_tokenizer.model', 'wb') as f:
+    f.write(old_model.SerializeToString())

src_dev/src_tokenizer/spm_tokenizer_code/step5_test_spm_tokenizer.py

@@ -0,0 +1,15 @@
+import sentencepiece as spm
+sp = spm.SentencePieceProcessor(model_file='tokenizer_part/dolly-tokenizer/tokenizer.model')
+
+sp_en= spm.SentencePieceProcessor(model_file='tokenizer_part/dolly-tokenizer/tokenizer_en.model')
+sentence = 'A Frenchman from Paris'
+pieces = sp.encode_as_pieces(sentence)
+pieces_en = sp_en.encode_as_pieces(sentence)
+print(pieces)  # ['▁A', '▁F', 'ren', 'ch', 'man', '▁from', '▁P', 'ar', 'is']
+print(pieces_en) #'▁A', '▁F', 'ren', 'ch', 'man', '▁from', '▁P', 'aris'
+tokens = sp.encode(sentence)
+tokens_en = sp_en.encode(sentence)
+print(tokens)  # [338, 493, 587, 349, 807, 429, 373, 279, 282]
+print(tokens_en) #364, 490, 5981, 398, 2036, 32046, 386, 9020] 
+# 解码后的结果应等于编码前的输入
+# assert sp.decode(tokens) == sentence