Moving to lazy root imports to make config loading snappy #2073
Conversation
for more information, see https://pre-commit.ci
|
@JackUrb seems some thunder tests are failing: |
|
I think this is pretty dubious. It starts with not capturing the side-effects of importing as caught by the CI, but probably also impacts typecheckers etc. If you have the - arguably somewhat special - need to import litgpt.config without importing litgpt, how about you add the litgpt path to PYTHONPATH and import config that way? |
From what I understand, relying on import side-effects is frequently considered non-pythonic, and even the most common pattern using them (the registry pattern) can get away with doing an explicit discovery call instead. For typecheckers, this pattern works perfectly fine with in case some other typecheckers rely on static information. The way the thunder test does this right now is really not transparent - without knowing that
This isn't terrible to do for a one-off case, but it makes litgpt imports always a bit weird. Long-term it's a stumbling block whenever someone on our team is setting up their environment. |
for more information, see https://pre-commit.ci
|
Local testing against thunder's |
My math would be that 99% of litgpt users don't need the "don't import config only". So the remaining 1% could be served by adding a localized hack for their use: import os, importlib
def import_litgpt_config():
litgpt_spec = importlib.util.find_spec('litgpt')
config_path = os.path.join(litgpt_spec.submodule_search_locations[0], 'config.py')
config_spec = importlib.util.spec_from_file_location('litgpt_config', config_path)
module = importlib.util.module_from_spec(config_spec)
config_spec.loader.exec_module(module)
return module
litgpt_config = import_litgpt_config()which is a straightforward adaptation of the pertinent example for importing a source file directly in the official Python documentation. Maybe we could add this as a recipe to the top of the config file and call it a day. (I don't disagree with pushing the torch import to the single thing using it btw.) WDYT? |
|
cc: @lantiga ^^ |
|
We don't just use Is there still anything preventing the lazy solution from covering the original use cases and your concerns otherwise? I don't disagree that most users will never need this particular functionality, but I don't see what prevents this from being useable for the 100%. |
Like or do you need anything more elaborate to get the typechecking? I'm trying to find the right balance here. The PR as is prevents tab completion in ipython after |
I love your perspective :) Really appreciated 👍 |
Ah this is a good catch, I didn't realize that interactive shells rely on
It's just a little more elaborate - to get things correct our module needs to expose anything we may use as a type directly, as we can't use the imported module to typehint things ( # fast_litgpt.py
import importlib
import os
from typing import TYPE_CHECKING
def import_litgpt_config():
litgpt_spec = importlib.util.find_spec("litgpt")
config_path = os.path.join(litgpt_spec.submodule_search_locations[0], "config.py")
config_spec = importlib.util.spec_from_file_location("litgpt_config", config_path)
module = importlib.util.module_from_spec(config_spec)
config_spec.loader.exec_module(module)
return module
def import_litgpt_args():
litgpt_spec = importlib.util.find_spec("litgpt")
args_path = os.path.join(litgpt_spec.submodule_search_locations[0], "args.py")
args_spec = importlib.util.spec_from_file_location("litgpt_args", args_path)
module = importlib.util.module_from_spec(args_spec)
args_spec.loader.exec_module(module)
return module
litgpt_config = import_litgpt_config()
litgpt_args = import_litgpt_args()
if TYPE_CHECKING:
import litgpt.args
import litgpt.config
litgpt_config = litgpt.config
litgpt_args = litgpt.args
TrainArgs = litgpt_args.TrainArgs
EvalArgs = litgpt_args.EvalArgs
Config = litgpt_config.Config
__all__ = [
"litgpt_config",
"litgpt_args",
"TrainArgs",
"EvalArgs",
"Config",
]This then seems to work as desired, so I'm ok with doing this even if it's something we'll have to drill to folks internally (don't import from Personally I still feel like there's value in getting this to work out-of-box, but understood if y'all are worried about compatibility risks. I'll open a separate PR just making the torch import in |
|
Problem
While the goal of #2034 was great, it didn't go far enough to ensure that imports would actually avoid pulling the whole litgpt tree. At the moment the behavior is something like:
For simple scripts that queue jobs elsewhere, this import delay is pretty brutal. The crux of the issue is in the
litgpt.__init__module, which is always getting imported no matter what.Changes
This PR edits
__init__.pyto move all of the key imports into a__getattr__, such that the imports can be done lazily when someone actually wants tofrom litgpt import LLMor similar, but doesn't cause all these imports to trigger when you're trying to just get an arg or config class.It also delays the load of
torchinside ofConfigfor similar reasons, though this is less critical (saves a few seconds).Testing
Config imports are (comparatively) blazing fast:
Can still import everything with the old semantics: