From e4d021426165d26c8200fd854e8681366c653336 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Fri, 31 Jul 2026 22:46:53 +0000 Subject: [PATCH] fix(hf_ptq): use sequential device_map for DiffusionGemma DiffusionGemma ties weights between its encoder and decoder. Loading it with device_map="auto" (balanced) can place the two sides of a tied pair on different GPUs; the tie cannot then be honored and one side is left on the meta device, so the pre-quantization preview fails with: RuntimeError: Tensor.item() cannot be called on meta tensors Detect DiffusionGemma configs in get_model and select device_map="sequential", which keeps tied modules together. This mirrors the existing per-model handling for bart and t5, where device_map="auto" similarly mis-shards tied encoder/decoder weights. Multi-GPU only; single-GPU runs were unaffected. Previously this required passing --use_seq_device_map manually. Fixes NVBug 6524370 Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/example_utils.py | 27 +++++ tests/examples/hf_ptq/test_example_utils.py | 103 ++++++++++++++++++++ 2 files changed, 130 insertions(+) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d74ffb34efb..d1031dd6084 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -304,6 +304,20 @@ def is_speculative(hf_config): ) +def is_diffusion_gemma(hf_config) -> bool: + """Check if the model architecture is DiffusionGemma. + + Underscores are ignored: the family is spelled ``diffusion_gemma`` in configs + and ``DiffusionGemma`` in class names. The nested ``text_config`` is checked too, + since multi-modal wrappers keep the family name there. + """ + names = [] + for cfg in (hf_config, getattr(hf_config, "text_config", None)): + names.append(getattr(cfg, "model_type", None) or "") + names.extend(getattr(cfg, "architectures", None) or []) + return any("diffusiongemma" in name.lower().replace("_", "") for name in names) + + def get_tokenizer(ckpt_path, trust_remote_code=False, **kwargs) -> PreTrainedTokenizerBase: print(f"Initializing tokenizer from {ckpt_path}") @@ -696,6 +710,19 @@ def get_model( model_kwargs = config_kwargs.copy() model_kwargs.setdefault("dtype", "auto") + # DiffusionGemma ties encoder/decoder weights. device_map "auto" (balanced) can split + # a tied pair across GPUs, leaving one side on the meta device and breaking generation. + # Sequential packs the model onto GPU 0 first (up to gpu_mem_percentage), keeping tied + # modules together for checkpoints that fit; larger ones can still spill and split a + # tied pair, and need an explicit single-device map. Multi-GPU only: a single-GPU split + # cannot separate a tied pair, and sequential would needlessly cap max_memory there. + if device != "cpu" and torch.cuda.device_count() > 1 and is_diffusion_gemma(hf_config): + print( + "Detected DiffusionGemma model. Using device_map='sequential'; the balanced " + "'auto' mapping can split its tied encoder/decoder weights across GPUs." + ) + use_seq_device_map = True + if use_seq_device_map: device_map = "sequential" # If we use sequential, set max_memory limit to ensure that the model does not occupy the full GPU diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 00621ec6125..57b9676ee92 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -316,3 +316,106 @@ def from_pretrained(*args, **kwargs): else: assert "trust_remote_code" not in calls["from_config"] assert calls["from_pretrained"]["trust_remote_code"] is True + + +@pytest.mark.parametrize( + ("model_type", "architecture", "device_count", "expected_device_map"), + [ + # DiffusionGemma ties encoder/decoder weights; "auto" can split a tied pair + # across GPUs, so multi-GPU loads must fall back to "sequential". + ("diffusion_gemma", "DiffusionGemmaForConditionalGeneration", 2, "sequential"), + # Detection must also work off ``architectures`` alone, without ``model_type``. + (None, "DiffusionGemmaForConditionalGeneration", 2, "sequential"), + # Single GPU cannot split a tied pair, so it keeps the unrestricted "auto" map. + ("diffusion_gemma", "DiffusionGemmaForConditionalGeneration", 1, "auto"), + # "gemma" is a substring of "diffusiongemma"; other Gemmas must not match. + ("gemma3", "Gemma3ForCausalLM", 2, "auto"), + ], +) +def test_get_model_device_map_for_diffusion_gemma( + monkeypatch, model_type, architecture, device_count, expected_device_map +): + calls = {} + hf_config = SimpleNamespace( + architectures=[architecture], + dtype=torch.float16, + model_type=model_type, + torch_dtype=torch.bfloat16, + ) + + class FakeModel: + def eval(self): + calls["eval"] = True + + def parameters(self): + return iter(()) + + class FakeArchitecture: + @staticmethod + def _from_config(config, **kwargs): + return FakeModel() + + @staticmethod + def from_pretrained(*args, **kwargs): + calls["from_pretrained"] = kwargs + return FakeModel() + + monkeypatch.setattr( + example_utils.AutoConfig, "from_pretrained", lambda *args, **kwargs: hf_config + ) + # Set rather than delete: ``transformers`` lazy-imports, so a deleted real class + # (e.g. Gemma3ForCausalLM) reappears on the next ``hasattr`` and the real one loads. + # raising=False: DiffusionGemma may not exist in the installed transformers. + monkeypatch.setattr(example_utils.transformers, architecture, FakeArchitecture, raising=False) + monkeypatch.setattr(example_utils, "is_nemotron_vl", lambda config: False) + monkeypatch.setattr(example_utils, "is_speculative", lambda config: False) + monkeypatch.setattr(example_utils, "init_empty_weights", lambda include_buffers: nullcontext()) + monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 1024}) + monkeypatch.setattr(example_utils, "infer_auto_device_map", lambda model, max_memory: {"": 0}) + monkeypatch.setattr(torch.cuda, "device_count", lambda: device_count) + + example_utils.get_model("checkpoint", device="cuda", trust_remote_code=True) + + assert calls["from_pretrained"]["device_map"] == expected_device_map + # Sequential caps per-GPU memory; "auto" must stay unrestricted. + if expected_device_map == "sequential": + assert calls["from_pretrained"]["max_memory"] == {0: 1024 * 0.8} + else: + assert "max_memory" not in calls["from_pretrained"] + + +@pytest.mark.parametrize( + ("hf_config", "expected"), + [ + (SimpleNamespace(model_type="diffusion_gemma", architectures=None), True), + (SimpleNamespace(model_type=None, architectures=["DiffusionGemmaForCausalLM"]), True), + (SimpleNamespace(model_type="gemma3", architectures=["Gemma3ForCausalLM"]), False), + # Multi-modal wrappers keep the family name on the nested ``text_config``. + ( + SimpleNamespace( + model_type="multimodal", + architectures=["SomeWrapperForConditionalGeneration"], + text_config=SimpleNamespace(model_type="diffusion_gemma"), + ), + True, + ), + ( + SimpleNamespace( + model_type="multimodal", + text_config=SimpleNamespace(architectures=["DiffusionGemmaForCausalLM"]), + ), + True, + ), + # A non-DiffusionGemma nested config must not match. + ( + SimpleNamespace( + model_type="multimodal", text_config=SimpleNamespace(model_type="gemma3") + ), + False, + ), + # Stub configs may omit either attribute entirely. + (SimpleNamespace(), False), + ], +) +def test_is_diffusion_gemma(hf_config, expected): + assert example_utils.is_diffusion_gemma(hf_config) is expected