Skip to content

Fix LoRA offload detection for alignment hooks (#10914)#14271

Open
Bekhouche wants to merge 2 commits into
huggingface:mainfrom
Bekhouche:fix-10914-lora-offload-detection
Open

Fix LoRA offload detection for alignment hooks (#10914)#14271
Bekhouche wants to merge 2 commits into
huggingface:mainfrom
Bekhouche:fix-10914-lora-offload-detection

Conversation

@Bekhouche

Copy link
Copy Markdown

What does this PR do?

Fixes #10914.

LoRA loading previously treated every Accelerate AlignDevicesHook as evidence that sequential CPU offloading was enabled. Alignment-only hooks, such as those installed for quantized models, were therefore removed and followed by an unintended call to enable_sequential_cpu_offload(). This could move a GPU-resident model to CPU and significantly slow inference.

This change:

  • only recognizes AlignDevicesHook instances with offload=True as sequential CPU offloading;
  • inspects nested module hooks, preserving detection of genuine sequential offloading;
  • preserves alignment-only hooks;
  • adds regression coverage for both behaviors.

This updates the approach from the stale attempt in #8750 by checking the hook's explicit offload state instead of inferring it only from the hook type.

Tests

conda run -n conda3.12 env PYTHONPATH=src python -m pytest tests/lora/test_lora_loader_utils.py -q

8 passed

Only treat hooks that actively offload weights as sequential CPU offloading, preserving alignment-only hooks while retaining genuine offload handling.

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Model getting offloaded to CPU without user's intention

1 participant