ModularStarEncoder-1B Pre-trained model

ModularStarEncoder-1B is an encoder pre-trained on The Stack v2. ModularStarEncoder is a modular pre-trained encoder with five exit points, allowing users to perform multiple exit fine-tuning depending on downstream tasks. We built ModularStarEncoder on top of StarCoder-2, reducing its size from 15B to 1B parameters in bfloat16. Our architecture consists of 36 hidden layers, each with 16 attention heads and 4 key-value heads, using Grouped Query Attention (GQA). The model employs Rotary Positional Encoding (RoPE) with a base period theta = 10^-6 and features a hidden dimensionality of 1024 with an intermediate size of 12,288.

To enhance efficiency, we replaced the causal self-attention layers with bidirectional self-attention. Unlike StarCoder-2, which uses sliding window attention, we opted for full attention to ensure greater modularity, avoiding the receptive field constraints of sliding window mechanisms. Additionally, we extended the maximum input length to 2048 tokens, accommodating longer code snippets compared to prior code encoders such as StarEncoder.

Finally, our implementation integrates FlashAttention V2 for faster inference.

• Paper: One Model to Train them All: Hierarchical Self-Distillation for Enhanced Early Layer Embeddings
• Languages: 600+ Programming languages

For the finetuned version on code-to-code and text-to code see ModularStarEncoder-finetuned.

How to use

from transformers import AutoModel
from transformers import AutoTokenizer

#import the model
model = AutoModel.from_pretrained("andreagurioli1995/ModularStarEncoder", trust_remote_code=True)

#import the tokenizer, the tokenizer applies LEFT padding!
tokenizer = AutoTokenizer.from_pretrained("andreagurioli1995/ModularStarEncoder")

 

code_snippet = "your code to embed here"

#You should follow this pattern to embed a snippet of code 
sentence =  f"{tokenizer.sep_token}{code_snippet}{tokenizer.cls_token}"

#Tokenizing your sentence
tokenized_sensence = tokenizer(sentence, return_tensors="pt",truncation=True, max_length=2048)

#Embedding the tokenized sentence
embedded_sentence = model(**tokenized_sensence)

You will get as an output six elements:

• last_hidden_state: the representation of the last hidden state from the model;
• hidden_states: raw representation from all the hidden states of the model, without pooling, normalization, and projection
• loss: loss value if a ground truth is given (None if used in inference)
• prediction_logits: prediction scores from masked language modeling head
• seq_relationship_scores: prediction scores of in-context loss (concatenate multiple samples with the separator token if you want a meaningful score)
• attentions: attention scores from the encoder

Training

We pre-trained ModularStarEncoder with a batch size of 3.99M tokens for 245,000 training steps, processing 1T tokens. The pre-training and fine-tuning were conducted on 512 NVIDIA Ampere (64GB) GPUs using the Leonardo supercomputer, requiring 450,000 GPU working hours.

Hyperparameter	Value
Hidden size	1024
Max. position embeddings	2048
Num. of attention heads	12
Num. of key values heads	4
Num. of hidden layers	36
Attention	GQA
Num. of parameters	≈1B
Training tokens	≈1T
Loss function	MLM + In-Context loss
Multi-layer loss	yes

Licence

The model is licensed under the BigCode OpenRAIL-M v1 license agreement. You can find the full agreement here.

Citation

@article{gurioli2025modeltrainallhierarchical,
      title={One Model to Train them All: Hierarchical Self-Distillation for Enhanced Early Layer Embeddings}, 
      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}, 
}