[NVBug: 6524370] use sequential device_map for DiffusionGemma (#2041)

### What does this PR do?

Type of change: Bug fix

Fixes NVBug 6524370.

`DiffusionGemma` ties weights between its encoder and decoder.
`get_model` loads with `device_map="auto"`
(`examples/hf_ptq/example_utils.py`), and `"auto"` is an alias for
`"balanced"` — accelerate splits the model evenly across all visible
GPUs by size, with no awareness of tied parameters. On multi-GPU it can
place the two sides of a tied pair on different devices; the tie then
cannot be honored and one side is left on the `meta` device.

The pre-quantization preview in `pre_quantize` then reaches `(input_ids
== self.config.image_token_id).any()` in `generation_diffusion_gemma.py`
and fails:

```
RuntimeError: Tensor.item() cannot be called on meta tensors
```

This is multi-GPU-only by construction: with one visible GPU the
balanced split is trivial, nothing is separated, and nothing lands on
`meta`.

This PR detects DiffusionGemma configs in `get_model` and selects
`device_map="sequential"`, which fills one GPU before spilling to the
next and so keeps tied modules together. It mirrors the existing
per-model handling for `bart` and `t5`, where `device_map="auto"`
similarly mis-shards tied encoder/decoder weights.

Detection reads `model_type` and `architectures` from the config and
ignores underscores, since the family is spelled `diffusion_gemma` in
the Transformers module path and `DiffusionGemma` in the class name.

### Usage

No API change. Previously this needed the flag passed manually:

```bash
python hf_ptq.py --model <diffusion-gemma-ckpt> --recipe <recipe> \
  --export_path <out> --trust_remote_code --use_seq_device_map
```

It is now selected automatically, and the model load logs:

```
Detected DiffusionGemma model. Using device_map='sequential'; the balanced
'auto' mapping can split its tied encoder/decoder weights across GPUs.
```

Passing `--use_seq_device_map` explicitly still works and is unaffected.

### Testing

- Reproduced on 4x GB200 with `diffusiongemma-26B-A4B-it` and the
`nvfp4_experts_only` recipe; `--use_seq_device_map` resolves the crash,
confirming the device-mapping cause.
- Validated on oci-hsg (4x GB200): with this patch and no CLI flag,
`diffusiongemma-26B-A4B-it` loads correctly and the meta-tensor crash no
longer reproduces.
- `is_diffusion_gemma` checked against both config spellings,
`architectures=None`, `architectures=[]`, and a `gemma3` negative to
confirm no over-match — `get_model_type` already orders `DiffusionGemma`
before `Gemma` for exactly this substring-collision reason.
- `pre-commit run --files examples/hf_ptq/example_utils.py` passes
(ruff, ruff-format, mypy, bandit).

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ❌ — no existing unit coverage
for `get_model` device-map selection; happy to add a config-level test
for `is_diffusion_gemma` if wanted.
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A
- Did you get Claude approval on this PR?: ❌ — pending

### Additional Information

NVBug 6524370. Same class of failure as the existing `t5` workaround in
`get_model`; a general "any tied encoder/decoder model" rule was
considered but rejected, since `tie_word_embeddings=True` holds for most
decoder-only LLMs where `auto` is fine and forcing sequential would
regress large-model runs.

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved DiffusionGemma model loading on multi-GPU systems by keeping
related model weights together.
* Added more reliable DiffusionGemma model recognition across supported
configurations.
* Preserved existing automatic device allocation for single-GPU systems
and other supported models.
* Improved loading reliability by applying appropriate memory limits
during multi-GPU setup.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Juhi Mittal <juhim@nvidia.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: juhi10071998

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

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