From 7a43f6246dccf0d10c4353940d4904ab11d406b3 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Mon, 20 Jul 2026 01:46:59 -0700 Subject: [PATCH 1/2] Reduce FLUX int8 test peak memory with sequential offload --- tests/quantization/bnb/test_mixed_int8.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/quantization/bnb/test_mixed_int8.py b/tests/quantization/bnb/test_mixed_int8.py index 981881c005b7..d264dd572f2b 100644 --- a/tests/quantization/bnb/test_mixed_int8.py +++ b/tests/quantization/bnb/test_mixed_int8.py @@ -605,7 +605,14 @@ def _setup_slow_flux(self): transformer=transformer_8bit, torch_dtype=torch.float16, ) - self.pipeline_8bit.enable_model_cpu_offload() + # On devices with <= 24 GB VRAM, enable_model_cpu_offload can OOM because it + # moves an entire sub-model to the accelerator at once. Fall back to + # sequential (per-layer) CPU offload in that case. + _, total_mem = torch.accelerator.get_memory_info() + if total_mem <= 24 * (1024**3): + self.pipeline_8bit.enable_sequential_cpu_offload() + else: + self.pipeline_8bit.enable_model_cpu_offload() yield del self.pipeline_8bit @@ -650,7 +657,7 @@ def test_lora_loading(self): expected_slice = np.array([0.3916, 0.3916, 0.3887, 0.4243, 0.4155, 0.4233, 0.4570, 0.4531, 0.4248]) max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3 + assert max_diff < 2e-3 @require_transformers_version_greater("4.44.0") From 43ebe9972ef276b9b27fb5a4cc1f58ea51c2f53b Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Mon, 20 Jul 2026 19:01:27 -0700 Subject: [PATCH 2/2] Gate FLUX int8 slow test with require_big_accelerator --- tests/quantization/bnb/test_mixed_int8.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/tests/quantization/bnb/test_mixed_int8.py b/tests/quantization/bnb/test_mixed_int8.py index d264dd572f2b..a67d26d8cad4 100644 --- a/tests/quantization/bnb/test_mixed_int8.py +++ b/tests/quantization/bnb/test_mixed_int8.py @@ -41,6 +41,7 @@ load_pt, numpy_cosine_similarity_distance, require_accelerate, + require_big_accelerator, require_bitsandbytes_version_greater, require_peft_backend, require_peft_version_greater, @@ -590,6 +591,7 @@ def get_dummy_tensor_inputs(device=None, seed: int = 0): @require_transformers_version_greater("4.44.0") +@require_big_accelerator class TestSlowBnb8bitFlux(Base8bitTests): @pytest.fixture(autouse=True) def _setup_slow_flux(self): @@ -605,14 +607,7 @@ def _setup_slow_flux(self): transformer=transformer_8bit, torch_dtype=torch.float16, ) - # On devices with <= 24 GB VRAM, enable_model_cpu_offload can OOM because it - # moves an entire sub-model to the accelerator at once. Fall back to - # sequential (per-layer) CPU offload in that case. - _, total_mem = torch.accelerator.get_memory_info() - if total_mem <= 24 * (1024**3): - self.pipeline_8bit.enable_sequential_cpu_offload() - else: - self.pipeline_8bit.enable_model_cpu_offload() + self.pipeline_8bit.enable_model_cpu_offload() yield del self.pipeline_8bit @@ -657,7 +652,7 @@ def test_lora_loading(self): expected_slice = np.array([0.3916, 0.3916, 0.3887, 0.4243, 0.4155, 0.4233, 0.4570, 0.4531, 0.4248]) max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 2e-3 + assert max_diff < 1e-3 @require_transformers_version_greater("4.44.0")