Pretraining with FlexBert but can't export to HF

[asafam]

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!

[asafam]

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
    )

[asafam]

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