DiffusionGemma weights remain on 'meta' device with device_map="auto"
Author: HuangBugWei
Pull request: no

Hello Hugging Face team, 
Thank you for your amazing work on the `transformers` library. I encountered an issue when trying to run the new `DiffusionGemma` model using `device_map="auto"`:

```python
from transformers import DiffusionGemmaForBlockDiffusion, AutoProcessor

MODEL_ID = "google/diffusiongemma-26B-A4B-it"

# Load model
processor = AutoProcessor.from_pretrained(MODEL_ID, local_files_only=True)
model = DiffusionGemmaForBlockDiffusion.from_pretrained(
    MODEL_ID,
    dtype="auto",
    device_map="auto",
    local_files_only=True
)


# Prompt
message = [
    {"role": "user", "content": "Why is the sky blue?"}
]

# Process input
input_ids = processor.apply_chat_template(
    message,
    tokenize=True,
    add_generation_prompt=True,
    return_dict=True,
    return_tensors="pt"
).to(model.device)
output = model.generate(**input_ids, max_new_tokens=512)

# Parse output
text = processor.decode(output[0], skip_special_tokens=False)

```

### Environment
- **transformers:** 5.11.0
- **accelerate:** 1.13.0
- **torch:** 2.11.0+cu126
- **flash_attn:** 2.8.3
- **OS:** Linux (Python 3.10.12)


### Error Stack Trace

Running the code above results in the following error during `model.generate()`:

```text
File "transformers/models/diffusion_gemma/generation_diffusion_gemma.py", line 645, in generate
    and input_ids is not None
File "torch/_meta_registrations.py", line 8086, in meta_local_scalar_dense
    raise RuntimeError("Tensor.item() cannot be called on meta tensors")
RuntimeError: Tensor.item() cannot be called on meta tensors

```

I have inspected the device allocation using `model.named_parameters()` right before the generation step:

```python
for name, param in model.named_parameters():
    print(f"{name}, {param.device}")
# model.encoder.language_model.layers.0.self_attn.v_proj.weight, meta
```

I noticed that multiple sub-modules (such as `model.encoder.language_model.embed_tokens.weight` and various layers under `model.decoder.self_conditioning`) still remain on the `meta` device instead of being properly materialized to `cuda`.

It seems like `device_map="auto"` is missing or skipping these specific blocks during the dispatch/load phase for this new architecture.

Any insights or fixes on this would be highly appreciated. Thank you for your time!

Comments:
Hi @HuangBugWei 👋 

I've copied your snippet to my dev setup (H100 GPU), and it ran fine on my end. I have the same package versions as you, so it may be hardware-dependent. 

What's your hardware setup?

Hi @gante, 

I ran it on 2 l40s GPUs.

@HuangBugWei can you also print out and show the output of `model.hf_device_map`?

Hi @zucchini-nlp ,

I tried printing `model.hf_device_map` for `DiffusionGemmaForBlockDiffusion`, and it actually exposed the root cause! It throws an `AttributeError`:

```bash
Traceback (most recent call last):
  File "/share/home/defaultuser/project/diffusiongemma.py", line 20, in <module>
    print(model.hf_device_map)
          ^^^^^^^^^^^^^^^^^^^
  File "/opt/python-3.11.12-gpu/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1968, in __getattr__
    raise AttributeError(
AttributeError: 'DiffusionGemmaForBlockDiffusion' object has no attribute 'hf_device_map'
```

To check my environment, I ran a autoregressive `gemma-4-26B-A4B-it` model under the exact same 2x L40S setup, and it successfully returned the expected device map:

```bash
# Standard gemma output (works perfectly):
{'model.vision_tower': 0, 'model.language_model.embed_tokens': 0, 'lm_head': 0, ... 'model.embed_vision': 1}
```

I'm not entirely sure why this happens, but it seems hf_device_map is not being created for this new architecture under my setup.

Hopefully, this info helps with your investigation. Please let me know if you need me to run any more tests!

@zucchini-nlp This is something we noticed at dev time, but that we hadn't the time (and knowledge) to fix: multi-device setups seem to hang.

I don't have a multi-GPU setup, so not sure how to help 👀  My suspicion: since the model has the same structure as a gemma4moe, and it works fine there, it may be due to the model-specific weight sharing?

Interesting, let me try to reproduce it with multiple small gpus, though inspecting the code I can't pinpoint anything weird

Hey @HuangBugWei, I experienced this exact same symptom (weights stranded on the `meta` device) while deploying this locally, and I managed to track down the root cause.

The issue is a conflict between how `transformers` defines the tied weights for `DiffusionGemma` and how `accelerate` handles them during device mapping.

### The Root Cause
In `transformers/models/diffusion_gemma/modeling_diffusion_gemma.py`, `DiffusionGemmaModel._tied_weights_keys` uses regex patterns (e.g., `r"encoder.language_model.layers\.(?:[^.]+\.)*weight"`). 

Because `accelerate` does not fully support matching regex patterns during its module split hierarchy calculations, it fails to link the tied weights properly. In your case, this caused `device_map="auto"` to silently skip assigning those specific modules to a device, leaving them stranded on `meta` and triggering your `RuntimeError` during `model.generate()`.

If you run this in a highly constrained memory environment (like a single 24GB GPU) where `accelerate` is forced to aggressively split the encoder across devices, this silent failure actually erupts into a hard crash inside `get_module_size_with_ties`.

### Minimal Reproduction of the Core Bug
You can trigger the underlying `IndexError` instantly without downloading weights by forcing `accelerate` to split the model on the `meta` device:

```python
import torch
from transformers import DiffusionGemmaForBlockDiffusion, DiffusionGemmaConfig
from accelerate import infer_auto_device_map

config = DiffusionGemmaConfig.from_pretrained("google/diffusiongemma-26B-A4B-it")

with torch.device("meta"):
    model = DiffusionGemmaForBlockDiffusion(config)

# Force a split to trigger the IndexError
max_memory = {0: "22GiB", "cpu": "50GiB"} 
device_map = infer_auto_device_map(
    model,
    max_memory=max_memory,
    no_split_module_classes=model._no_split_modules,
    dtype=torch.bfloat16
)
```

### The Workaround

Until this is patched, you have to completely bypass `accelerate`'s size calculations by passing a hardcoded `device_map` dictionary.

If you are running this on a shared workstation, this split keeps the heavy embeddings and initial layers hot on the GPU (~11.1 GB VRAM) while offloading the rest to system RAM:

```python
device_map = {}
device_map["model.encoder.language_model.embed_tokens"] = 0
device_map["model.decoder.embed_tokens"] = 0
device_map["model.encoder.vision_tower"] = 0
device_map["model.encoder.embed_vision"] = 0
device_map["model.decoder.self_conditioning"] = 0
device_map["lm_head"] = 0

# Keep layers 0-5 on GPU, offload 6-29 to CPU
for i in range(6):
    device_map[f"model.encoder.language_model.layers.{i}"] = 0
    device_map[f"model.decoder.layers.{i}"] = 0

for i in range(6, 30):
    device_map[f"model.encoder.language_model.layers.{i}"] = "cpu"
    device_map[f"model.decoder.layers.{i}"] = "cpu"

device_map["model.encoder.language_model.norm"] = "cpu"
device_map["model.decoder.norm"] = "cpu"

model = DiffusionGemmaForBlockDiffusion.from_pretrained(
    "google/diffusiongemma-26B-A4B-it",
    torch_dtype=torch.bfloat16,
    device_map=device_map
)
```

Cc @ArthurZucker, @muellerzr — it seems the fix here is either explicitly expanding `_tied_weights_keys` to list all 30 layers by string name in `transformers`, or patching `accelerate` to handle regex hierarchies more gracefully.


Yep, also seeing the issue with multi-gpu when the whole module is tied since it is not being split further. Doesn't seem to be a regex issue, all tied groups are inferred correctly when building the device map

I will raise a PR soon