[ROCm][Tests] Reduce host memory when downcasting FP32 HF references (#57869)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Sep 21, 2026 at 09:42 UTC
7488b5780f585fc0a00a6ab28b6a2477696574d1
1 file changed
+30
tests/conftest.py
+30
@@ -504,6 +504,36 @@ class HfRunner:
504
if model_cls is not None:
505
_fix_v4_tied_weights_keys(model_cls)
506
507
+ from transformers.integrations import is_deepspeed_zero3_enabled
508
+
509
+ # On ROCm, avoid a converted CPU copy when downcasting FP32 references.
510
+ # Same-dtype CPU loads can stay lazily mapped and use less host RAM.
511
+ # Leave customized, remote, and quantized loading unchanged.
512
+ if (
513
+ current_platform.is_rocm()
514
+ and self.device == "cuda"
515
+ and auto_cls is AutoModelForCausalLM
516
+ and revision is None
517
+ and model_kwargs.keys() == {"dtype"}
518
+ and self.config.dtype == torch.float32
519
+ and model_kwargs["dtype"] in (torch.float16, torch.bfloat16)
520
+ and not (trust_remote_code and hasattr(self.config, "auto_map"))
521
+ and getattr(self.config, "quantization_config", None) is None
522
+ and getattr(
523
+ self.config.get_text_config(decoder=True),
524
+ "quantization_config",
525
+ None,
526
+ )
527
+ is None
528
+ and not is_deepspeed_zero3_enabled()
529
+ ):
530
+ model_kwargs = {
531
+ **model_kwargs,
532
+ "device_map": torch.device(
533
+ "cuda", torch.accelerator.current_device_index()
534
+ ),
535
+ }
536
+
537
model = cast(
538
nn.Module,
539
auto_cls.from_pretrained(