Skip to content

fix(training): pin the trainer's model device_map to the resolved device - #33

Merged
log0u7 merged 1 commit into
mainfrom
fix/bnb-cpu-offload
Oct 26, 2026
Merged

log0u7 merged 1 commit into
mainfrom
fix/bnb-cpu-offload

Conversation

@log0u7

@log0u7 log0u7 commented Oct 26, 2026

Copy link
Copy Markdown
Owner

Summary

  • QLoRA/LoRA trainer model load kwargs now pin plan.device (device_map): the default auto dispatch sent quantized modules to the CPU on a busy 10GB card - an 8B QLoRA load failed with 'some modules are dispatched on the CPU or the disk' (found by the first 8B run on the RTX 3080).
  • Extracted _model_init_kwargs(plan) for testability; unit test added.

Issue for this PR

N/A

Checklist

  • Conventional Commits
  • Test added
  • make lint + mypy green locally

The default auto dispatch sent quantized modules to the CPU on a busy
10GB card: an 8B QLoRA load failed with 'some modules are dispatched on
the CPU or the disk'. model_init_kwargs now pins plan.device (extracted
into _model_init_kwargs for testability).
@log0u7
log0u7 merged commit 71a53ad into main Oct 26, 2026
8 checks passed
@log0u7
log0u7 deleted the fix/bnb-cpu-offload branch October 26, 2026 20:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant