[NVBug: 6524370] use sequential device_map for DiffusionGemma

[juhi10071998]

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:

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.
• 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).

Draft pending an end-to-end PTQ run on 4x GB200 with the patch applied and no CLI flag.

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?: 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

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b7c072fa-e147-4372-9966-fd1b2c0b853a

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 69.96%. Comparing base (a23390d) to head (7929ddd).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2041      +/-   ##
==========================================
- Coverage   69.97%   69.96%   -0.02%     
==========================================
  Files         519      519              
  Lines       59399    59399              
==========================================
- Hits        41565    41559       -6     
- Misses      17834    17840       +6     

Flag	Coverage Δ
unit	55.15% <ø> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.