RoFormerForCausalLM, BigBirdForCausalLM, RemBertForCausalLM, MegatronBertForCausalLM attend to future tokens since v5.9.0 — CLM training loss collapses #48745

System Info

- `transformers` version: 5.18.0.dev0
- Platform: Linux-7.0.0-31-generic-x86_64-with-glibc2.39
- Python version: 3.10.20
- Huggingface_hub version: 1.31.0
- Safetensors version: 0.8.0
- Accelerate version: 1.13.0
- Accelerate config: not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.12.1+cu130 (CUDA)
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: no, the script below is CPU-only
- GPU type: NVIDIA GeForce RTX 3060

Who can help?

@vasqu @Cyrilvallez

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Training RoFormerForCausalLM (is_decoder=True) as a plain causal LM, the loss fell to almost zero within a few hundred steps on random tokens. The decoder is not causal: each position sees the tokens after it, so its label is visible. BertLMHeadModel is fine; these four are not.

main at 5474a55e92, CPU:

import torch
from transformers import (
    BertConfig, BertLMHeadModel,
    BigBirdConfig, BigBirdForCausalLM,
    MegatronBertConfig, MegatronBertForCausalLM,
    RemBertConfig, RemBertForCausalLM,
    RoFormerConfig, RoFormerForCausalLM,
)

torch.manual_seed(0)
kw = dict(vocab_size=99, hidden_size=32, num_hidden_layers=2, num_attention_heads=4, intermediate_size=37, is_decoder=True)
models = {
    "bert": BertLMHeadModel(BertConfig(**kw)),
    "big_bird": BigBirdForCausalLM(BigBirdConfig(attention_type="original_full", **kw)),
    "megatron_bert": MegatronBertForCausalLM(MegatronBertConfig(**kw)),
    "rembert": RemBertForCausalLM(RemBertConfig(input_embedding_size=32, output_embedding_size=32, **kw)),
    "roformer": RoFormerForCausalLM(RoFormerConfig(embedding_size=32, **kw)),
}

# 1. Do the logits at positions 0..6 change when only token 7 changes?
ids = torch.randint(3, 99, (1, 8))
other = ids.clone()
other[0, -1] = (ids[0, -1] + 1) % 96 + 3
for name, model in models.items():
    model.eval()
    with torch.no_grad():
        delta = (model(input_ids=ids).logits[0, :-1] - model(input_ids=other).logits[0, :-1]).abs().max().item()
    print(f"{name:14} earlier logits move by {delta:.4f} when the last token changes")

# 2. Train as a causal LM on fresh random tokens: nothing is predictable, so the loss cannot go below ln(96) = 4.56
for name in ("bert", "roformer"):
    model = models[name].train()
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    for step in range(600):
        ids = torch.randint(3, 99, (8, 16))
        loss = model(input_ids=ids, labels=ids).loss
        opt.zero_grad(); loss.backward(); opt.step()
    print(f"{name:14} loss after 600 steps on random tokens: {loss.item():.2f}")

bert           earlier logits move by 0.0000 when the last token changes
big_bird       earlier logits move by 0.0007 when the last token changes
megatron_bert  earlier logits move by 0.0149 when the last token changes
rembert        earlier logits move by 0.0009 when the last token changes
roformer       earlier logits move by 0.0011 when the last token changes
bert           loss after 600 steps on random tokens: 4.57
roformer       loss after 600 steps on random tokens: 0.11

Pretrained weights (from_pretrained(..., is_decoder=True)): junnyu/roformer_chinese_base as RoFormerForCausalLM moves earlier logits by 6.5 when a later token changes, google/bigbird-roberta-base as BigBirdForCausalLM by 9.0. Greedy generate() with and without use_cache also disagree.

RoFormerModel.forward and the other three always call create_bidirectional_mask (modeling_roformer.py:727), where BertModel branches on is_decoder and calls create_causal_mask (modeling_bert.py:696-708). Before #43924 all went through get_extended_attention_mask, which had that branch built in; on the commit before it the script prints 0.0000 everywhere. A v5.9.0 regression, invisible to CI because these four are still excluded from the generation tests and test_decoder_model_past_with_large_inputs compares the embedding output, which no mask can change.

Expected behavior

With is_decoder=True the four should be causal, as up to v5.8 and as BERT still is: a change to token k leaves the logits before k untouched, and cached and cacheless generate() agree.

Nobody seems to have reported it — a loss that goes down too well is not an error — but any *ForCausalLM fine-tune of these families with is_decoder=True since May, and any EncoderDecoderModel with a RoFormer, RemBert or MegatronBert decoder, trains with the labels visible.

One caveat: the junnyu/roformer_chinese_sim_char_* checkpoints are SimBERT/UniLM-trained (bidirectional prompt), and the bug happens to give cached generate() exactly that, so they generate better on ≥ 5.9 than on 5.8 or with the fix — a prefix-LM mask, not is_decoder semantics.

I have the fix (BERT's is_decoder branch in each of the four) and a regression test ready locally — red on main per model, green with the fix. Happy to open the PR if a maintainer wants it.