---
library_name: transformers
datasets:
- bigcode/the-stack-v2-train-full-ids
license: bigcode-openrail-m
---

# ModularStarEncoder-1B Pre-trained model

<!-- Provide a quick summary of what the model is/does. -->

ModularStarEncoder-1B is an encoder pre-trained on [The Stack v2](https://huggingface.co/datasets/bigcode/the-stack-v2-train). 
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](https://huggingface.co/bigcode/starcoder2-15b), 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](https://huggingface.co/bigcode/starcoder2-15b), 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](https://huggingface.co/bigcode/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](https://arxiv.org/abs/2503.03008)
- **Languages:** 600+ Programming languages


For the finetuned version on code-to-code and text-to code see [ModularStarEncoder-finetuned](https://huggingface.co/modularStarEncoder/ModularStarEncoder-finetuned).


### How to use
```python
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

<!-- Provide a longer summary of what this model is. -->
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](https://arxiv.org/abs/2307.16885) 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](https://huggingface.co/spaces/bigcode/bigcode-model-license-agreement).


# 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}, 
}
```