From 7488b5780f585fc0a00a6ab28b6a2477696574d1 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 09:42:47 -0500 Subject: [PATCH] [ROCm][Tests] Reduce host memory when downcasting FP32 HF references (#57869) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/conftest.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index 508e9fa3d5..cd1d11917b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -504,6 +504,36 @@ class HfRunner: if model_cls is not None: _fix_v4_tied_weights_keys(model_cls) + from transformers.integrations import is_deepspeed_zero3_enabled + + # On ROCm, avoid a converted CPU copy when downcasting FP32 references. + # Same-dtype CPU loads can stay lazily mapped and use less host RAM. + # Leave customized, remote, and quantized loading unchanged. + if ( + current_platform.is_rocm() + and self.device == "cuda" + and auto_cls is AutoModelForCausalLM + and revision is None + and model_kwargs.keys() == {"dtype"} + and self.config.dtype == torch.float32 + and model_kwargs["dtype"] in (torch.float16, torch.bfloat16) + and not (trust_remote_code and hasattr(self.config, "auto_map")) + and getattr(self.config, "quantization_config", None) is None + and getattr( + self.config.get_text_config(decoder=True), + "quantization_config", + None, + ) + is None + and not is_deepspeed_zero3_enabled() + ): + model_kwargs = { + **model_kwargs, + "device_map": torch.device( + "cuda", torch.accelerator.current_device_index() + ), + } + model = cast( nn.Module, auto_cls.from_pretrained(