Skip to content

GH-3630: persist trust_remote_code across save/load round-trips#3697

Open
haoyu-haoyu wants to merge 1 commit intoflairNLP:masterfrom
haoyu-haoyu:fix/persist-trust-remote-code
Open

GH-3630: persist trust_remote_code across save/load round-trips#3697
haoyu-haoyu wants to merge 1 commit intoflairNLP:masterfrom
haoyu-haoyu:fix/persist-trust-remote-code

Conversation

@haoyu-haoyu
Copy link
Copy Markdown

Models that need trust_remote_code=True (EuroBERT, OLMo, etc.) train fine but crash when the trainer tries to reload the best checkpoint, because the custom config class is not in CONFIG_MAPPING at that point.

Root cause: trust_remote_code flows into __init__ via **kwargs and gets consumed by the HF calls, but is never stored. On save, to_params() does not include it. On reload, from_params() hits CONFIG_MAPPING[model_type] which raises KeyError for custom types.

The fix:

  • Store trust_remote_code as an instance attribute
  • Emit it from to_params() when True
  • In from_params(), when the model type is unknown and trust_remote_code was saved, use AutoConfig.from_pretrained() to re-download the config (which registers the custom class). _name_or_path is stripped from the overrides to avoid clobbering the freshly-resolved repo id.
  • Pass trust_remote_code through to create_from_state() so the tokenizer and model loading calls also get it

Closes #3630

Models that require trust_remote_code=True (EuroBERT, OLMo, etc.) train
fine but crash on reload because the custom config class is not in
CONFIG_MAPPING.  This happens because the flag was consumed by __init__
but never saved.

- Store trust_remote_code as an instance attribute
- Include it in to_params() when True
- In from_params(), when the model type is not in CONFIG_MAPPING and
  trust_remote_code was saved, re-download the config from the hub so
  the custom class gets registered.  Strip _name_or_path from the
  overrides to avoid clobbering the freshly-resolved repo id.
- Pass trust_remote_code through to create_from_state() so tokenizer
  and model loading also get it
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.

[Bug]: Cannot Load Saved HF Model EuroBert, but training is possible

1 participant