KeyError: 'default' for discrete diffusion language model LLaDA2

[ksasi]

Describe the bug

Hi,

The following code block from the documentation (https://huggingface.co/docs/diffusers/main/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) is giving key error :

model_id = "inclusionAI/LLaDA2.1-mini"

model = AutoModelForCausalLM.from_pretrained(
 model_id, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
)

Reproduction

mport torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

from diffusers import BlockRefinementScheduler, LLaDA2Pipeline

model_id = "inclusionAI/LLaDA2.1-mini"

model = AutoModelForCausalLM.from_pretrained(
model_id, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
scheduler = BlockRefinementScheduler()

Logs

---------------------------------------------------------------------------
KeyError Traceback (most recent call last)
[/tmp/ipykernel_1800/3647447919.py](https://localhost:8080/#) in <cell line: 0>()
 7 model_id = "inclusionAI/LLaDA2.1-mini"
 8 
----> 9 model = AutoModelForCausalLM.from_pretrained(
 10 model_id, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
 11 )

4 frames
[/usr/local/lib/python3.12/dist-packages/transformers/models/auto/auto_factory.py](https://localhost:8080/#) in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
 363 model_class.register_for_auto_class(auto_class=cls)
 364 model_class = add_generation_mixin_to_remote_model(model_class)
--> 365 return model_class.from_pretrained(
 366 pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
 367 )

[/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py](https://localhost:8080/#) in from_pretrained(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, use_safetensors, weights_only, *model_args, **kwargs)
 4070 with ContextManagers(model_init_context):
 4071 # Let's make sure we don't run the init function of buffer modules
-> 4072 model = cls(config, *model_args, **model_kwargs)
 4073 
 4074 if hf_quantizer is not None: # replace module with quantized modules (does not touch weights)

[~/.cache/huggingface/modules/transformers_modules/inclusionAI/LLaDA2_dot_1_hyphen_mini/f21be037104f6e044e1a86b6d8864a6b85cc868e/modeling_llada2_moe.py](https://localhost:8080/#) in __init__(self, config)
 960 def __init__(self, config: LLaDA2MoeConfig):
 961 super().__init__(config)
--> 962 self.model = LLaDA2MoeModel(config)
 963 self.vocab_size = config.vocab_size
 964 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

[~/.cache/huggingface/modules/transformers_modules/inclusionAI/LLaDA2_dot_1_hyphen_mini/f21be037104f6e044e1a86b6d8864a6b85cc868e/modeling_llada2_moe.py](https://localhost:8080/#) in __init__(self, config)
 781 self._use_flex_attention = config._attn_implementation == "flex_attention"
 782 self.norm = LLaDA2MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
--> 783 self.rotary_emb = LLaDA2MoeRotaryEmbedding(config=config)
 784 self.gradient_checkpointing = False
 785 # Initialize weights and apply final processing

[~/.cache/huggingface/modules/transformers_modules/inclusionAI/LLaDA2_dot_1_hyphen_mini/f21be037104f6e044e1a86b6d8864a6b85cc868e/modeling_llada2_moe.py](https://localhost:8080/#) in __init__(self, config, device)
 106 
 107 self.config = config
--> 108 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
 109 
 110 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)

KeyError: 'default'

System Info

0.38.0.dev0 (from git+https://github.com/huggingface/diffusers@f2be8bd6b3dc4035bd989dc467f15d86bf3c9c12)

Who can help?

No response

[matdou]

Ok, so I dug a bit about this issue. This isn't a diffusers issue, the issue is in transformers/modeling_rope_utils.py.

The model's remote code (modeling_llada2_moe.py) imports ROPE_INIT_FUNCTIONS directly from transformers and does:

self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] # rope_type = 'default'

But in transformers 5.4.0, ROPE_INIT_FUNCTIONS is defined as:

ROPE_INIT_FUNCTIONS = {
 "linear": ...,
 "dynamic": ...,
 "yarn": ...,
 "longrope": ...,
 "llama3": ...,
 # "default" is missing
}

Yet the transformers docs (same file, line 575-576) list explicitly 'default' as a valid rope_type. So 'default' was either removed or never added in 5.x. Either way it's a bug in tranfsormers.

Minimal repro (no model download):

from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
rope_init_fn = ROPE_INIT_FUNCTIONS["default"] # KeyError: 'default'

The fix should be adding 'default' back to ROPE_INIT_FUNCTIONS in huggingface/transformers. Happy to be corrected if I'm missing something or to open a PR on transformers if this is confirmed.

[matdou]

Could you share your transformers version?