[2025.06.17] update train code

src/configuration_dolly.py

@@ -0,0 +1,87 @@
+from transformers import PretrainedConfig
+from transformers import logging
+from transformers.modeling_rope_utils import rope_config_validation
+
+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=32000,
+        hidden_size=2048,
+        intermediate_size=6144,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=8,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=40960,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=True,
+        rope_theta=1000000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        use_sliding_window=False,
+        sliding_window=None,
+        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
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["DollyConfig"]

src/modeling_dolly.py

@@ -0,0 +1,925 @@
+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
+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)
+        return self.weight*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.max_seq_len_cached = config.max_position_embeddings
+        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)
+        self.original_inv_freq = self.inv_freq
+        
+    @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"
+        ): 
+            logger.warning_once(
+                f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
+                "unexpected results may be encountered."
+            )
+
+    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_()
+        elif isinstance(module, DollyRMSNorm):
+            module.weight.data.fill_(1.0)
+
+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/step0_train_tokenizer_jsonl.py

@@ -1,9 +1,9 @@
 import json
 import glob
-from tokenizers import Tokenizer
+from tokenizers import Tokenizer,Tokenizer, normalizers, pre_tokenizers,  processors, decoders
+from tokenizers.normalizers import NFC
 from tokenizers.models import BPE
 from tokenizers.trainers import BpeTrainer
