Attribute error on _NotYetLoadedTensor loading checkpoint into quantized model #20119
Description
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-2Error 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?