Skip to content

model: add RBLNModernBertForMaskedLM - #625

Merged
rebel-thkim merged 3 commits into
RBLN-SW:devfrom
warpspaceinc:feat/modernbert-maskedlm
Aug 6, 2026
Merged

model: add RBLNModernBertForMaskedLM#625
rebel-thkim merged 3 commits into
RBLN-SW:devfrom
warpspaceinc:feat/modernbert-maskedlm

Conversation

@ho4040

@ho4040 ho4040 commented Jul 7, 2026

Copy link
Copy Markdown

Type of Change

  • New Model Support

Changes Overview

Adds RBLNModernBertForMaskedLM (+ RBLNModernBertForMaskedLMConfig) for the
ModernBERT encoder, following the existing RBLNRobertaForMaskedLM pattern.

ModernBERT needs two model-specific adjustments over the generic encoder path:

  • Force SDPA at load (get_pytorch_model). ModernBERT selects its attention
    backend from config._attn_implementation and prefers HuggingFace's
    flash_attention_2 (the CUDA-only flash-attn kernels) when that package is
    present; rebel-compiler cannot compile that path for RBLN. We pin
    attn_implementation="sdpa" so compilation is deterministic regardless of the
    host environment and the compiler lowers scaled_dot_product_attention to its
    own kernel. (This is unrelated to RBLN's own flash_attn attn_impl, which
    is a decoder KV-cache feature and does not apply to encoders.)
  • A dedicated compile wrapper (ModernBertModelWrapper) forwards the 2D
    padding 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;
RBLNAutoModelForMaskedLM resolves it through the existing naming convention.
A test mirroring TestBertForMaskedLM is 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:

  • Env: ATOM+ (RBLN-CA22), rebel-compiler 0.11.0, optimum-rbln 0.11.0.post1,
    transformers 5.8.1, torch 2.11.0
  • Compile: jhu-clsp/mmBERT-base, export=True, max_seq_len=128,
    batch_size=1, fp32 → compiled_model.rbln produced
  • Accuracy vs CPU eager (SDPA): top-1 agreement = 1.00 over all real
    token positions; masked-token prediction matches ("Paris"); masked-slot
    logit 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

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
rebel-kblee self-requested a review July 15, 2026 08:52

@rebel-kblee rebel-kblee left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM 👍 Thanks for contribute!

@rebel-thkim
rebel-thkim merged commit ee9ce75 into RBLN-SW:dev Aug 6, 2026
9 of 14 checks passed
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants