Attribute error on _NotYetLoadedTensor loading checkpoint into quantized model

[rasbt]

Bug description

When upgrading the lightning version from 2.3.0.dev20240428 to 2.3.3, we encounter an AttributeError: '_NotYetLoadedTensor' object has no attribute 'data'.

What version are you seeing the problem on?

master

How to reproduce the bug

litgpt generate --quantize bnb.nf4 checkpoints/microsoft/phi-2

Error messages and logs

⚡ main ~/litgpt2 litgpt generate --quantize bnb.nf4 checkpoints/microsoft/phi-2 
{'checkpoint_dir': PosixPath('checkpoints/microsoft/phi-2'),
 'compile': False,
 'max_new_tokens': 50,
 'num_samples': 1,
 'precision': None,
 'prompt': 'What food do llamas eat?',
 'quantize': 'bnb.nf4',
 'temperature': 0.8,
 'top_k': 50,
 'top_p': 1.0}
Loading model 'checkpoints/microsoft/phi-2/lit_model.pth' with {'name': 'phi-2', 'hf_config': {'name': 'phi-2', 'org': 'microsoft'}, 'scale_embeddings': False, 'block_size': 2048, 'vocab_size': 50257, 'padding_multiple': 512, 'padded_vocab_size': 51200, 'n_layer': 32, 'n_head': 32, 'head_size': 80, 'n_embd': 2560, 'rotary_percentage': 0.4, 'parallel_residual': True, 'bias': True, 'lm_head_bias': True, 'n_query_groups': 32, 'shared_attention_norm': True, 'norm_class_name': 'LayerNorm', 'norm_eps': 1e-05, 'mlp_class_name': 'GptNeoxMLP', 'gelu_approximate': 'tanh', 'intermediate_size': 10240, 'rope_condense_ratio': 1, 'rope_base': 10000, 'n_expert': 0, 'n_expert_per_token': 0, 'rope_n_elem': 32}
Time to instantiate model: 0.25 seconds.
Traceback (most recent call last):
  File "/home/zeus/miniconda3/envs/cloudspace/bin/litgpt", line 8, in <module>
    sys.exit(main())
  File "/teamspace/studios/this_studio/litgpt2/litgpt/__main__.py", line 71, in main
    CLI(parser_data)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/jsonargparse/_cli.py", line 119, in CLI
    return _run_component(component, init.get(subcommand))
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/jsonargparse/_cli.py", line 204, in _run_component
    return component(**cfg)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/teamspace/studios/this_studio/litgpt2/litgpt/generate/base.py", line 255, in main
    load_checkpoint(fabric, model, checkpoint_path)
  File "/teamspace/studios/this_studio/litgpt2/litgpt/utils.py", line 362, in load_checkpoint
    model.load_state_dict(state_dict, strict=strict)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/lightning/fabric/wrappers.py", line 168, in load_state_dict
    return self._original_module.load_state_dict(state_dict=state_dict, strict=strict, **kwargs)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2139, in load_state_dict
    load(self, state_dict)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2127, in load
    load(child, child_state_dict, child_prefix)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2121, in load
    module._load_from_state_dict(
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1991, in _load_from_state_dict
    hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 72, in __call__
    return self.hook(*args, **kwargs)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/lightning/fabric/plugins/precision/bitsandbytes.py", line 166, in _quantize_on_load_hook
    quantize_fn(weight)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/lightning/fabric/plugins/precision/bitsandbytes.py", line 320, in quantize_
    if weight.data.dtype == torch.uint8:
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/lightning/fabric/utilities/load.py", line 166, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: '_NotYetLoadedTensor' object has no attribute 'data'

Environment

Current environment

* CUDA:
        - GPU:
                - NVIDIA A10G
        - available:         True
        - version:           12.1
* Lightning:
        - lightning:         2.3.3
        - lightning-cloud:   0.5.70
        - lightning-sdk:     0.1.10
        - lightning-utilities: 0.11.3.post0
        - pytorch-lightning: 2.3.3
        - torch:             2.2.1+cu121
        - torchmetrics:      1.3.1
        - torchvision:       0.17.1+cu121
* Packages:
        - absl-py:           2.1.0
        - accelerate:        0.32.1
        - aiohttp:           3.9.5
        - aiosignal:         1.3.1
        - annotated-types:   0.7.0
        - anyio:             4.4.0
        - argon2-cffi:       23.1.0
        - argon2-cffi-bindings: 21.2.0
        - arrow:             1.3.0
        - asttokens:         2.4.1
        - async-lru:         2.0.4
        - async-timeout:     4.0.3
        - attrs:             23.2.0
        - babel:             2.15.0
        - backoff:           2.2.1
        - beautifulsoup4:    4.12.3
        - bitsandbytes:      0.42.0
        - bleach:            6.1.0
        - boto3:             1.34.142
        - botocore:          1.34.142
        - cachetools:        5.3.3
        - certifi:           2024.7.4
        - cffi:              1.16.0
        - chardet:           5.2.0
        - charset-normalizer: 3.3.2
        - click:             8.1.7
        - colorama:          0.4.6
        - comm:              0.2.2
        - contourpy:         1.2.1
        - cycler:            0.12.1
        - dataproperty:      1.0.1
        - datasets:          2.20.0
        - debugpy:           1.8.2
        - decorator:         5.1.1
        - defusedxml:        0.7.1
        - dill:              0.3.8
        - dnspython:         2.6.1
        - docstring-parser:  0.16
        - email-validator:   2.2.0
        - evaluate:          0.4.2
        - exceptiongroup:    1.2.1
        - executing:         2.0.1
        - fastapi:           0.111.0
        - fastapi-cli:       0.0.4
        - fastjsonschema:    2.20.0
        - filelock:          3.15.4
        - fire:              0.6.0
        - fonttools:         4.53.1
        - fqdn:              1.5.1
        - frozenlist:        1.4.1
        - fsspec:            2024.5.0
        - google-auth:       2.32.0
        - google-auth-oauthlib: 1.2.1
        - grpcio:            1.64.1
        - h11:               0.14.0
        - hf-transfer:       0.1.6
        - httpcore:          1.0.5
        - httptools:         0.6.1
        - httpx:             0.27.0
        - huggingface-hub:   0.24.0
        - idna:              3.7
        - importlib-resources: 6.4.0
        - ipykernel:         6.26.0
        - ipython:           8.17.2
        - ipywidgets:        8.1.1
        - isoduration:       20.11.0
        - jedi:              0.19.1
        - jinja2:            3.1.4
        - jmespath:          1.0.1
        - joblib:            1.4.2
        - json5:             0.9.25
        - jsonargparse:      4.31.0
        - jsonlines:         4.0.0
        - jsonpointer:       3.0.0
        - jsonschema:        4.23.0
        - jsonschema-specifications: 2023.12.1
        - jupyter-client:    8.6.2
        - jupyter-core:      5.7.2
        - jupyter-events:    0.10.0
        - jupyter-lsp:       2.2.5
        - jupyter-server:    2.14.1
        - jupyter-server-terminals: 0.5.3
        - jupyterlab:        4.2.0
        - jupyterlab-pygments: 0.3.0
        - jupyterlab-server: 2.27.2
        - jupyterlab-widgets: 3.0.11
        - kiwisolver:        1.4.5
        - lightning:         2.3.3
        - lightning-cloud:   0.5.70
        - lightning-sdk:     0.1.10
        - lightning-utilities: 0.11.3.post0
        - litdata:           0.2.17
        - litgpt:            0.4.5
        - litserve:          0.1.3
        - lm-eval:           0.4.3
        - lxml:              5.2.2
        - markdown:          3.6
        - markdown-it-py:    3.0.0
        - markupsafe:        2.1.5
        - matplotlib:        3.8.2
        - matplotlib-inline: 0.1.7
        - mbstrdecoder:      1.1.3
        - mdurl:             0.1.2
        - mistune:           3.0.2
        - more-itertools:    10.3.0
        - mpmath:            1.3.0
        - multidict:         6.0.5
        - multiprocess:      0.70.16
        - nbclient:          0.10.0
        - nbconvert:         7.16.4
        - nbformat:          5.10.4
        - nest-asyncio:      1.6.0
        - networkx:          3.3
        - nltk:              3.8.1
        - notebook-shim:     0.2.4
        - numexpr:           2.10.1
        - numpy:             1.26.4
        - nvidia-cublas-cu12: 12.1.3.1
        - nvidia-cuda-cupti-cu12: 12.1.105
        - nvidia-cuda-nvrtc-cu12: 12.1.105
        - nvidia-cuda-runtime-cu12: 12.1.105
        - nvidia-cudnn-cu12: 8.9.2.26
        - nvidia-cufft-cu12: 11.0.2.54
        - nvidia-curand-cu12: 10.3.2.106
        - nvidia-cusolver-cu12: 11.4.5.107
        - nvidia-cusparse-cu12: 12.1.0.106
        - nvidia-nccl-cu12:  2.19.3
        - nvidia-nvjitlink-cu12: 12.5.82
        - nvidia-nvtx-cu12:  12.1.105
        - oauthlib:          3.2.2
        - orjson:            3.10.6
        - overrides:         7.7.0
        - packaging:         24.1
        - pandas:            2.1.4
        - pandocfilters:     1.5.1
        - parso:             0.8.4
        - pathvalidate:      3.2.0
        - peft:              0.11.1
        - pexpect:           4.9.0
        - pillow:            10.4.0
        - pip:               24.1.2
        - platformdirs:      4.2.2
        - portalocker:       2.10.1
        - prometheus-client: 0.20.0
        - prompt-toolkit:    3.0.47
        - protobuf:          4.23.4
        - psutil:            6.0.0
        - ptyprocess:        0.7.0
        - pure-eval:         0.2.2
        - pyarrow:           17.0.0
        - pyarrow-hotfix:    0.6
        - pyasn1:            0.6.0
        - pyasn1-modules:    0.4.0
        - pybind11:          2.13.1
        - pycparser:         2.22
        - pydantic:          2.8.2
        - pydantic-core:     2.20.1
        - pygments:          2.18.0
        - pyjwt:             2.8.0
        - pyparsing:         3.1.2
        - pytablewriter:     1.2.0
        - python-dateutil:   2.9.0.post0
        - python-dotenv:     1.0.1
        - python-json-logger: 2.0.7
        - python-multipart:  0.0.9
        - pytorch-lightning: 2.3.3
        - pytz:              2024.1
        - pyyaml:            6.0.1
        - pyzmq:             26.0.3
        - referencing:       0.35.1
        - regex:             2024.5.15
        - requests:          2.32.3
        - requests-oauthlib: 2.0.0
        - rfc3339-validator: 0.1.4
        - rfc3986-validator: 0.1.1
        - rich:              13.7.1
        - rouge-score:       0.1.2
        - rpds-py:           0.19.0
        - rsa:               4.9
        - s3transfer:        0.10.2
        - sacrebleu:         2.4.2
        - safetensors:       0.4.3
        - scikit-learn:      1.3.2
        - scipy:             1.11.4
        - send2trash:        1.8.3
        - sentencepiece:     0.2.0
        - setuptools:        69.5.1
        - shellingham:       1.5.4
        - simple-term-menu:  1.6.4
        - six:               1.16.0
        - sniffio:           1.3.1
        - soupsieve:         2.5
        - sqlitedict:        2.1.0
        - stack-data:        0.6.3
        - starlette:         0.37.2
        - sympy:             1.13.0
        - tabledata:         1.3.3
        - tabulate:          0.9.0
        - tcolorpy:          0.1.6
        - tensorboard:       2.15.1
        - tensorboard-data-server: 0.7.2
        - termcolor:         2.4.0
        - terminado:         0.18.1
        - threadpoolctl:     3.5.0
        - tinycss2:          1.3.0
        - tokenizers:        0.19.1
        - tomli:             2.0.1
        - torch:             2.2.1+cu121
        - torchmetrics:      1.3.1
        - torchvision:       0.17.1+cu121
        - tornado:           6.4.1
        - tqdm:              4.66.4
        - tqdm-multiprocess: 0.0.11
        - traitlets:         5.14.3
        - transformers:      4.42.4
        - triton:            2.2.0
        - typepy:            1.3.2
        - typer:             0.12.3
        - types-python-dateutil: 2.9.0.20240316
        - typeshed-client:   2.7.0
        - typing-extensions: 4.12.2
        - tzdata:            2024.1
        - ujson:             5.10.0
        - uri-template:      1.3.0
        - urllib3:           2.2.2
        - uvicorn:           0.30.1
        - uvloop:            0.19.0
        - watchfiles:        0.22.0
        - wcwidth:           0.2.13
        - webcolors:         24.6.0
        - webencodings:      0.5.1
        - websocket-client:  1.8.0
        - websockets:        12.0
        - werkzeug:          3.0.3
        - wheel:             0.43.0
        - widgetsnbextension: 4.0.11
        - word2number:       1.1
        - xxhash:            3.4.1
        - yarl:              1.9.4
        - zstandard:         0.23.0
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - 
        - processor:         x86_64
        - python:            3.10.10
        - release:           5.15.0-1064-aws
        - version:           #70~20.04.1-Ubuntu SMP Fri Jun 14 15:42:13 UTC 2024

More info

Now, one thing to keep in mind is that we have the bitsandbytes version pinned to bitsandbytes==0.42.0 because 0.43 results in the following issue:

/teamspace/studios/this_studio/litgpt2/litgpt/generate/base.py:207: UserWarning: LitGPT only supports bitsandbytes v0.42.0. This may result in errors when using quantization.
  warnings.warn(
Loading model 'checkpoints/microsoft/phi-2/lit_model.pth' with {'name': 'phi-2', 'hf_config': {'name': 'phi-2', 'org': 'microsoft'}, 'scale_embeddings': False, 'block_size': 2048, 'vocab_size': 50257, 'padding_multiple': 512, 'padded_vocab_size': 51200, 'n_layer': 32, 'n_head': 32, 'head_size': 80, 'n_embd': 2560, 'rotary_percentage': 0.4, 'parallel_residual': True, 'bias': True, 'lm_head_bias': True, 'n_query_groups': 32, 'shared_attention_norm': True, 'norm_class_name': 'LayerNorm', 'norm_eps': 1e-05, 'mlp_class_name': 'GptNeoxMLP', 'gelu_approximate': 'tanh', 'intermediate_size': 10240, 'rope_condense_ratio': 1, 'rope_base': 10000, 'n_expert': 0, 'n_expert_per_token': 0, 'rope_n_elem': 32}
Time to instantiate model: 0.24 seconds.
Traceback (most recent call last):
  File "/home/zeus/miniconda3/envs/cloudspace/bin/litgpt", line 8, in <module>
    sys.exit(main())
  File "/teamspace/studios/this_studio/litgpt2/litgpt/__main__.py", line 71, in main
    CLI(parser_data)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/jsonargparse/_cli.py", line 119, in CLI
    return _run_component(component, init.get(subcommand))
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/jsonargparse/_cli.py", line 204, in _run_component
    return component(**cfg)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/teamspace/studios/this_studio/litgpt2/litgpt/generate/base.py", line 252, in main
    model = fabric.setup_module(model)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 308, in setup_module
    module = self._move_model_to_device(model=module, optimizers=[])
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 976, in _move_model_to_device
    model = self.to_device(model)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 526, in to_device
    self._strategy.module_to_device(obj)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/lightning/fabric/strategies/single_device.py", line 59, in module_to_device
    module.to(self.root_device)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1152, in to
    return self._apply(convert)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 802, in _apply
    module._apply(fn)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 825, in _apply
    param_applied = fn(param)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1150, in convert
    return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 324, in to
    return self._quantize(device)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 289, in _quantize
    w_4bit, quant_state = bnb.functional.quantize_4bit(
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/bitsandbytes/functional.py", line 1234, in quantize_4bit
    raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {A.dtype}")
ValueError: Blockwise quantization only supports 16/32-bit floats, but got torch.uint8

I wonder if these are issues that can or need to be addressed one at a time. Supporting lightning 2.3.3 with bitsandbytes==0.42.0 first to restore the litgpt quantization as is, and then see how we can upgrade to the most recent bitsandbytes version.

Any thoughts?

cc @awaelchli @carmocca

[rasbt]

github-actions bot added the ver: 2.2.x label

It's actually v2.3.3