Load dynamic module (remote code) only once if code isn't change

[XuehaiPan]

What does this PR do?

Fixes #30370 (comment)

Add an indicator __transformers_module_hash__ to the remote code module.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ydshieh

[ydshieh]

Hi @XuehaiPan, thank you for taking the action.

My first question is: what if the target (remote code) is changed since the last time it has been loaded.

I don't know very well what is the context mentioned in #30370 (comment), but IIRD, that happens within a single python process (i.e. multiple loading of the same module). And if that is the case, my question on top of this comment have to be considered, right?

Maybe @tmm1 could explain a bit more the situation, ideally with a code snippet to demonstrate the issue.

[XuehaiPan]

My first question is: what if the target (remote code) is changed since the last time it has been loaded.

@ydshieh Thanks for raising this. I change this indicator to the hash of the source code.

[ydshieh]

Nice. I still prefer for @tmm1 to elaborate the issue a bit more though.

But in the meantime, could you tell me more about the usage of threading.Lock() in this PR 🙏 ?

[XuehaiPan]

But in the meantime, could you tell me more about the usage of threading.Lock() in this PR 🙏 ?

It's just for thread safety. We are modifying the global variable sys.modules.

[tmm1]

prefer for @tmm1 to elaborate the issue a bit more though

sure, here is a simple example:

import sys
from accelerate import init_empty_weights
from transformers import AutoModelForCausalLM

def load_modeling_code(model_name):
    with init_empty_weights():
        model = AutoModelForCausalLM.from_pretrained(
            model_name, trust_remote_code=True
        )
        return sys.modules[model.__class__.__module__]

model_name = "deepseek-ai/DeepSeek-Coder-V2-Lite-Base"

# 1. patch modeling code
mod = load_modeling_code(model_name)
import liger_kernel
mod.CrossEntropyLoss = liger_kernel.transformers.cross_entropy.LigerCrossEntropyLoss
mod.DeepseekV2RMSNorm = liger_kernel.transformers.rms_norm.LigerRMSNorm
mod.DeepseekV2MLP.forward = liger_kernel.transformers.swiglu.LigerSwiGLUMLP.forward

# 2. check patch
mod = load_modeling_code(model_name)
print(mod.CrossEntropyLoss == liger_kernel.transformers.cross_entropy.LigerCrossEntropyLoss)

# 3. create and use model
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)

currently on main branch, the patch is missing in step (2) and (3)

[XuehaiPan]

It looks reasonable to assume the import statement returns the cached module object if it has already been imported.

Ideally, we may need to simulate this:

sys.path.insert(0, HF_MODULES_CACHE)
module = importlib.import_module(name)  # returns the same module object if it has already been imported (i.e. return sys.modules[name])
assert sys.path[0] == HF_MODULES_CACHE
del sys.path[0]

[ydshieh]

Hi @tmm1 Thank you for providing the detailed information.

I am wondering if it would make (more) sense for the patching code to have a factory method like

def load_modeling_code(model_name):
    with init_empty_weights():
        model = AutoModelForCausalLM.from_pretrained(
            model_name, trust_remote_code=True
        )
        return sys.modules[model.__class__.__module__]

def patch_modeling_code(model_name):

    # 1. patch modeling code
    mod = load_modeling_code(model_name)
    import liger_kernel
    mod.CrossEntropyLoss = liger_kernel.transformers.cross_entropy.LigerCrossEntropyLoss
    mod.DeepseekV2RMSNorm = liger_kernel.transformers.rms_norm.LigerRMSNorm
    mod.DeepseekV2MLP.forward = liger_kernel.transformers.swiglu.LigerSwiGLUMLP.forward

    return mod

model_name = "deepseek-ai/DeepSeek-Coder-V2-Lite-Base"
mod = patch_modeling_code(model_name)

But I am open to what @XuehaiPan have done so far.

[LysandreJik]

cc @Rocketknight1 as well :)

[tmm1]

I am wondering if it would make (more) sense for the patching code to have a factory method like

it could work, but often the user wants to apply some patch then invoke a well known trainer framework. such a framework would not accept a factory callback

more importantly, the factory would still need to call AutoModelForCausalLM.from_pretrained two times:

1. to download and load the remote code, so it is available inside sys.modules for patching
2. without init_empty_weights(), and to actually use the patched code

without this PR, the second invocation will remove the patch. so it would be impossible to get a model object that actually is patched.

[ydshieh]

OK, I didn't know

# 3. create and use model
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)

in your previous comment is actually needed. Makes sense now.

I will let @Rocketknight1 to have a look on the changes of this PR too.

[Rocketknight1]

It seems clean to me! I can think of two failure cases:

1. This will not detect changes in other files (for example, if the model also has utils/functions in a separate file from the main modeling file, the hash of the main modeling file will not change and so it will not be reimported.
2. I think this could cause a regression if users load a remote_code model, monkey-patch methods, and then reload the model to clear their changes, since now their changes will persist. However, I suspect there are no users depending on this weird behaviour, so it's not a serious problem.

Both of these are very minor issues that cannot really be fixed in this PR, and I don't think they should block it, so I'm happy with it!

[ydshieh]

Hi @XuehaiPan It would be great if you can add a corresponding test, like a (simple version of) patched module (like what @tmm1 has) will remain the same.

[XuehaiPan]

It would be great if you can add a corresponding test

Done.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ydshieh]

very nice. a few nit comments

[ydshieh]

In my previous comment

        # The configuration file is cached in the snapshot directory. So the module file is not changed after dumping

in other files (other than tests/models/auto/test_configuration_auto.py)

here is not about configuration

I can go with the current version, just let me know if you want to add it back (with the correct comments).

What I mean is that we have to change the comment The configuration file to the actual file name (like the image processor file). Otherwise, that part of test is actually ok.

[XuehaiPan]

I can go with the current version, just let me know if you want to add it back (with the correct comments).

@ydshieh I add a commit to address this.

[ydshieh]

@Rocketknight1 in case if you want to take a final look. (and/or feel free to ping a core maintainer 🙏 )

[Rocketknight1]

I'm happy with it, I think! The test failure seems unrelated, and I like the core goal of allowing the same model to be imported twice without wasting time, and without getting two different output classes.

cc @XuehaiPan maybe rebase onto main and see if that fixes the test?

cc @LysandreJik for core maintainer approval!

[LysandreJik]

Thanks! This looks good to me if approved by the esteemed @Rocketknight1

[Rocketknight1]

@XuehaiPan let me know if you're happy to merge, or if there's anything else you want to tweak first!

[XuehaiPan]

@XuehaiPan let me know if you're happy to merge, or if there's anything else you want to tweak first!

@Rocketknight1 I think this is the final version of the PR.

[Rocketknight1]

Merging, in that case. Thank you for the PR!