model: add RBLNModernBertForMaskedLM - #625
Merged
Merged
Conversation
Add ModernBERT MaskedLM support following the RBLNRobertaForMaskedLM pattern. Force SDPA at load (FlashAttention-2 is CUDA-only) and use a dedicated wrapper that passes the 2D mask through so ModernBERT builds its own full + sliding-window masks. Register exports and add a test mirroring TestBertForMaskedLM. Validated on ATOM+ (RBLN-CA22).
rebel-kblee
self-requested a review
July 15, 2026 08:52
rebel-kblee
approved these changes
Jul 15, 2026
rebel-kblee
left a comment
Contributor
There was a problem hiding this comment.
LGTM 👍 Thanks for contribute!
rebel-dkhong
pushed a commit
that referenced
this pull request
Aug 10, 2026
Co-authored-by: rebel-kblee <119555851+rebel-kblee@users.noreply.github.com> Co-authored-by: rebel-thkim <157466331+rebel-thkim@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Type of Change
Changes Overview
Adds
RBLNModernBertForMaskedLM(+RBLNModernBertForMaskedLMConfig) for theModernBERT encoder, following the existing
RBLNRobertaForMaskedLMpattern.ModernBERT needs two model-specific adjustments over the generic encoder path:
get_pytorch_model). ModernBERT selects its attentionbackend from
config._attn_implementationand prefers HuggingFace'sflash_attention_2(the CUDA-onlyflash-attnkernels) when that package ispresent; rebel-compiler cannot compile that path for RBLN. We pin
attn_implementation="sdpa"so compilation is deterministic regardless of thehost environment and the compiler lowers
scaled_dot_product_attentionto itsown kernel. (This is unrelated to RBLN's own
flash_attnattn_impl, whichis a decoder KV-cache feature and does not apply to encoders.)
ModernBertModelWrapper) forwards the 2Dpadding mask unchanged instead of pre-expanding it with
_prepare_4d_attention_mask. ModernBERT alternates full and local(sliding-window) attention and builds both 4D masks internally via
create_bidirectional_mask/create_bidirectional_sliding_window_mask;pre-expanding would feed the global mask to the sliding layers and drop the
local windowing.
The class is registered in the three
__init__export lists;RBLNAutoModelForMaskedLMresolves it through the existing naming convention.A test mirroring
TestBertForMaskedLMis added(
hf-internal-testing/tiny-random-ModernBertForMaskedLM).Motivation and Context
ModernBERT (including the multilingual mmBERT) currently has no official
optimum-rbln support. We built and validated this adapter end-to-end on real
ATOM+ hardware:
transformers 5.8.1, torch 2.11.0
jhu-clsp/mmBERT-base,export=True,max_seq_len=128,batch_size=1, fp32 →compiled_model.rblnproducedtoken positions; masked-token prediction matches (
"Paris"); masked-slotlogit cosine similarity 0.957
We also published a technical write-up of the port, showcasing running an
unsupported encoder on ATOM+:
We'd love ModernBERT to become officially supported so we can point readers at
first-party support. Happy to adjust to your conventions, extend sequence-length
/ bucketing coverage, or hand off for an accuracy review.
Related Issues