Skip to content

from_pretrained loads wrong custom module after save_pretrained #45698

Closed
#45729
Closed
from_pretrained loads wrong custom module after save_pretrained#45698
#45729
Labels
bug

Description

@nurpax
nurpax
opened on Apr 29, 2026

System Info

Transformers version: 5.7.0.dev0
Python version: 3.13.3
Platform: Linux-6.17.0-20-generic-x86_64-with-glibc2.39
Machine: x86_64

Note, running transformers env fails with: NameError: name 'CompletionCreateParamsStreaming' is not defined

Who can help?

@Cyrilvallez (model loading)

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

With the attached main.py and custom_model.py:

1. Create two local saved models

python main.py save --path=pretrained_a --magic="Magic A"
python main.py save --path=pretrained_b --magic="Magic B"

Note: these setup commands print <MAGIC> because they instantiate the local source file before patching the saved custom_model.py.

2. Control case: loading both models works

python main.py load-and-load --path-a=pretrained_a --path-b=pretrained_b

This correctly prints Magic A from pretrained_a and Magic B from pretrained_b:

Load "pretrained_a"
Model says:  Magic A
Source path: HF_MODULES_CACHE/transformers_modules/pretrained_a/custom_model.py
Source says: Magic A
Load "pretrained_b"
Model says:  Magic B
Source path: HF_MODULES_CACHE/transformers_modules/pretrained_b/custom_model.py
Source says: Magic B

3. Failing case: saving first changes the later load

python main.py save-and-load --save-path=pretrained_a --save-magic="Magic A" --load-path=pretrained_b
Save "pretrained_a" with "Magic A"
Model says:  <MAGIC>
Source path: custom_model.py
Source says: <MAGIC>
Load "pretrained_b"
Model says:  <MAGIC>
Source path: custom_model.py
Source says: <MAGIC>

Actual result: pretrained_b is loaded from local custom_model.py and prints <MAGIC>.

custom_model.py
main.py

Expected behavior

Expected result: pretrained_b should be loaded from HF_MODULES_CACHE/transformers_modules/pretrained_b/custom_model.py and print Magic B, like it does in the control case.

Activity

  1. nurpax
    added
    bug
    on Apr 29, 2026

  2. Rocketknight1 commented on Apr 29, 2026

    @Rocketknight1
    Rocketknight1
    on Apr 29, 2026
    Member

    Since it's a related area, can you rebase after #45642 is merged and check the issue is still present?

  3. nurpax commented on Apr 29, 2026

    @nurpax
    nurpax
    on Apr 29, 2026
    ContributorAuthor

    @Rocketknight1 This looks related to #45642, but I think it is a different problem.

    #45642 fixes collisions inside the dynamic module cache: two different local model directories could previously map to the same cached module path. After that PR, the control case correctly gets distinct hashed cache paths, e.g.

    HF_MODULES_CACHE/transformers_modules/pretrained_a/<hash>/custom_model.py
HF_MODULES_CACHE/transformers_modules/pretrained_b/<hash>/custom_model.py

    This issue is not that pretrained_a and pretrained_b collide in the dynamic module cache. In the failing case, Transformers does not reach the dynamic module cache for pretrained_b at all. After saving/registering the local custom model earlier in the same Python process, AutoModel.from_pretrained(..., trust_remote_code=True) resolves to the already-registered local custom_model.py class instead of loading the class declared by pretrained_b's auto_map.

    So #45642 changes the expected cached source path, but the failing behavior remains:

    Expected:
    HF_MODULES_CACHE/transformers_modules/pretrained_b/<hash>/custom_model.py

    Actual:
    custom_model.py

  4. nurpax commented on Apr 29, 2026

    @nurpax
    nurpax
    on Apr 29, 2026
    ContributorAuthor

    Sorry, forgot to include in the previous comment, that I did also test on the latest main with the fixes. The issue is still present!

  5. Rocketknight1 commented on Apr 30, 2026

    @Rocketknight1
    Rocketknight1
    on Apr 30, 2026
    Member

    Yeah, this does seem like a real bug. I think it was introduced in #45094 - cc @HanFa @hmellor. I can try drafting a fix, but if anyone has touched this code recently and can see an obvious approach, let me know!

  6. sebastianvlad1 commented on Apr 30, 2026

    @sebastianvlad1
    sebastianvlad1
    on Apr 30, 2026
    Contributor

    Hi @Rocketknight1, I've already started working on this if you don't mind me taking a shot at it.

  7. sebastianvlad1 commented on Apr 30, 2026

    @sebastianvlad1
    sebastianvlad1
    on Apr 30, 2026
    Contributor

    My plan is to add a small exception in AutoConfig.from_pretrained that if the local directory already contains the module file referenced in auto_map, it should use that saved module instead of any globally registered class for the same model_type.

  8. Prachi-kushwaha commented on May 2, 2026

    @Prachi-kushwaha
    Prachi-kushwaha
    on May 2, 2026

    @nurpax i tried to reproduce this on transformers==5.7.0.dev0 with a clean environment (Python 3.12 CPU-only)
    Both the control case and the failing case behave correctly for me:

  9. nurpax commented on May 4, 2026

    @nurpax
    nurpax
    on May 4, 2026
    ContributorAuthor

    Hi @Prachi-kushwaha, I can repro the bug using the latest transformers main as follows.

    python -m venv .venv
source .venv/bin/activate
pip install torch pytest-xdist 'git+https://github.com/huggingface/transformers.git@6f90cbb5729ed94ae05d385c359fa277172663b9'
wget https://github.com/user-attachments/files/27194366/main.py
wget https://github.com/user-attachments/files/27194367/custom_model.py

python main.py save --path=pretrained_a --magic="Magic A"
python main.py save --path=pretrained_b --magic="Magic B"

# Control case, works as expected:

python main.py load-and-load --path-a=pretrained_a --path-b=pretrained_b


Load "pretrained_a"
Model says:  Magic A
Source path: HF_MODULES_CACHE/transformers_modules/pretrained_a/a7ae1d1ead4c31cc/custom_model.py
Source says: Magic A
Load "pretrained_b"
Model says:  Magic B
Source path: HF_MODULES_CACHE/transformers_modules/pretrained_b/8f15c219f628be07/custom_model.py
Source says: Magic B

# Failing case

python main.py save-and-load --save-path=pretrained_a --save-magic="Magic A" --load-path=pretrained_b


Save "pretrained_a" with "Magic A"
Model says:  <MAGIC>
Source path: custom_model.py
Source says: <MAGIC>
Load "pretrained_b"
Model says:  <MAGIC>
Source path: custom_model.py
Source says: <MAGIC>

    The "Model says", "Source says" lines do not match the outputs of the control case.

  10. hmellor commented on May 5, 2026

    @hmellor
    hmellor
    on May 5, 2026
    Member

    I'm a little confused, if you don't want AutoConfig to use the class that was explicitly registered using AutoConfig.register with model_type="custom", why are you registering it?

  11. nurpax commented on May 6, 2026

    @nurpax
    nurpax
    on May 6, 2026
    ContributorAuthor

    I think the core issue is that AutoConfig.register() / AutoModel.register() are process-global, while auto_map is checkpoint-local.

    Once a package import has registered a model_type, that registration affects every later from_pretrained() call in the same process. But the desired code-resolution policy is often per checkpoint.

    For some loads, registered/local code should win: local debugging, patched installed code, or trust_remote_code=False.

    For others, the checkpoint's auto_map should win: the checkpoint may point to the exact code snapshot saved with that checkpoint, and using the current process's registered implementation can silently load incompatible code.

    I don't think this is an unusual edge case. The docs show Auto*.register() as the way to integrate custom models with Auto classes, while register_for_auto_class() / auto_map is the way to make saved custom checkpoints reloadable/shareable. So users naturally end up in environments where:

    • a model type is globally registered in the current process, and
    • checkpoints also contain auto_map metadata pointing at checkpoint-specific code.

    The precedence between these two mechanisms - process-global registration and checkpoint-local auto_map metadata - does not appear to be clearly documented, which makes mixed use semantics ambiguous.

    Maybe the fix is not to change the default, but to expose the precedence explicitly, for example with a per-call option likeprefer_remote_code=False (or prefer_auto_map=True?), while preserving the current default behavior.

  12. sebastianvlad1 commented on May 6, 2026

    @sebastianvlad1
    sebastianvlad1
    on May 6, 2026
    Contributor

    Let me know if my PR looks good or if you’d prefer I go with an explicit  prefer_auto_map  parameter instead.

  13. sebastianvlad1 commented on May 20, 2026

    @sebastianvlad1
    sebastianvlad1
    on May 20, 2026
    Contributor

    Hi @Rocketknight1 @Cyrilvallez , just checking in to see if there's still interest in getting this fixed. Happy to help in any way!

  14. nurpax commented on May 27, 2026

    @nurpax
    nurpax
    on May 27, 2026
    ContributorAuthor

    I still think the point in my above comment stands. @sebastianvlad1, if you're interested in updating your PR, I'd go for the prefer_auto_map right away, since that's a more general solution to this issue.

    We have cases in our production training code where we could use the extra control. But I can't speak on behalf of the maintainers if they're interested in this or not.

  15. hmellor commented on May 27, 2026

    @hmellor
    hmellor
    on May 27, 2026
    Member

    For others, the checkpoint's auto_map should win: the checkpoint may point to the exact code snapshot saved with that checkpoint, and using the current process's registered implementation can silently load incompatible code.

    Actually, the reason for #45094 is because the code in checkpoints is often {stale/broken on newer Transformers} and we want the register to override it.

  16. sebastianvlad1 commented on May 29, 2026

    @sebastianvlad1
    sebastianvlad1
    on May 29, 2026
    Contributor

    @nurpax I’ve made the PR update already. If you have time, feel free to have a look and tell me if anything should be changed.

  17. github-actions commented on Jun 23, 2026

    @github-actions
    github-actionsBot
    on Jun 23, 2026 – with GitHub Actions
    Contributor

    This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

    Please note that issues that do not follow the contributing guidelines are likely to be ignored.

  18. github-actions
    closed this as completedon Jul 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bug

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions