Skip to content

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

Merged
juhi10071998 merged 1 commit into
mainfrom
nvbug_6524370
Aug 3, 2026
Merged

[NVBug: 6524370] use sequential device_map for DiffusionGemma#2041
juhi10071998 merged 1 commit into
mainfrom
nvbug_6524370

Conversation

@juhi10071998

@juhi10071998 juhi10071998 commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

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

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.

@copy-pr-bot

copy-pr-bot Bot commented Jul 31, 2026

Copy link
Copy Markdown

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

coderabbitai Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

The change detects DiffusionGemma models from Hugging Face configuration metadata. Multi-GPU, non-CPU DiffusionGemma loads now use sequential device mapping. Tests cover detection and device-map selection.

Changes

DiffusionGemma device mapping

Layer / File(s) Summary
DiffusionGemma detection, loading, and validation
examples/hf_ptq/example_utils.py, tests/examples/hf_ptq/test_example_utils.py
Adds is_diffusion_gemma. Multi-GPU, non-CPU DiffusionGemma loads use sequential mapping with capped memory. Tests cover DiffusionGemma detection and mapping behavior.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 6
✅ Passed checks (6 passed)
Check name Status Explanation
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns ✅ Passed PR changes add only DiffusionGemma detection/device mapping and tests; no new unsafe loads, pickle, eval/exec, # nosec, or dependency manifests. trust_remote_code remains caller-controlled.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes the main change: using sequential device mapping for DiffusionGemma models.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch nvbug_6524370

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

@juhi10071998 juhi10071998 changed the title fix(hf_ptq): use sequential device_map for DiffusionGemma [NVBug: 6524370] use sequential device_map for DiffusionGemma Jul 31, 2026
@codecov

codecov Bot commented Jul 31, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 66.83%. Comparing base (14b20c0) to head (e4d0214).

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2041      +/-   ##
==========================================
- Coverage   67.00%   66.83%   -0.18%     
==========================================
  Files         520      520              
  Lines       59545    59545              
==========================================
- Hits        39900    39796     -104     
- Misses      19645    19749     +104     
Flag Coverage Δ
examples 43.08% <ø> (-0.22%) ⬇️
unit 55.25% <ø> (ø)

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.

@juhi10071998
juhi10071998 marked this pull request as ready for review August 3, 2026 20:38
@juhi10071998
juhi10071998 requested a review from a team as a code owner August 3, 2026 20:38
@juhi10071998
juhi10071998 requested a review from cjluo-nv August 3, 2026 20:38

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Warning

CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.

Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.

👉 Steps to fix this

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/hf_ptq/example_utils.py`:
- Around line 715-720: Update the conditional around
is_diffusion_gemma(hf_config) so automatic use of use_seq_device_map is limited
to multi-GPU runs, preserving existing single-GPU memory behavior;
alternatively, retain the mapping and add a regression test covering a
near-capacity one-GPU run and its max_memory handling.
- Around line 710-721: Update the pack-quantized model loading path guarded by
has_pack_quantized_config(hf_config) to use the selected device map, including
use_seq_device_map for DiffusionGemma, instead of hardcoding device_map="auto".
Pass through the computed max_memory value as well, preserving the sequential
mapping and GPU memory limits.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 0d1bb738-0b8a-48a3-b0a3-44a783508295

📥 Commits

Reviewing files that changed from the base of the PR and between 3882635 and 3a02bfa.

📒 Files selected for processing (1)
  • examples/hf_ptq/example_utils.py

Comment thread examples/hf_ptq/example_utils.py Outdated
Comment thread examples/hf_ptq/example_utils.py Outdated

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Small, well-motivated fix: auto-select device_map="sequential" for DiffusionGemma in examples/hf_ptq/example_utils.get_model, mirroring the existing per-model bart/t5 handling. Placement (before the use_seq_device_map block, so max_memory capping still applies) is correct, and the multi-GPU + non-CPU guard is right. Three non-blocking points:

  1. No test, although the harness already exists. The PR body says "no existing unit coverage for get_model device-map selection", but tests/examples/hf_ptq/test_example_utils.py::test_get_model_uses_expected_dtype_kwarg already monkeypatches AutoConfig.from_pretrained, init_empty_weights, get_max_memory, and infer_auto_device_map and asserts on the from_pretrained kwargs. A case with a diffusion_gemma config + patched torch.cuda.device_count() asserting device_map == "sequential" (plus a gemma3 negative for the substring-collision concern the PR body describes as manually verified) is a handful of lines there. Worth adding since the whole fix is a detection heuristic.

  2. The comment overstates the guarantee. sequential only keeps a tied pair together when the model fits on the first GPU; because use_seq_device_map also caps max_memory to gpu_mem_percentage (0.8), a model that doesn't fit under that cap will still spill and can still split the tied pair onto meta. The t5/bart precedent it cites actually uses device_map=None (single device), which is the stronger guarantee. Either soften the comment or explain why sequential is sufficient for this family.

  3. Minor robustness/dedup: hf_config.architectures is accessed directly while the rest of this file uses getattr(config, "architectures", []); and DiffusionGemma detection now lives in three places (MODEL_NAME_TO_TYPE, _reorder_canonical_first, and this new helper).

No licensing changes, no prompt-injection content in the PR text.

Comment thread examples/hf_ptq/example_utils.py Outdated
Underscores are ignored: the family is spelled ``diffusion_gemma`` in configs
and ``DiffusionGemma`` in class names.
"""
names = [getattr(hf_config, "model_type", None) or "", *(hf_config.architectures or [])]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

Two small things:

  • hf_config.architectures is accessed directly here, while the rest of this file consistently uses getattr(config, "architectures", []) (see _is_multimodal_config, is_nemotron_vl). Configs built in tests / non-PretrainedConfig stubs won't have the attribute and will raise AttributeError. Suggest *(getattr(hf_config, "architectures", None) or []).
  • DiffusionGemma detection now exists in three places: MODEL_NAME_TO_TYPE (modelopt/torch/export/model_utils.py), _reorder_canonical_first in the same file (which does the same "diffusiongemma" / "diffusion_gemma" normalization on the class name), and this helper. Consider putting the config-level predicate next to the others in model_utils.py so future spellings only need updating once.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in b45f9c36 — switched to *(getattr(hf_config, "architectures", None) or []).

Worth noting the file is split rather than consistent here: _is_multimodal_config (179) and is_nemotron_vl (204) use getattr, but is_speculative (302) and get_model (792-794) access .architectures directly. Taking the defensive form anyway.

On consolidating into model_utils.py: agreed there is drift risk, though the two existing spots key off the class name (type(model).__name__) while this one keys off the config, so it is not a pure duplicate. Leaving it here to keep this an examples-only bugfix — happy to do the move as a follow-up if you would rather have it now.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirming the owner call: consolidating the DiffusionGemma predicate into model_utils.py will be a follow-up PR, not this one. Keeping this change scoped to the examples-level bugfix so it can land for NVBug 6524370.

The drift risk is real and noted — MODEL_NAME_TO_TYPE, _reorder_canonical_first, and is_diffusion_gemma would all need updating if a new spelling appears.

Comment thread examples/hf_ptq/example_utils.py Outdated

# 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 mapping keeps tied modules together. Multi-GPU only: a single-GPU split

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

"Sequential mapping keeps tied modules together" holds only when the model fits on the first GPU. Because setting use_seq_device_map = True also caps max_memory to gpu_mem_percentage (0.8 by default), a checkpoint larger than 0.8 x one GPU will still spill to the next device and can still land one side of a tied pair on meta — i.e. the same crash, just harder to hit. The t5/bart precedent referenced in the PR body actually sets device_map = None (whole model on one device), which is the stronger guarantee. Please either soften this comment to state the limitation, or explain why sequential is always sufficient for this family (e.g. size bound of the released checkpoints).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair — the comment claimed more than the code delivers. Softened in b45f9c36:

# 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.

You are right that this is a mitigation, not a guarantee, and that t5/bart use the stronger device_map = None. I did not follow them there because None puts the whole model on one device with no headroom cap, which would regress large checkpoints that legitimately need the split. For the reported case (26B-A4B, ~52 GB vs 0.8 x 186 GB on GB200) sequential is comfortably sufficient; beyond that bound the limitation is now stated rather than implied.

Comment thread examples/hf_ptq/example_utils.py
@juhi10071998
juhi10071998 requested a review from a team as a code owner August 3, 2026 21:21

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Warning

CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.

Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.

👉 Steps to fix this

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/examples/hf_ptq/test_example_utils.py`:
- Around line 321-342: Add a parameterized multi-GPU case to
test_get_model_device_map_for_diffusion_gemma with model_type=None,
architecture="DiffusionGemmaForConditionalGeneration", and
expected_device_map="sequential", so the device-map selection is verified when
DiffusionGemma is identified solely by architectures.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 21fab311-75ef-4a84-84d7-1b929860eeeb

📥 Commits

Reviewing files that changed from the base of the PR and between fafe1db and b45f9c3.

📒 Files selected for processing (2)
  • examples/hf_ptq/example_utils.py
  • tests/examples/hf_ptq/test_example_utils.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • examples/hf_ptq/example_utils.py

Comment thread tests/examples/hf_ptq/test_example_utils.py

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Re-review of #2041 (auto device_map="sequential" for DiffusionGemma in examples/hf_ptq/example_utils.get_model). The head commit resolves the substantive items from the previous round; one flagged-Major item is labeled "addressed" but is not present in the code, so a human should confirm scope.

Addressed since last round (verified against the file at head):

  • Missing test → tests/examples/hf_ptq/test_example_utils.py now has test_get_model_device_map_for_diffusion_gemma (multi-GPU → sequential + max_memory == {0: 819.2}; single-GPU → auto with no max_memory; gemma3 negative → auto) and test_is_diffusion_gemma (both spellings, missing attrs). I traced get_model with these stubs — the patched AutoConfig/init_empty_weights/get_max_memory/infer_auto_device_map/torch.cuda.device_count path reaches from_pretrained and is_model_on_gpu cleanly (empty parameters()), so the tests exercise the real selection logic rather than a stub.
  • Overstated comment → softened to state the "checkpoints that fit" limitation and the need for an explicit single-device map beyond it.
  • hf_config.architectures → now *(getattr(hf_config, "architectures", None) or []).
  • Single-GPU regression → gated on torch.cuda.device_count() > 1, so single-GPU keeps unrestricted auto (pinned by a test case).

Remaining for owner judgment:

  • 💬 CodeRabbit's Major finding on the has_pack_quantized_config(hf_config) branch is annotated "✅ Addressed in commit 6ef39d9", but the code at head still hardcodes device_map="auto" (and drops max_memory) in that branch. Published DiffusionGemma compressed-tensors checkpoints exist, and such a load would bypass this fix and can still split the tied pair onto meta — i.e. the same crash on a different entry path. It's pre-existing code and arguably out of scope for this bugfix, but the resolution label doesn't match the diff, so please confirm intent.
  • 💬 Author replied on consolidation: the two existing spots key off type(model).__name__ while this keys off the config, so it's not a pure duplicate; happy to move to model_utils.py as a follow-up — still worth an owner call, since DiffusionGemma spelling normalization now lives in three places (MODEL_NAME_TO_TYPE, _reorder_canonical_first, is_diffusion_gemma) and future spellings must be updated in all of them.
  • Minor, optional: CodeRabbit's last suggestion (a multi-GPU model_type=None + architectures=["DiffusionGemma..."] device-map case) is not added; test_is_diffusion_gemma covers architecture-only detection at the predicate level, which I consider adequate.

No licensing changes and no prompt-injection content in the PR text or comments.

@juhi10071998

Copy link
Copy Markdown
Contributor Author

/claude review

Comment on lines +712 to +723
# 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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] No escape hatch: this override is unconditional, so a multi-GPU DiffusionGemma user can no longer get device_map="auto" at all.

The forced switch changes two things for every multi-GPU DiffusionGemma load — device_map autosequential, and max_memory unset → capped at gpu_mem_percentage (0.8) per GPU. Packing onto GPU 0 first means activations and calibration buffers now compete for the same device rather than spreading across all GPUs, so a checkpoint that happened to load fine under auto (small enough that the balanced split didn't separate the tied pair) can now OOM during calibration with no way to opt back out. That's a real, if narrow, behavior change with no CLI path around it.

The print makes it visible rather than silent, and for the reported 26B-A4B-on-GB200 case the trade is clearly right — so this is non-blocking. But the surrounding per-model handling in this function (bart, t5, mxfp4) is all for cases that are broken under the alternative, whereas here auto is merely risky. Worth considering a --no_seq_device_map / tri-state flag so the override is a default rather than a hard rule, or at minimum noting in examples/llm_ptq/README.md that DiffusionGemma multi-GPU is force-sequential and how to work around an OOM (fewer visible GPUs + explicit single-device map).

Happy to defer this to a follow-up if you'd rather keep the bugfix minimal.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed on the analysis, deferring to a follow-up. The distinction you draw is the right one: bart/t5/mxfp4 special-case loads that are broken under the alternative, whereas auto here is risky rather than always fatal, so a hard override is a stronger stance than the surrounding code takes.

Keeping this PR scoped to the NVBug 6524370 fix. A tri-state flag is new CLI surface and deserves its own review; the README note belongs with it so the documented workaround matches whatever the flag ends up being. Tracking alongside the model_utils.py consolidation follow-up.

Noting the escape hatch that exists today: CUDA_VISIBLE_DEVICES=0 bypasses the override entirely, since it is gated on torch.cuda.device_count() > 1.

Comment thread examples/hf_ptq/example_utils.py

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Claude review passed — no blocking issues found. LGTM

Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3

Scope

Full review of both changed files (examples/hf_ptq/example_utils.py, tests/examples/hf_ptq/test_example_utils.py), plus surrounding get_model context and the DiffusionGemma detection sites in modelopt/torch/export/model_utils.py.

What I verified

  • Placement is correct. The new block sits after model_kwargs is built and before the if use_seq_device_map: block, so setting use_seq_device_map = True still picks up the max_memory cap. The later bart/t5 device_map = None overrides can't collide with DiffusionGemma.
  • Guard is right. device != "cpu" and torch.cuda.device_count() > 1 matches the failure mode described in the PR body — a single-GPU balanced split can't separate a tied pair, and gating on multi-GPU avoids regressing single-GPU runs with an unnecessary 0.8 memory cap.
  • No over-match. "gemma" in "diffusiongemma" is the real collision risk here (the same one MODEL_NAME_TO_TYPE orders around at model_utils.py:38-40), and the substring direction is checked correctly: "diffusiongemma" in name, not the reverse. gemma3/Gemma3ForCausalLM does not match, and that's pinned by a test.
  • Defensive attribute access. getattr(hf_config, "model_type", None) or "" and *(getattr(hf_config, "architectures", None) or []) handle None, [], and missing-attribute configs — all four covered by test_is_diffusion_gemma.
  • Tests exercise the real path. I traced get_model under the stubs: the patched AutoConfig / init_empty_weights / get_max_memory / infer_auto_device_map / torch.cuda.device_count chain reaches from_pretrained and is_model_on_gpu cleanly (parameters() returns an empty iterator), and infer_auto_device_map returning {"": 0} keeps on_cpu False so the auto cases legitimately assert max_memory absent. The assertions land on real selection logic, not on a stub.
  • Backward compatible. No signature or default change; --use_seq_device_map still works and is only ever turned on, never off.

Suggestions posted (all non-blocking)

  1. No escape hatch (example_utils.py:712-723) — the override is unconditional, so a multi-GPU DiffusionGemma load can no longer get auto at all. Packing onto GPU 0 plus the 0.8 cap means a checkpoint that loaded fine under auto could now OOM during calibration with no CLI way back. The print makes it visible and for the reported 26B case the trade is clearly right; worth a follow-up flag or a README note.
  2. is_diffusion_gemma skips nested text_config (example_utils.py:307-317) — its two neighbours in this file (get_original_hf_quant_method, has_pack_quantized_config) both walk (config, config.text_config). A conditional-generation wrapper config that only names the family on the inner config would slip past and silently keep auto. Speculative for published checkpoints, since top-level architectures covers them.
  3. Stale test module docstring (test_example_utils.py:15-19, outside the diff so noted here) — still claims the file covers only load_mtp_weights, though it already covered get_original_hf_quant_method, _resolve_init_config, and get_model dtype selection, and now adds device-map selection. One-line fix while you're in the file.

Note on the earlier has_pack_quantized_config item

The previous round flagged that branch hardcoding device_map="auto"; it still does at head. I'm not counting it as blocking: it's pre-existing code that this diff doesn't touch, and the existing --use_seq_device_map flag is equally inert there — so it's a standing limitation of that branch rather than a gap introduced by this fix. Reasonable to leave for a follow-up.

Risk

Low. Examples-only change, additive, single narrowly-gated branch with an explanatory print, and the detection heuristic is now covered by eight parameterized cases.

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Re-review of #2041 (auto device_map="sequential" for DiffusionGemma in examples/hf_ptq/example_utils.get_model). Everything substantive from the prior rounds is now resolved in-tree, and the two items left open are explicitly deferred with reasoning the author has stated on-thread.

Verified at head:

  • Nested text_config (previous SUGGESTION)is_diffusion_gemma now iterates (hf_config, getattr(hf_config, "text_config", None)), matching get_original_hf_quant_method (625) and has_pack_quantized_config (746). Three new test_is_diffusion_gemma cases cover nested model_type, nested architectures, and a nested gemma3 negative. getattr(None, ...) on a missing text_config degrades to ""/[] correctly, and the bare SimpleNamespace() case pins that.
  • Missing test (critical, prior round)test_get_model_device_map_for_diffusion_gemma now has four parametrized cases: diffusion_gemma @2 GPUs → sequential + max_memory == {0: 819.2}; model_type=None + architecture-only @2 GPUs → sequential (CodeRabbit's last request, addressed in 6acad706); @1 GPU → auto with no max_memory; gemma3 @2 GPUs → auto. I traced get_model under the stubs — AutoConfig/init_empty_weights/get_max_memory/infer_auto_device_map/torch.cuda.device_count all patch the real path, infer_auto_device_map → {"": 0} keeps on_cpu False so the auto cases legitimately assert max_memory absent, and the empty parameters() keeps is_model_on_gpu from blowing up. Real selection logic, not a stub.
  • Single-GPU regression + overstated comment → still fixed (torch.cuda.device_count() > 1 gate; comment states the "checkpoints that fit" limitation).
  • Placement → the block sits before if use_seq_device_map: so the max_memory cap applies, and after it the bart/t5 device_map = None overrides can't collide.

Deferred with stated rationale (accepting as reasonable, not blocking):

  • has_pack_quantized_config still hardcodes device_map="auto" and drops max_memory, so a pack-quantized DiffusionGemma checkpoint can still hit the meta-tensor crash. Author confirmed this is pre-existing (that branch already ignores --use_seq_device_map for every pack-quantized model) and that fixing it changes loading behavior for all such models, so it belongs in its own PR. The author also correctly flagged that CodeRabbit's "✅ Addressed in commit 6ef39d9" badge on that thread is wrong — good catch, and it removes the label/diff mismatch I raised last round.
  • Consolidating the DiffusionGemma predicate into model_utils.py and a tri-state escape-hatch flag (+ README note) are both tracked as follow-ups; CUDA_VISIBLE_DEVICES=0 is a workable stopgap given the multi-GPU gate.

Nit, non-blocking: the test module docstring (lines 15-19) still claims the file covers only load_mtp_weights, though it now also covers get_original_hf_quant_method, _resolve_init_config, and two get_model selection paths — one-line fix if you touch the file again.

No licensing changes (existing headers only). No prompt-injection content in the PR text or comments.

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) <noreply@anthropic.com>
Signed-off-by: Juhi Mittal <juhim@nvidia.com>
@juhi10071998 juhi10071998 self-assigned this Aug 3, 2026
@juhi10071998
juhi10071998 merged commit 302a7ad into main Aug 3, 2026
41 checks passed
@juhi10071998
juhi10071998 deleted the nvbug_6524370 branch August 3, 2026 23:34
@github-actions

github-actions Bot commented Aug 3, 2026

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-08-03 23:34 UTC

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants