library_name: transformers datasets:

• bigcode/the-stack-v2-train-full-ids license: bigcode-openrail-m

ModularStarEncoder-1B 预训练模型

ModularStarEncoder-1B 是基于 The Stack v2 数据集预训练的编码器。该模型采用模块化设计，包含五个出口点，支持根据下游任务进行多出口微调。我们在 StarCoder-2 基础上进行改造，将参数量从150亿缩减至10亿（bfloat16精度）。

模型架构包含36个隐藏层，每层配置16个注意力头和4个键值头，采用分组查询注意力机制（GQA）。使用旋转位置编码（RoPE），基础周期θ=10^-6，隐藏维度为1024，中间层维度为12,288。

为提升效率，我们将因果自注意力层替换为双向自注意力层。与采用滑动窗口注意力的 StarCoder-2 不同，我们选择全局注意力机制以确保更好的模块化特性，避免滑动窗口的感知范围限制。同时将最大输入长度扩展至2048个token，支持处理比 StarEncoder 等前代代码编码器更长的代码片段。

模型集成FlashAttention V2加速推理过程。

• 论文: 《一模型统御全局：分层自蒸馏增强早期层表征》
• 支持语言: 600+种编程语言

代码转代码和文本转代码的微调版本参见 ModularStarEncoder-finetuned。

使用指南

from transformers import AutoModel
from transformers import AutoTokenizer

# 加载模型
model = AutoModel.from_pretrained("andreagurioli1995/ModularStarEncoder", trust_remote_code=True)

# 加载分词器（采用左填充）
tokenizer = AutoTokenizer.from_pretrained("andreagurioli1995/ModularStarEncoder")

code_snippet = "待嵌入的代码"

# 代码片段嵌入格式模板
sentence = f"{tokenizer.sep_token}{code_snippet}{tokenizer.cls_token}"

# 分词处理
tokenized_sensence = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=2048)

# 生成嵌入表示
embedded_sentence = model(**tokenized_sensence)

输出包含六个要素：

• last_hidden_state：模型最终隐藏层表征
• hidden_states：未经池化/归一化/投影的所有隐藏层原始表征
• loss：存在真实标签时的损失值（推理时为None）
• prediction_logits：掩码语言建模头的预测分数
• seq_relationship_scores：上下文内损失的预测分数（如需有效评分需用分隔符连接多个样本）
• attentions：编码器注意力分数

训练详情

我们在512张NVIDIA Ampere（64GB）GPU上使用 Leonardo 超算进行预训练，累计消耗45万GPU工时。训练批次规模为399万token，共执行24.5万训练步，处理总token量达1万亿。

许可协议

本模型采用BigCode OpenRAIL-M v1许可证，完整协议参见此处。

引用格式

@article{gurioli2025modeltrainallhierarchical,
      title={一模型统御全局：分层自蒸馏增强早期层表征}, 
      author={Andrea Gurioli and Federico Pennino and João Monteiro and Maurizio Gabbrielli},
      year={2025},
      eprint={2503.03008},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2503.03008}, 
}