NVIDIA-NeMo / NVIDIA-NeMo/Automodel

Checkpoint loading fails with quantized base weights (QLoRA)

Open
#1,020 0 comments 0 reactions 1 assignee View on GitHub

@adil-a is already working on this.

Since Jan 12, 2026.

bug ckpt
Dominant language
Python
Stars
963
Forks
318
Avg merge
3d 20h
Merged PRs (30d)
143

Description

Describe the bug
As shown in the attached stack trace, currently the checkpointer ignores completely the quantization parameter, which causes it to restore the full accuracy weights (leading to OOM in this case). Instead, during loading it should quantize weights on-the-fly and avoid OOMing.

Steps/Code to reproduce bug

Please list minimal steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

Expected behavior

A clear and concise description of what you expected to happen.

Additional context

[rank0]:   File "/mnt/4tb/auto/Automodel3/examples/llm_finetune/finetune.py", line 33, in <module>█████████████████████████████████████████████████████████████████████████████████████████████████████▊  | 4.90G/4.97G [02:38<00:00, 112MB/s]
[rank0]:     main()
[rank0]:   File "/mnt/4tb/auto/Automodel3/examples/llm_finetune/finetune.py", line 28, in main
[rank0]:     recipe.setup()
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/recipes/llm/train_ft.py", line 1011, in setup
[rank0]:     model, model_state_dict_keys, self.optimizer, self.loss_fn, self.param_info = build_model_and_optimizer(
[rank0]:                                                                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/recipes/llm/train_ft.py", line 319, in build_model_and_optimizer
[rank0]:     checkpointer.load_base_model(
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/components/checkpoint/checkpointing.py", line 342, in load_base_model
[rank0]:     to_empty_parameters_only(model, device=device)
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/components/checkpoint/checkpointing.py", line 716, in to_empty_parameters_only
[rank0]:     return _apply(model, lambda t: torch.empty_like(t, device=device, dtype=dtype), recurse=recurse)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/components/checkpoint/checkpointing.py", line 781, in _apply
[rank0]:     _apply(child, fn, recurse=recurse)
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/components/checkpoint/checkpointing.py", line 781, in _apply
[rank0]:     _apply(child, fn, recurse=recurse)
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/components/checkpoint/checkpointing.py", line 781, in _apply
[rank0]:     _apply(child, fn, recurse=recurse)
[rank0]:   [Previous line repeated 2 more times]
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/components/checkpoint/checkpointing.py", line 805, in _apply
[rank0]:     param_applied = fn(param)
[rank0]:                     ^^^^^^^^^
[rank0]:   File "/mnt/4tb/auto/Automodel3/nemo_automodel/components/checkpoint/checkpointing.py", line 716, in <lambda>
[rank0]:     return _apply(model, lambda t: torch.empty_like(t, device=device, dtype=dtype), recurse=recurse)
[rank0]:                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/4tb/auto/Automodel3/.venv/lib/python3.12/site-packages/torch/_prims_common/wrappers.py", line 309, in _fn
[rank0]:     result = fn(*args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/4tb/auto/Automodel3/.venv/lib/python3.12/site-packages/torch/_refs/__init__.py", line 5055, in empty_like
[rank0]:     return torch.empty_permuted(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 896.00 MiB. GPU 0 has a total capacity of 47.50 GiB of which 176.62 MiB is free. Including non-PyTorch memory, this process has 47.32 GiB memory in use. Of the allocated memory 46.98 GiB is allocated by PyTorch, and 11.31 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.