NVIDIA-NeMo / NVIDIA-NeMo/Automodel
Checkpoint loading fails with quantized base weights (QLoRA)
@adil-a is already working on this.
Since Jan 12, 2026.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.