load_best_model_at_end reloads PEFT adapter weights onto CUDA and can OOM under low remaining GPU memory

[DogWala]

System Info

• transformers version: local current checkout (5.3.0.dev0)
• Python: 3.12.12
• PyTorch: 2.10.0+cu128
• CUDA available: True
• CUDA device count: 8
• torchvision: 0.25.0+cu128
• Pillow: 12.1.1
• PEFT: 0.18.1

I can also provide the full transformers env output if needed.

Who can help?

@SunMarc @BenjaminBossan

Information

The problem arises when using:

• TrainingArguments(load_best_model_at_end=True)
• a PEFT / LoRA model
• the built-in best checkpoint reload path at the end of training
• a situation where very little GPU memory remains at that point

This is not based on the original private training pipeline. I reduced it to a minimal local repro using:

• local transformers.Trainer
• local peft
• local Qwen3.5-0.8B
• local GSM8K parquet files

Keeping GPU memory usage relatively high during training is a common and intentional setup when trying to maximize utilization.

Under the same training setup, disabling load_best_model_at_end does not cause a failure, while enabling it can trigger a late OOM during the best-model reload stage that should ideally not happen.

Tasks

This is not from an official example script. It is a minimal self-contained repro intended to isolate the bug.

Reproduction

I reduced this to a minimal repro in a separate folder that does not depend on the original training pipeline.

The key idea is:

1. train a PEFT / LoRA model with load_best_model_at_end=True
2. keep tail GPU memory very tight near the end of training
3. let the built-in best-model reload happen
4. observe that PEFT adapter loading goes directly to CUDA and can OOM at that stage

The one-shot repro script runs two cases:

1. default built-in best-model reload path
2. CPU-safe control run

The default repro parameters are:

• tail_margin_mb=32
• allocator_slack_mb=300
• lora_r=2048
• lora_alpha=4096
• train_samples=2
• eval_samples=2
• max_steps=1
• eval_steps=1
• max_length=128
• dtype=fp16

Command:

bash run_oom_repro_latest_transformers.sh

The default path runs:

python oom_repro_large_tail_qwen35_gsm8k.py \
  --single-run \
  --tail-margin-mb 32 \
  --allocator-slack-mb 300 \
  --best-load-device default \
  --lora-r 2048 \
  --lora-alpha 4096 \
  --model-path <LOCAL_MODEL_PATH> \
  --gsm8k-root <LOCAL_GSM8K_ROOT> \
  --output-root <LOCAL_OUTPUT_ROOT>/oom_runs_default \
  --train-samples 2 \
  --eval-samples 2 \
  --max-steps 1 \
  --eval-steps 1 \
  --max-length 128 \
  --dtype fp16

The control run only changes the best-load device to CPU:

python oom_repro_large_tail_qwen35_gsm8k.py \
  --single-run \
  --tail-margin-mb 32 \
  --allocator-slack-mb 300 \
  --best-load-device cpu \
  --lora-r 2048 \
  --lora-alpha 4096 \
  --model-path <LOCAL_MODEL_PATH> \
  --gsm8k-root <LOCAL_GSM8K_ROOT> \
  --output-root <LOCAL_OUTPUT_ROOT>/oom_runs_cpu \
  --train-samples 2 \
  --eval-samples 2 \
  --max-steps 1 \
  --eval-steps 1 \
  --max-length 128 \
  --dtype fp16

Observed result:

• default path: status = "cuda_oom"
• CPU-safe control path: status = "ok"

The important part is that training itself finishes. The failure happens during the built-in best-model reload stage.

This is also the key reason I consider this a bug rather than just expected memory pressure: under the same settings, training without load_best_model_at_end is fine, but enabling load_best_model_at_end introduces a late OOM because adapter weights are reloaded onto CUDA at the end.

From the recorded events in the default-path repro:

• trainer._load_best_model.enter
  • is_in_train = true
  • optimizer_is_none = false
  • lr_scheduler_is_none = false
  • callback_optimizer_is_none = false
• peft.load_adapter.enter
  • torch_device = None
• peft.load_peft_weights.enter
  • device = "cuda"

So under low remaining GPU memory, the PEFT best-load path observed in this repro is loading adapter weights onto CUDA during load_best_model_at_end, and that can trigger a late OOM.

Key event excerpts from the failing run:

{"event": "trainer._load_best_model.enter", "is_in_train": true, "optimizer_is_none": false, "lr_scheduler_is_none": false, "callback_optimizer_is_none": false, "callback_lr_scheduler_is_none": false}

{"event": "peft.load_adapter.enter", "torch_device": "None"}

{"event": "peft.load_peft_weights.enter", "device": "cuda"}

The actual OOM in the failing repro was:

CUDA out of memory. Tried to allocate 32.00 MiB. GPU 0 has a total capacity of 47.40 GiB of which 17.19 MiB is free.
Including non-PyTorch memory, this process has 47.38 GiB memory in use.
Of the allocated memory 46.96 GiB is allocated by PyTorch, and 93.13 MiB is reserved by PyTorch but unallocated.

Full traceback from the failing default-path run:

