unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to archive

Integrate the UltraBERT architecture

Language: Python+1666 -0feature
concepts demonstrated
Class inheritance and method overridingCustom model architecture integrationAttention mechanisms with variable window sizesRotary positional embeddings
The Issue & Context

Before this PR, the Transformers library lacked support for the UltraBERT architecture, which requires handling distinct rotary embedding setups per layer to match its local and global attention mechanisms.

Root Cause Analysis

UltraBERT's architecture differs from similar models by using two unique rotary embedding theta values—one for local attention windows and another for global attention—necessitating multiple rotary embeddings stored and managed simultaneously, which was not supported in existing code.

Solution Strategy & Architectural Pattern

This PR introduced UltraBERT as a new model in Transformers, leveraging existing ModernBERT and EuroBERT modular code as superclasses to minimize duplication. It implemented a rotary embedding ModuleDict in modeling_ultrabert.py to initialize and store separate rotary embeddings for local and global attention layers, enabling correct per-layer rotary embedding application. The PR also added config, pretrained model integration, and heads for MaskedLM and SequenceClassification.

Key Takeaway for Contributors

When integrating models with layer-specific parameter variations—like distinct rotary embeddings—designing modular, flexible parameter containers (e.g., ModuleDict keyed by embedding type) allows clean management without duplicating entire model code, facilitating maintainability and extensibility.