Moving to lazy root imports to make config loading snappy

[JackUrb]

Problem

While the goal of https://github.com/Lightning-AI/litgpt/pull/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:

from litgpt.args import EvalArgs
...
import time:      3605 |   11016312 |     litgpt.api
import time:      1768 |   11032303 |   litgpt
import time:      1919 |   11034222 | litgpt.args

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__.py to move all of the key imports into a __getattr__, such that the imports can be done lazily when someone actually wants to from litgpt import LLM or 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 torch inside of Config for similar reasons, though this is less critical (saves a few seconds).

Testing

Config imports are (comparatively) blazing fast:

>>> from litgpt import Config
...
import time:     10049 |     186252 | litgpt.config

# new shell
>>> from litgpt.args import EvalArgs
...
import time:      2795 |      28360 | litgpt.args

Can still import everything with the old semantics:

>>> from litgpt import *
...
>>> LLM
<class 'litgpt.api.LLM'>
>>> Config
<class 'litgpt.config.Config'>
>>> GPT
<class 'litgpt.model.GPT'>
>>> PromptStyle
<class 'litgpt.prompts.PromptStyle'>
>>> Tokenizer
<class 'litgpt.tokenizer.Tokenizer'>

[Borda]

@JackUrb seems some thunder tests are failing:

=========================== short test summary info ============================
FAILED tests/ext_thunder/test_thunder_networks.py::test_quantization - AttributeError: module 'litgpt' has no attribute 'config'
FAILED tests/ext_thunder/test_thunder_networks.py::test_memory_litgpt_llama3 - AttributeError: module 'litgpt' has no attribute 'config'
FAILED tests/ext_thunder/test_thunder_networks.py::test_checkpointing_thunderfx - AttributeError: module 'litgpt' has no attribute 'config'

[t-vi]

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?

[JackUrb]

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.

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 pylance at least, but you can be more explicit by doing an:

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    # actually do all the imports

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 litgpt is causing a cascade of imports, how would one know that the config submodule is going to be available? Still, this can be resolved for now by just adding config (and the other previously imported submodules) to the lazy import registry to preserve the old behavior that it would be available.

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?

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.

[JackUrb]

Local testing against thunder's litgpt_model imports after my most recent change:

Python 3.11.12 (main, Apr  9 2025, 04:04:00) [Clang 20.1.0 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import litgpt
>>> litgpt.config.name_to_config.update({})
>>> Config = litgpt.Config
>>> GPT = litgpt.GPT
>>> RMSNorm = litgpt.model.RMSNorm
>>> CausalSelfAttention = litgpt.model.CausalSelfAttention
>>> LLaMAMLP = litgpt.model.LLaMAMLP
>>> build_rope_cache = litgpt.model.build_rope_cache
>>> apply_rope = litgpt.model.apply_rope
>>> Block = litgpt.model.Block
>>>

[t-vi]

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?

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.

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?

[Borda]

cc: @lantiga ^^

[JackUrb]

We don't just use litgpt.config, we also use litgpt.args (so we'd have to patch both). And unfortunately this solution would leave us with no longer having the nice typechecking when we want to use these (without another layer of hacking at least).

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%.

[t-vi]

We don't just use litgpt.config, we also use litgpt.args (so we'd have to patch both). And unfortunately this solution would leave us with no longer having the nice typechecking when we want to use these (without another layer of hacking at least).

Like

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import litgpt.config as litgpt_config

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 import litgpt.

[Borda]

I'm trying to find the right balance here. The PR as is prevents tab completion in ipython after import litgpt.

I love your perspective :) Really appreciated 👍

[JackUrb]

The PR as is prevents tab completion in ipython after import litgpt.

Ah this is a good catch, I didn't realize that interactive shells rely on __dir__ for tab autocomplete. Overriding this things now work (following the module load delay) with ipython tab complete, though I'll admit the delay coming later in this circumstance feels a bit unexpected. (Still, this could be avoided by importing any of the modules that trigger the deeper loads before trying to tab autocomplete).

or do you need anything more elaborate to get the typechecking?

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 (train_args: litgpt_args.TrainArgs isn't happy and causes train_args to have no type hinting anywhere). We could go the route of having TYPE_CHECKING imports everywhere, and use the quoted classes (train_arg: "TrainArgs") to remedy this, but that'd be clunky. Instead I can get thing together with the following, exposing everything we use at the top level:

# 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 litgpt in any config/launcher scripts, import from fast_litgpt for these!).

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 config lazy, to keep this PR as an archive of the more complete approach in case it's something you want to reinvestigate.

[Borda]

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 config lazy, to keep this PR as an archive of the more complete approach in case it's something you want to reinvestigate.

@t-vi @lantiga what is your judgment?

[Borda]

making it as draft until we hear back from @lantiga or @t-vi