Pretraining with FlexBert but can't export to HF

@warner-benjamin

Hi @warner-benjamin and team,
Thank you for your fantastic work on ModernBERT and FlexBert!

I’ve trained a FlexBert model using your codebase and I’m now trying to convert it to Hugging Face Transformers format for easier sharing and downstream use. I found the convert_to_hf.py script referenced in some previous issues and used it, along with the model definitions from this YAML.

The model loads correctly with AutoModel and AutoTokenizer, but when I run a forward pass, the outputs are all NaN.

from transformers import AutoModel, AutoConfig

model_path = "</path/to/model/from/convert_to_hf>"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path)
inputs = tokenizer("Hello world", return_tensors="pt")
with torch.no_grad():
    outputs = test_model(**inputs) # <- returns all NaN

Do you have any advice or best practices for converting models to HF format, or tips on what might be causing this issue? Any guidance would be greatly appreciated!

Thanks in advance!

If anyone gets this error, I think I managed to solve it:
The issue was that I needed to add a config.json to the below snippet from the conversion function to run successfully:

config_json_path = f"{target_path}/config.json"
with open(config_json_path, "r") as f:
    config_dict = json.load(f)
    config_dict = update_config(
        config_dict, bos_token_id, eos_token_id, cls_token_id, pad_token_id, sep_token_id, max_length, torch_dtype
    )

I fixed the script to load some of the configuration from the yaml file I used to train with:

with open(yaml_config) as f:
      config = OmegaConf.load(f)
      model_config = OmegaConf.to_container(config.model.model_config, resolve=True)
      config_dict = cast(DictConfig, model_config)
      config_dict = update_config(
          config_dict, bos_token_id, eos_token_id, cls_token_id, pad_token_id, sep_token_id, max_length, vocab_size, torch_dtype
      )

Some adjustments are required in the update_config method as well to accommodate this change:

def update_config(
    source_config: dict,
    bos_token_id: int,
    eos_token_id: int,
    cls_token_id: int,
    pad_token_id: int,
    sep_token_id: int,
    max_length: int,
    vocab_size: int,
    torch_dtype: TorchDtype,
) -> dict:
    target_config = {
        "_name_or_path": "ModernBERT-base",
        "architectures": ["ModernBertForMaskedLM"],
        "attention_bias": source_config["attn_out_bias"],
        "attention_dropout": source_config["attention_probs_dropout_prob"],
        "bos_token_id": bos_token_id,
        "classifier_activation": source_config.get("head_class_act", source_config["hidden_act"]),
        "classifier_bias": source_config.get("head_class_bias", False),
        "classifier_dropout": source_config.get("head_class_dropout", 0.0),
        "classifier_pooling": source_config.get("classifier_pooling", "mean"),
        "cls_token_id": cls_token_id,
        "decoder_bias": source_config.get("decoder_bias", True),
        "deterministic_flash_attn": source_config.get("deterministic_fa2", False),
        "embedding_dropout": source_config["embed_dropout_prob"],
        "eos_token_id": eos_token_id,
        "global_attn_every_n_layers": source_config["global_attn_every_n_layers"],
        "global_rope_theta": source_config["rotary_emb_base"],
        "gradient_checkpointing": source_config.get("gradient_checkpointing", False),
        "hidden_activation": source_config["hidden_act"],
        "hidden_size": source_config["hidden_size"],
        "initializer_cutoff_factor": source_config.get("init_cutoff_factor", 2.0),
        "initializer_range": source_config.get("initializer_range", 0.02),
        "intermediate_size": source_config["intermediate_size"],
        "layer_norm_eps": source_config["norm_kwargs"]["eps"],
        "local_attention": source_config["sliding_window"],
        "local_rope_theta": source_config.get("local_attn_rotary_emb_base", source_config["rotary_emb_base"]),
        "max_position_embeddings": max_length,  # Override with first config value
        "mlp_bias": source_config["mlp_in_bias"],
        "mlp_dropout": source_config["mlp_dropout_prob"],
        "model_type": "modernbert",
        "norm_bias": source_config["norm_kwargs"]["bias"],
        "norm_eps": source_config["norm_kwargs"]["eps"],
        "num_attention_heads": source_config["num_attention_heads"],
        "num_hidden_layers": source_config["num_hidden_layers"],
        "pad_token_id": pad_token_id,
        "position_embedding_type": source_config.get("position_embedding_type", "absolute"),
        "sep_token_id": sep_token_id,
        "tie_word_embeddings": source_config.get("tie_word_embeddings", True),
        "torch_dtype": torch_dtype.value,
        "transformers_version": "4.48.0",
        "vocab_size": vocab_size,
    }
    return target_config

Navigation Menu

Repository navigation

Pretraining with FlexBert but can't export to HF #234

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Pretraining with FlexBert but can't export to HF #234

Description

Activity

asafam commented on May 25, 2025

asafam commented on May 25, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions