Use device_map="auto" in single file tests to support large models on limited GPU memory#13816
Conversation
Signed-off-by: jiqing-feng <[email protected]>
Signed-off-by: jiqing-feng <[email protected]>
|
Hi @sayakpaul . Would you please review this PR? Thanks! |
|
The failed CI ("LoRA tests with PEFT main") is unrelated to this PR. The same failure is happening across other PRs as well. It appears to be a flaky/known issue with PEFT main. All other 12 checks pass. |
|
@DN6 , could you pls help take a second look? Thx. |
|
Hi @DN6 . The failed CI is not related to my changes. Would you please review the PR? Thanks! |
|
Hi @DN6 . Please let me know what need to be changed before merging this PR. Thanks! |
|
The failed CI is not related to my changes. |
|
Hi @sayakpaul . I already got the approval. Please let me know what need to be changed before merging. Thanks! |
|
Why do you folks keep pinging multiple maintainers multiple times on multiple PRs? |
|
Hi @sayakpaul, apologies for the repeated pings — that wasn't my intention to spam you or the team. I just wanted to keep this small test-only PR from going stale, but I completely understand it adds noise on your end. I'll hold off on further pings and wait patiently. The PR is already approved by DN6 and the only failing CI ("LoRA tests with PEFT main") is an unrelated flaky failure seen across other PRs. Please take your time, and thanks a lot for your work maintaining the project. |
|
Failing tests are unrelated. |
…on limited GPU memory (#13816) * fix flux tests OOM on 24G GPU Signed-off-by: jiqing-feng <[email protected]> * revert wrong change Signed-off-by: jiqing-feng <[email protected]> --------- Signed-off-by: jiqing-feng <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
Problem
Single file loading tests (
SingleFileTesterMixin) useddevice=torch_deviceordevice_map=torch_device, forcing the entire model onto a single GPU. For large models like FLUX.1-dev (~12B params, ~24GB in bf16), this fails on single 24GB GPUs — especiallytest_single_file_model_configwhich loads two models simultaneously.Changes
tests/models/testing_utils/single_file.pytest_single_file_model_config:device=torch_device→device_map="auto"test_single_file_model_parameters:device_map=str(torch_device)/device=torch_device→device_map="auto"test_single_file_loading_with_device_map:device_map=torch_device→device_map="auto"tests/models/transformers/test_models_transformer_flux.pyTestFluxSingleFile: addedtorch_dtype = torch.bfloat16to halve memory usagetests/single_file/test_model_flux_transformer_single_file.pytest_device_map_cuda→test_device_map_auto:device_map="cuda"→device_map="auto", addedtorch_dtype=torch.bfloat16Why
device_map="auto"instead of CPU offloadenable_model_cpu_offload()is a pipeline-level API, not available for individual modelfrom_single_fileloading.device_map="auto"is the model-level solution — accelerate automatically places weights on GPU and offloads the rest to CPU RAM when GPU memory is insufficient.