Traceback (most recent call last):
  File "<LOCAL_REPRO_ROOT>/oom_repro_large_tail_qwen35_gsm8k.py", line 226, in run_single
    train_result = trainer.train()
                   ^^^^^^^^^^^^^^^
  File "<LOCAL_TRANSFORMERS_ROOT>/src/transformers/trainer.py", line 1424, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "<LOCAL_TRANSFORMERS_ROOT>/src/transformers/trainer.py", line 1522, in _inner_training_loop
    return self._finalize_training(trial, num_train_samples, start_time)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<LOCAL_TRANSFORMERS_ROOT>/src/transformers/trainer.py", line 1841, in _finalize_training
    self._load_best_model()
  File "<LOCAL_REPRO_ROOT>/oom_repro_large_tail_qwen35_gsm8k.py", line 142, in _load_best_model
    return super()._load_best_model()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<LOCAL_REPRO_ROOT>/mini_repro_qwen35_gsm8k.py", line 259, in _load_best_model
    out = super()._load_best_model()
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<LOCAL_TRANSFORMERS_ROOT>/src/transformers/trainer.py", line 3449, in _load_best_model
    model.load_adapter(self.state.best_model_checkpoint, active_adapter)
  File "<LOCAL_REPRO_ROOT>/mini_repro_qwen35_gsm8k.py", line 185, in wrapped_load_adapter
    return original_load_adapter(self, model_id, adapter_name, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<LOCAL_VENV_SITE_PACKAGES>/peft/peft_model.py", line 1362, in load_adapter
    adapters_weights = load_peft_weights(
                       ^^^^^^^^^^^^^^^^^^
  File "<LOCAL_REPRO_ROOT>/mini_repro_qwen35_gsm8k.py", line 200, in wrapped_load_peft_weights
    out = original_load_peft_weights(model_id, device=device, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<LOCAL_VENV_SITE_PACKAGES>/peft/utils/save_and_load.py", line 693, in load_peft_weights
    adapters_weights = safe_load_file(filename, device=device)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<LOCAL_VENV_SITE_PACKAGES>/safetensors/torch.py", line 338, in load_file
    result[k] = f.get_tensor(k)
                ^^^^^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 32.00 MiB. GPU 0 has a total capacity of 47.40 GiB of which 17.19 MiB is free. Including non-PyTorch memory, this process has 47.38 GiB memory in use. Of the allocated memory 46.96 GiB is allocated by PyTorch, and 93.13 MiB is reserved by PyTorch but unallocated.

For comparison, under the same memory pressure, using a CPU-safe reload path succeeds:

model.load_adapter(
    ...
    torch_device="cpu",
    ...
)

In the successful control run:

• peft.load_peft_weights.enter shows device = "cpu"
• the run completes successfully
• status = "ok"

Expected behavior

When load_best_model_at_end=True is used with a PEFT model, the best-model reload stage should not fail solely because adapter weights requires more GPU memory to reload on GPU.

Possible fix

One possible fix would be to make the PEFT load_best_model_at_end path follow a more memory-safe loading pattern similar to the non-PEFT model case: load the best model state on CPU first like:

model.load_adapter(
    self.state.best_model_checkpoint,
    active_adapter,
    torch_device="cpu",
    autocast_adapter_dtype=False,
)

oom_repro_large_tail_qwen35_gsm8k.py
run_oom_repro_latest_transformers.sh

[pankajbaid567]

Hi, I'd like to work on this issue.

From the repro and stack trace it appears that Trainer._load_best_model() calls model.load_adapter() without specifying torch_device, which results in PEFT loading adapter weights directly on CUDA.

Under low GPU memory this causes an OOM during the best-model reload stage.

My current idea is to modify the PEFT path in _load_best_model() to load adapters on CPU first using:

model.load_adapter(..., torch_device="cpu")

and then move the model back to the training device.

I'll investigate whether this needs to be guarded behind a PEFT check or handled generically.

Please let me know if this direction sounds reasonable before I open a PR.

[DogWala]

Yes, that's what I'm currently using.

[BenjaminBossan]

Thanks for reporting the issue. I agree that it's very annoying to have an error at the very end of the training run. What is unclear to me is where the memory pressure comes from. During training, we have the following situation:

1. Base model weights loaded on GPU.
2. LoRA weights loaded on GPU.
3. Gradients for LoRA weights loaded on GPU.
4. Hidden states required for backwards path loaded on GPU.
5. Possibly optimizer states loaded on GPU.

When we load the best model, we should only have:

1. Base model weights loaded on GPU.
2. LoRA weights loaded on GPU.

How is it possible that this OOMs when training works fine? Unless I'm missing something, there must be something else going on, like memory not being freed or accidentally using the wrong dtype.

One possible fix would be to make the PEFT load_best_model_at_end path follow a more memory-safe loading pattern similar to the non-PEFT model case: load the best model state on CPU first like:

model.load_adapter(
self.state.best_model_checkpoint,
active_adapter,
torch_device="cpu",
autocast_adapter_dtype=False,
)

If this is about autocast_adapter_dtype=False, that would change the dtype, which could indeed safe memory if the base model is loaded in, say, fp16 or bf16 (PEFT by default loads the adapter weights in float32 but this argument leads to the same dtype being used as in the base model weight). However, if the training was done with float32, we could see different behavior due to the lower precision, so I'm not sure if we want to enable this by default.

If this is about torch_device="cpu", AFAICT, this argument is not supported, so I doubt that it helps.

[DogWala]

Thanks for the detailed reply. I agree that, in principle, it is surprising to see an OOM at best-model reload time if the training loop itself completed successfully. The key point from the repro is that this is not happening in a fully cleaned-up post-training state.

In the failing repro, Trainer._load_best_model() is entered while training-related state is still alive. The instrumentation shows:

• is_in_train = true
• optimizer_is_none = false
• lr_scheduler_is_none = false
• callback_optimizer_is_none = false

So this is not yet a state where optimizer / scheduler / callback-held training state has already been torn down. On top of that, the repro intentionally keeps tail GPU occupancy high to simulate a realistic high-utilization training setup.

Under the same training configuration:

• training completes successfully
• disabling load_best_model_at_end does not fail
• enabling load_best_model_at_end fails during the final adapter reload

So the issue is that turning on load_best_model_at_end introduces an additional late failure mode under memory pressure.

I agree that the most principled solution would be to release all unnecessary training-time allocations before loading the best model. However, from looking at the current Trainer lifecycle, that is not easy to clean optimizer / scheduler / callback-held state before _load_best_model() runs, and changing that lifecycle would be a broader and more invasive fix.

---

The immediate failure mechanism in the repro is that the PEFT reload path first materializes a temporary adapter state_dict on CUDA before copying it into the existing adapter parameters. In the failing run, the recorded events show:

• trainer._load_best_model.enter
• peft.load_adapter.enter
• peft.load_peft_weights.enter with device="cuda"

and the traceback ends in:

safe_load_file(filename, device=device)

So the extra memory pressure is:

• base model weights on GPU
• existing adapter weights on GPU
• remaining training-time state / allocator pressure still alive
• plus a temporary CUDA copy of the adapter checkpoint tensors created by load_peft_weights(..., device="cuda")

That temporary CUDA materialization is what pushes the run over the edge in the repro.

---

Also, to clarify one point that may have caused confusion: autocast_adapter_dtype=False is not the key part of this PR, and I am not changing that behavior in the PR.

I used autocast_adapter_dtype=False only in the repro control path to reduce extra variables when comparing the failing CUDA-side adapter load against a CPU-first load path. The actual change proposed in the PR is about the device used for adapter checkpoint loading, not about changing adapter dtype behavior.

---

On the torch_device point, unless I am missing something, I think there may be a mismatch about which load_adapter() implementation is actually being called here.

In this repro, the model passed to Trainer is a PEFT model created with get_peft_model(...), so the relevant load_adapter() implementation appears to be PEFT's PeftModel.load_adapter, rather than only the Transformers adapter integration layer.

As far as I know, in PEFT, PeftModel.load_adapter(...) does support torch_device:

• peft_model.py

And when torch_device is omitted, the PEFT path appears to fall back to device inference before loading adapter weights:

• save_and_load.py
• other.py

So my understanding is that the PR is not introducing a new unsupported argument, but passing an argument that is already supported by the PEFT load_adapter() implementation that this repro resolves to. But if I am missing something about the exact path taken here, I am happy to be corrected.

[BenjaminBossan]

I agree that the most principled solution would be to release all unnecessary training-time allocations before loading the best model. However, from looking at the current Trainer lifecycle, that is not easy to clean optimizer / scheduler / callback-held state before _load_best_model() runs, and changing that lifecycle would be a broader and more invasive fix.

Honestly, I would argue to do something along these lines. The fact that loading the PEFT weight here causes the OOM is in a sense coincidence, it could also happen in other situations, right? So I think a more general solution would be desired. I'm not too knowledgeable about Trainer but at a quick glance, it looks like we should be able to clean up most things, maybe add a call to empty_cache(), and only then call _load_best_model(). @SunMarc WDYT?

Also, to clarify one point that may have caused confusion: autocast_adapter_dtype=False is not the key part of this PR, and I am not changing that behavior in the PR.

Got it.

On the torch_device point, unless I am missing something, I think there may be a mismatch about which load_adapter() implementation is actually being called here.

Ah yes, I was looking at the transformers integration. In this case, I think it's an okay solution, as the final adapter weights will be loaded to the same device as the base weights, so things like inference should work. Why exactly this reduces memory pressure, I'm not sure.

[DogWala]

In this repro, the model passed to Trainer is a PEFT model created with get_peft_model(...), so the relevant load_adapter() implementation appears to be PEFT's PeftModel.load_adapter, rather than only the Transformers adapter integration layer.

trainer.py#L3440 calls _is_peft_model(model) as a check, so the model on this path should be a PEFT model. In the current version, the torch_device argument here should be fine.

[DogWala]

Do you think this issue should be fixed in the trainer, or should we just leave it as it is for now?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[SunMarc]

We do have a release_memory function in Trainer + trainer.accelerator.clear(). Maybe this is what you are looking for, I don't mind calling it before reloading the model.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.