-from tokenizers.pre_tokenizers import ByteLevel
 from transformers import PreTrainedTokenizerFast
 
 SPECIAL_TOKENS = [
@@ -65,17 +65,50 @@ def read_corpus(files, batch_size=1000):
             if batch:
                 yield batch
 
-tokenizer = Tokenizer(BPE(unk_token=None))
+tokenizer = Tokenizer(BPE(
-tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
+    unk_token=None,
+    continuing_subword_prefix="",
+    end_of_word_suffix="",
+    fuse_unk=False,
+    byte_fallback=False
+))
+tokenizer.normalizer = normalizers.Sequence([NFC()])
+regex_pattern = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+"""
+tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
+    pre_tokenizers.Split(
+        pattern=pre_tokenizers.Regex(regex_pattern),
+        behavior="isolated",
+        invert=False
+    ),
+    pre_tokenizers.ByteLevel(
+        add_prefix_space=False,
+        trim_offsets=False,
+        use_regex=False
+    )
+])
 
-trainer = BpeTrainer(
+tokenizer.post_processor = processors.ByteLevel(
-    vocab_size=151643,  
+    add_prefix_space=False,
-    min_frequency=2,    
+    trim_offsets=False,
-    special_tokens=SPECIAL_TOKENS,
+    use_regex=False
-    show_progress=True
 )
 
-corpus_files = glob.glob("corpus/*.jsonl")  
+tokenizer.decoder = decoders.ByteLevel(
+    add_prefix_space=False,
+    trim_offsets=False,
+    use_regex=False
+)
+
+trainer = BpeTrainer(
+    vocab_size=32000,  
+    min_frequency=1,    
+    special_tokens=SPECIAL_TOKENS,
+    show_progress=True,
+    continuing_subword_prefix="",
+    end_of_word_suffix=""
+)
+
+corpus_files = glob.glob("../corpus/*.jsonl")  
 tokenizer.train_from_iterator(
     iterator=read_corpus(corpus_files),
     trainer=trainer,
@@ -91,7 +124,7 @@ fast_tokenizer = PreTrainedTokenizerFast(
     add_prefix_space=False,
     clean_up_tokenization_spaces=False,
     errors="replace",
-    model_max_length=131072,
+    model_max_length=1024,
     split_special_tokens=False,
     unk_token=None,
     bos_token=None,
@@ -101,7 +134,7 @@ CHAT_TEMPLATE = """{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%-
 
 fast_tokenizer.chat_template = CHAT_TEMPLATE
 
-save_path = "weights/weights_tokenizer/hf_tokenizer_BBPE"
+save_path = "../weights/weights_tokenizer/hf_tokenizer_BBPE"
 fast_tokenizer.save_pretrained(save_path)
 
 print(f"Tokenizer saved to {save_path}")

src/step0_train_tokenizer_parquet.py

@@ -1,11 +1,11 @@
-import json
 import glob
-from tokenizers import Tokenizer
+from tokenizers import Tokenizer, normalizers, pre_tokenizers,  processors, decoders
+from tokenizers.normalizers import NFC
 from tokenizers.models import BPE
 from tokenizers.trainers import BpeTrainer
-from tokenizers.pre_tokenizers import ByteLevel
 from transformers import PreTrainedTokenizerFast
 import pyarrow.dataset as ds
+
 SPECIAL_TOKENS = [
     "<|endoftext|>",
     "<|im_start|>",
@@ -51,24 +51,57 @@ ADDITIONAL_SPECIAL_TOKENS = [
     "<|video_pad|>"
 ]
 
-
 def read_parquet_shards(file_pattern, batch_size=1000):
     dataset = ds.dataset(file_pattern, format='parquet')
     scanner = dataset.scanner(batch_size=batch_size)
     for batch in scanner.to_batches():
         yield batch.to_pandas()['text'].tolist()  
 
-tokenizer = Tokenizer(BPE(unk_token=None))
+tokenizer = Tokenizer(BPE(
-tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
+    unk_token=None,
+    continuing_subword_prefix="",
+    end_of_word_suffix="",
+    fuse_unk=False,
+    byte_fallback=False
+))
+tokenizer.normalizer = normalizers.Sequence([NFC()])
+regex_pattern = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+"""
+tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
+    pre_tokenizers.Split(
+        pattern=pre_tokenizers.Regex(regex_pattern),
+        behavior="isolated",
+        invert=False
+    ),
+    pre_tokenizers.ByteLevel(
+        add_prefix_space=False,
+        trim_offsets=False,
+        use_regex=False
+    )
+])
+
+tokenizer.post_processor = processors.ByteLevel(
+    add_prefix_space=False,
+    trim_offsets=False,
+    use_regex=False
+)
+
+tokenizer.decoder = decoders.ByteLevel(
+    add_prefix_space=False,
+    trim_offsets=False,
+    use_regex=False
+)
 
 trainer = BpeTrainer(
     vocab_size=32000,  
-    min_frequency=2,    
+    min_frequency=1,    
     special_tokens=SPECIAL_TOKENS,
-    show_progress=True
+    show_progress=True,
+    continuing_subword_prefix="",
+    end_of_word_suffix=""
 )
 
 corpus_files = glob.glob('corpus/parquet/train-0000*-of-00399.parquet' )
+
 tokenizer.train_from_iterator(
     iterator=read_parquet_shards(corpus_files),
     trainer=trainer,
@@ -84,17 +117,14 @@ fast_tokenizer = PreTrainedTokenizerFast(
     add_prefix_space=False,
     clean_up_tokenization_spaces=False,
     errors="replace",
-    model_max_length=131072,
+    model_max_length=1024,
     split_special_tokens=False,
     unk_token=None,
     bos_token=None,
 )
 
 CHAT_TEMPLATE = """{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}"""
-
 fast_tokenizer.chat_template = CHAT_TEMPLATE
-
 save_path = "weights/weights_tokenizer/hf_tokenizer_BBPE"
 fast_tokenizer.save_pretrained(save_path)
-
 print(f"Tokenizer saved to {save_path}")

src/step1_init_model_weights.py

@@ -0,0 +1,21 @@
+from configuration_dolly import DollyConfig
+from modeling_dolly import DollyModel
+
+model = DollyModel(DollyConfig())
+
+#保存配置和权重
+model.config.save_pretrained("../weights/weights_llm")
+model.save_pretrained("../weights/weights_llm")
+
+#加载配置和权重
+config = DollyConfig.from_pretrained("../weights/weights_llm")
+model = DollyModel.from_pretrained("../weights/weights_llm", config=config)
+
+#计算参数量
+model_size = sum(t.numel() for t in model.parameters())
+print(f"Model Size: {model_size/1000**2:.1f}M parameters")
+
+#测试模型
+# input_ids = torch.tensor([[101, 1024, 102]]).long() 
+# outputs = model(input_ids)
+# print(outputs)

src/step2_train_llm.py

@@ -0,0 +1,124 @@
+from modeling_dolly import DollyForCausalLM
+from configuration_dolly import DollyConfig
+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 = "../weights/weights_llm"  
+tokenizer_path=""
+resume_option = None
+first_time=True                     
+# resume_option = True                    
+# resume_option = "../checkpoints_ds/checkpoint-500" 
+if first_time:
+    model = DollyForCausalLM(DollyConfig()) 
+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("../corpus/*.jsonl")
+logger.info(f"找到训练语料: {data_files}")
+
+dataset = load_dataset('json', data_files=data_files, split='train')
+# dataset = load_dataset('text', data_files={'train': '../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="../checkpoints_ds",
+    deepspeed="../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="../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("../result/final")
+    tokenizer.save_pretrained("../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.py

@@ -19,10 +19,10 @@ logger = logging.getLogger(__name__)
 accelerator = Accelerator()
 set_seed(42)
 
-model_name = "/mnt/han.luzhi/dolly_llm/weights_llm"  
+model_name = "../weights_llm"  
 resume_option = None                     
 # resume_option = True                    
-# resume_option = "/mnt/han.luzhi/dolly_llm/checkpoints_ds/checkpoint-500"  
+# resume_option = "../checkpoints_ds/checkpoint-500"  
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
@@ -36,11 +36,11 @@ else:
     model = DollyForCausalLM.from_pretrained(model_name)
 model.gradient_checkpointing_enable()
 
-data_files = glob.glob("/mnt/han.luzhi/dolly_llm/corpus/*.jsonl")
+data_files = glob.glob("../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'})
+# dataset = load_dataset('text', data_files={'train': '../corpus/wikipedia.txt'})
 
 def preprocess_function(examples):
     return tokenizer(
@@ -65,15 +65,15 @@ data_collator = DataCollatorForLanguageModeling(
 )
 
 training_args = TrainingArguments(
-    output_dir="/mnt/han.luzhi/dolly_llm/checkpoints_ds",
+    output_dir="../checkpoints_ds",
-    deepspeed="/mnt/han.luzhi/dolly_llm/deepspeed_config/ds_stage2.json",  
+    deepspeed="../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_dir="../logs",
     logging_steps=20,
     save_total_limit=2,
     fp16=True,  
@@ -112,8 +112,8 @@ else:
 
 
 if accelerator.is_main_process: 
-    trainer.save_model("/mnt/han.luzhi/dolly_llm/result/final")
+    trainer.save_model("../result/final")
-    tokenizer.save_pretrained("/mnt/han.luzhi/dolly_llm/result/final")
+    tokenizer.save_pretrained("../result/final")
     logger.info("***** 训练完成，模型已保存 *****")
 
 ##  deepspeed --num_gpus=4 --master_port=12345 pretrain_transformers.py

src_dev/src_pretrain/pretrain_deepspeed_MoE.py

@@ -19,7 +19,7 @@ logger = logging.getLogger(__name__)
 accelerator = Accelerator()
 set_seed(42)
 
-model_name = "/mnt/han.luzhi/dolly/weights_MoE"  
+model_name = "../dolly/weights_MoE"  
 resume_option = None                     
 # resume_option = True                    
 # resume_option = "/mnt/han.luzhi/dolly_llm/checkpoints_ds/checkpoint-500"  
@@ -36,7 +36,7 @@ else:
     model = DollyMoEForCausalLM.from_pretrained(model_name)
 model.gradient_checkpointing_enable()
 
-data_files = glob.glob("/mnt/han.luzhi/dolly/corpus/*.jsonl")
+data_files = glob.glob("../dolly/corpus/*.jsonl")
 logger.info(f"找到训练语料: {data_files}")
 
 dataset = load_dataset('json', data_files=data_files, split='train')
@@ -65,15 +65,15 @@ data_collator = DataCollatorForLanguageModeling(
 )
 
 training_args = TrainingArguments(
-    output_dir="/mnt/han.luzhi/dolly/checkpoints_MoE",
+    output_dir="../checkpoints_MoE",
-    deepspeed="/mnt/han.luzhi/dolly/deepspeed_config/ds_stage2.json",  
+    deepspeed="../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_dir="../dolly/logs",
     logging_steps=20,
     save_total_limit=2,
     fp16=True,  
@@ -112,8 +112,8 @@ else:
 
 
 if accelerator.is_main_process: 
-    trainer.save_model("/mnt/han.luzhi/dolly/chechpoints_MoE/final")
+    trainer.save_model("../chechpoints_MoE/final")
-    tokenizer.save_pretrained("/mnt/han.luzhi/dolly/chechpoints_MoE/final")
+    tokenizer.save_pretrained("../chechpoints_MoE/final")
     logger.info("***** 训练完成，模型已保存 *****")
 
 ## deepspeed --num_gpus=4 --master_port=12345 pretrain_transformers.py

src_dev/src_pretrain/pretrain_transformers.py

@@ -17,11 +17,11 @@ logger = logging.getLogger(__name__)
 accelerator = Accelerator()
 set_seed(42)
 
-model_name = "/mnt/han.luzhi/dolly_llm/weights_llm"  
+model_name = "../weights_llm"  
 # resume_option = None                     
 # resume_option = True                    
-resume_option = "/mnt/han.luzhi/dolly_llm/result/checkpoint-59000"  
+resume_option = "../result/checkpoint-59000"  
-# resume_option = "/mnt/han.luzhi/dolly_llm/result/final" 
+# resume_option = "../result/final" 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
@@ -36,7 +36,7 @@ else:
 
 model.gradient_checkpointing_enable()
 
-dataset = load_dataset('text', data_files={'train': '/mnt/han.luzhi/dolly_llm/corpus/wikipedia.txt'})
+dataset = load_dataset('text', data_files={'train': '../corpus/wikipedia.txt'})
 
 def preprocess_function(examples):
     return tokenizer(
@@ -60,14 +60,14 @@ data_collator = DataCollatorForLanguageModeling(
 )
 
 training_args = TrainingArguments(
-    output_dir="/mnt/han.luzhi/dolly_llm/result",
+    output_dir="../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_dir="../logs",
     logging_steps=10,
     save_total_limit=2,
     fp16=True, 
@@ -102,8 +102,8 @@ logger.info(f" 总批大小 = {training_args.per_device_train_batch_size * accel
 trainer.train(resume_from_checkpoint=resume_option)
 
 if accelerator.is_main_process: 
-    trainer.save_model("/mnt/han.luzhi/dolly_llm/result/final")
+    trainer.save_model("../result/final")
-    tokenizer.save_pretrained("/mnt/han.luzhi/dolly_llm/result/final")
+    tokenizer.save_pretrained("../result/final")
     logger.info("***** 训练完成，模型已保存 *****")
 
 ##  python3 -m torch.distributed.launch --nproc_per_node=4 --master_port=12345 pretrain_transformers.py

src_dev/src_tokenizer/hf_tokenizer_code/step2_hf_train_from_scratch_BPEandBBPE.py

@@ -13,7 +13,7 @@ from tokenizers import (
 from transformers import PreTrainedTokenizerFast
 import logging
 from tqdm import tqdm
-log_dir = "/mnt/hanluzhi/dolly/logs"
+log_dir = "../logs"
 os.makedirs(log_dir, exist_ok=True)
 
 logging.basicConfig(

test.py

@@ -0,0 +1,54 @@
+import os
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+from datasets import Dataset, load_from_disk
+import deepspeed
+import numpy as np
+from glob import glob
+import logging
+from torch.utils.data import IterableDataset, DataLoader
+import pyarrow.parquet as pq
+# 配置参数
+MODEL_PATH = "models/qwen2"        # 本地模型路径
+TOKENIZED_DATA_DIR = "corpus"  # tokenized数据目录
+OUTPUT_DIR = "./output"            # 训练输出目录
+DEEPSPEED_CONFIG = "ds_config.json"  # DeepSpeed配置文件
+NUM_GPUS = 4                       # GPU数量
+PER_DEVICE_BATCH_SIZE = 4          # 每张GPU的batch size
+GRADIENT_ACCUM_STEPS = 8           # 梯度累积步数
+MAX_TRAIN_STEPS = 100000           # 最大训练步数
+SAVE_STEPS = 2000                  # 保存间隔步数
+LOGGING_STEPS = 50                 # 日志记录间隔
+LEARNING_RATE = 5e-5               # 学习率
+# 设置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+# ===== 1. 准备模型和分词器 =====
+logger.info("加载模型和分词器...")
+from transformers import AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained("/mnt/d/VMshare/LLM-from-scratch/weights/weights_tokenizer")
+file_path="tokenized_data/tokenized_train-00001-of-00399.parquet"
+table = pq.read_table(file_path)
+input_ids = table["input_ids"].to_numpy()
+attention_mask = table["attention_mask"].to_numpy()
+print(len(input_ids))
+for input_id in input_ids:
+    print(input_id)
+    print(len(input_id))
+    res=tokenizer.decode(input_id)
+    # print(res)
+    break
+
+file_path2="/mnt/d/VMshare/LLM-from-scratch/corpus/parquet/train-00001-of-00399.parquet"
+table2 = pq.read_table(file_path2)
+text = table2["text"]
+print(len(text))
+for value in text:
+    print(value)
+    break

test2.py

@@ -0,0 +1,169 @@
+import os
+import multiprocessing as mp
+from glob import glob
+import numpy as np
+from transformers import AutoTokenizer
+from datasets import load_dataset, Dataset
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pandas as pd
+from tqdm import tqdm
+import gc
+import psutil
+import time
+
+# 配置参数
+MODEL_PATH = "models/qwen2"  # 本地模型路径
+DATA_DIR = "datas"           # 原始数据目录
+OUTPUT_DIR = "tokenized_data"  # tokenized数据输出目录
+MAX_SEQ_LENGTH = 1024        # 最大序列长度
+NUM_PROC = min(32, psutil.cpu_count(logical=False))  # 物理核心数
+CHUNK_SIZE = 50000           # 每个处理块的大小
+MAX_MEMORY_PERCENT = 70      # 最大内存使用百分比
+
+# 创建输出目录
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+# 加载分词器
+tokenizer = AutoTokenizer.from_pretrained("/mnt/d/VMshare/LLM-from-scratch/weights/weights_tokenizer")
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# 获取所有parquet文件
+file_paths = ["/mnt/d/VMshare/LLM-from-scratch/corpus/parquet/train-00001-of-00399.parquet"]
+print(f"找到 {len(file_paths)} 个parquet文件")
+
+def memory_safe():
+    """检查内存是否安全"""
+    mem = psutil.virtual_memory()
+    return mem.percent < MAX_MEMORY_PERCENT
+
+def process_chunk(texts, tokenizer):
+    """处理文本块并返回tokenized结果"""
+    tokenized = tokenizer(
+        texts,
+        truncation=True,
+        max_length=MAX_SEQ_LENGTH,
+        padding="max_length",
+        return_tensors="np"
+    )
+    return tokenized["input_ids"], tokenized["attention_mask"]
+
+def process_file(file_path):
+    """处理单个大型parquet文件"""
+    try:
+        base_name = os.path.basename(file_path)
+        output_path = os.path.join(OUTPUT_DIR, f"tokenized_{base_name}")
+        
+        # 如果文件已处理则跳过
+        if os.path.exists(output_path):
+            print(f"文件已存在: {output_path}, 跳过")
+            return
+        
+        print(f"开始处理: {base_name}")
+        start_time = time.time()
+        
+        # 1. 获取文件总行数
+        dataset = load_dataset("parquet", data_files=file_path, split="train")
+        total_rows = len(dataset)
+        print(f"文件包含 {total_rows:,} 行")
+        
+        # 2. 创建空DataFrame用于收集结果
+        all_input_ids = []
+        all_attention_mask = []
+        
+        # 3. 分块处理
+        num_chunks = (total_rows + CHUNK_SIZE - 1) // CHUNK_SIZE
+        
+        for chunk_idx in range(num_chunks):
+            if not memory_safe():
+                print("内存使用过高，等待释放...")
+                time.sleep(30)
+                gc.collect()
+                if not memory_safe():
+                    print("内存仍然不足，跳过当前文件")
+                    return
+            
+            start_row = chunk_idx * CHUNK_SIZE
+            end_row = min((chunk_idx + 1) * CHUNK_SIZE, total_rows)
+            
+            # 读取当前块
+            chunk = dataset.select(range(start_row, end_row))
+            texts = chunk["text"]  # 根据实际字段名调整
+            
+            # 处理当前块
+            input_ids, attention_mask = process_chunk(texts, tokenizer)
+            
+            # 收集结果
+            all_input_ids.append(input_ids)
+            all_attention_mask.append(attention_mask)
+            
+            # 每处理5个块显示进度
+            if (chunk_idx + 1) % 5 == 0 or (chunk_idx + 1) == num_chunks:
+                elapsed = time.time() - start_time
+                rows_done = min((chunk_idx + 1) * CHUNK_SIZE, total_rows)
+                percent = rows_done / total_rows * 100
+                print(f"进度: {rows_done:,}/{total_rows:,} 行 ({percent:.1f}%) "
+                      f"耗时: {elapsed:.1f}秒")
+        
+        # 4. 合并所有块
+        all_input_ids = np.vstack(all_input_ids)
+        all_attention_mask = np.vstack(all_attention_mask)
+        
+        # 5. 创建PyArrow表并保存
+        table = pa.Table.from_arrays([
+            pa.array(all_input_ids.reshape(-1, MAX_SEQ_LENGTH).tolist()),
+            pa.array(all_attention_mask.reshape(-1, MAX_SEQ_LENGTH).tolist())
+        ], names=["input_ids", "attention_mask"])
+        
+        pq.write_table(table, output_path)
+        
+        # 6. 清理内存
+        del dataset, all_input_ids, all_attention_mask, table
+        gc.collect()
+        
+        elapsed = time.time() - start_time
+        print(f"完成处理: {base_name} -> {output_path} "
+              f"总耗时: {elapsed:.1f}秒, 平均速度: {total_rows/elapsed:.1f}行/秒")
+    
+    except Exception as e:
+        print(f"处理文件 {file_path} 时出错: {str(e)}")
+        # 清理可能的残留资源
+        if 'dataset' in locals():
+            del dataset
+        gc.collect()
+
+def process_file_wrapper(args):
+    """包装函数用于处理内存监控"""
+    file_path, idx, total = args
+    print(f"\n{'='*50}\n处理文件 [{idx+1}/{total}]: {os.path.basename(file_path)}\n{'='*50}")
+    process_file(file_path)
+    gc.collect()
+
+if __name__ == "__main__":
+    print(f"开始预处理 {len(file_paths)} 个文件，使用 {NUM_PROC} 个进程")
+    print(f"每个文件约18万行，总数据量约 {len(file_paths)*180000/1000000:.1f} 百万行")
+    
+    # 准备参数
+    file_args = [(fp, idx, len(file_paths)) for idx, fp in enumerate(file_paths)]
+    
+    # 使用进程池并行处理
+    with mp.Pool(processes=NUM_PROC) as pool:
+        results = list(tqdm(pool.imap(process_file_wrapper, file_args), total=len(file_paths)))
+    
+    print("所有文件处理完成!")
+    
+    # 生成元数据文件
+    metadata = {
+        "model": MODEL_PATH,
+        "tokenized_date": time.strftime("%Y-%m-%d %H:%M:%S"),
+        "max_seq_length": MAX_SEQ_LENGTH,
+        "num_files": len(file_paths),
+        "output_dir": OUTPUT_DIR
+    }
+    
+    with open(os.path.join(OUTPUT_DIR, "metadata.json"), "w") as f:
+        import json
+        json.dump(metadata, f, indent=2)
+    
+    print(f"元数据已保存到 {OUTPUT_DIR}/metadata.json")

utils/parquet2jsonl.py

@@ -0,0 +1,25 @@
+import pyarrow.parquet as pq
+import pyarrow as pa
+import json
+import os
+import glob
+
+input_pattern = "corpus/parquet/train-*-of-00399.parquet"
+output_dir = "corpus/jsonl"
+
+os.makedirs(output_dir, exist_ok=True)
+
+parquet_files = glob.glob(input_pattern)
+
+for parquet_file in parquet_files:
+    table = pq.read_table(parquet_file)
+    base_name = os.path.basename(parquet_file)
+    jsonl_file = os.path.join(output_dir, base_name.replace('.parquet', '.jsonl'))
+    
+    with open(jsonl_file, 'w', encoding='utf-8') as f:
+        for batch in table.to_batches():
+            for i in range(batch.num_rows):
+                row_dict = {col: batch[col][i].as_py() for col in batch.schema.names}
+                f.write(json.dumps(row_dict, ensure_ascii=False) + '\n')
+    
+    print(f"已转换 {parquet_file} 为 {jsonl_file}")

weights/weights_tokenizer/hf_tokenizer_BBPE/special_tokens_map.json

@@ -1,19 +0,0 @@
-{
-  "additional_special_tokens": [
-    "<|im_start|>",
-    "<|im_end|>",
-    "<|object_ref_start|>",
-    "<|object_ref_end|>",
-    "<|box_start|>",
-    "<|box_end|>",
-    "<|quad_start|>",
-    "<|quad_end|>",
-    "<|vision_start|>",
-    "<|vision_end|>",
-    "<|vision_pad|>",
-    "<|image_pad|>",
-    "<|video_pad|>"
-  ],
-  "eos_token": "<|im_end|>",
-  "pad_token": "<|endoftext|>"
-}

weights/weights_tokenizer/hf_tokenizer_BBPE/tokenizer_config.json

@@ -1,240 +0,0 @@
-{
-  "add_bos_token": false,
-  "add_prefix_space": false,
-  "added_tokens_decoder": {
-    "0": {
-      "content": "<|endoftext|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "1": {
-      "content": "<|im_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "2": {
-      "content": "<|im_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "3": {
-      "content": "<|object_ref_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "4": {
-      "content": "<|object_ref_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "5": {
-      "content": "<|box_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "6": {
-      "content": "<|box_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "7": {
-      "content": "<|quad_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8": {
-      "content": "<|quad_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "9": {
-      "content": "<|vision_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "10": {
-      "content": "<|vision_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "11": {
-      "content": "<|vision_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "12": {
-      "content": "<|image_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "13": {
-      "content": "<|video_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "14": {
-      "content": "<tool_call>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "15": {
-      "content": "</tool_call>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "16": {
-      "content": "<|fim_prefix|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "17": {
-      "content": "<|fim_middle|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "18": {
-      "content": "<|fim_suffix|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "19": {
-      "content": "<|fim_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "20": {
-      "content": "<|repo_name|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "21": {
-      "content": "<|file_sep|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "22": {
-      "content": "<tool_response>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "23": {
-      "content": "</tool_response>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "24": {
-      "content": "<think>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "25": {
-      "content": "</think>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
-  "additional_special_tokens": [
-    "<|im_start|>",
-    "<|im_end|>",
-    "<|object_ref_start|>",
-    "<|object_ref_end|>",
-    "<|box_start|>",
-    "<|box_end|>",
-    "<|quad_start|>",
-    "<|quad_end|>",
-    "<|vision_start|>",
-    "<|vision_end|>",
-    "<|vision_pad|>",
-    "<|image_pad|>",
-    "<|video_pad|>"
-  ],
-  "bos_token": null,
-  "chat_template":  "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set content = message.content %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in message.content %}\n                {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
-  "clean_up_tokenization_spaces": false,
-  "eos_token": "<|im_end|>",
-  "errors": "replace",
-  "extra_special_tokens": {},
-  "model_max_length": 131072,
-  "pad_token": "<|endoftext|>",
-  "split_special_tokens": false,
-  "tokenizer_class": "PreTrainedTokenizer",
-  "unk_token": null
-}