[Bugfix][Model] Support tensor parallelism for DiffusionGemma (#45719) (#46177)
Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Carlos Alvarado Martinez committed
Jun 26, 2026 at 16:05 UTC
701a23d99f405668158d1395e11c30107dd65b75
4 files changed
+69
-3
.buildkite/test_areas/lm_eval.yaml
+12
@@ -65,6 +65,18 @@ steps:
65
commands:
66
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
67
68
+- label: LM Eval Small Models (2xL4)
69
+ key: lm-eval-small-models-tp
70
+ timeout_in_minutes: 10
71
+ num_devices: 2
72
+ optional: true
73
+ source_file_dependencies:
74
+ - csrc/
75
+ - vllm/model_executor/layers/quantization
76
+ autorun_on_main: true
77
+ commands:
78
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
79
+
80
- label: LM Eval Large Models EP (2xB200)
81
key: lm-eval-large-models-ep-2xb200
82
timeout_in_minutes: 120
tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml
new
+10
@@ -0,0 +1,10 @@
1
+model_name: "RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic"
2
+accuracy_threshold: 0.84
3
+num_questions: 1319
4
+num_fewshot: 5
5
+startup_max_wait_seconds: 1200
6
+server_args: >-
7
+ --enforce-eager
8
+ --max-model-len 4096
9
+ --tensor-parallel-size 2
10
+ --attention-backend TRITON_ATTN
tests/evals/gsm8k/configs/models-small-tp.txt
new
+1
@@ -0,0 +1 @@
1
+DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml
vllm/model_executor/models/diffusion_gemma.py
+46
-3
@@ -28,6 +28,7 @@ from transformers import AutoModel
28
29
from vllm.config import VllmConfig
30
from vllm.config.compilation import CUDAGraphMode
31
+from vllm.distributed.parallel_state import get_tp_group
32
from vllm.logger import init_logger
33
from vllm.model_executor.layers.layernorm import RMSNorm
34
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -499,6 +500,13 @@ def _compiled_sample_step(
500
ST: int,
501
# Sampler config
502
entropy_bound: float,
503
+ # Tensor-parallel vocab sharding for the self-conditioning matmul.
504
+ # ``embed_weight`` is vocab-sharded ([vocab/tp, hidden]) while ``probs``
505
+ # spans the full vocab; [sc_vocab_start, sc_vocab_end) is this rank's slice.
506
+ sc_vocab_start: int,
507
+ sc_vocab_end: int,
508
+ tp_size: int,
509
+ tp_group_name: str,
510
) -> torch.Tensor:
511
"""Compiled decode step: temperature → Gumbel sample → probs/confidence →
512
accept/renoise → convergence, all as vectorized PyTorch ops.
@@ -629,7 +637,17 @@ def _compiled_sample_step(
637
# sc_embeds directly. Storing the [.., hidden] soft embed instead of the full
638
# [.., vocab] probs avoids a giant persistent buffer.
639
sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None]
632
- soft_embeds = torch.matmul(probs.to(embed_weight.dtype), embed_weight) * normalizer
640
+ # Self-conditioning soft embed = probs @ embed_tokens.weight. Under tensor
641
+ # parallelism the embedding is vocab-sharded ([vocab/tp, hidden]) while
642
+ # probs spans the full vocab, so each rank multiplies its local vocab slice
643
+ # [sc_vocab_start, sc_vocab_end) and the partials are summed across ranks.
644
+ local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype)
645
+ soft_embeds = torch.matmul(
646
+ local_probs, embed_weight[: sc_vocab_end - sc_vocab_start]
647
+ )
648
+ if tp_size > 1:
649
+ soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name)
650
+ soft_embeds = soft_embeds * normalizer
651
sc_embeds[decode_slots] = soft_embeds * sc_keep
652
653
# Overwrite canvas with argmax for newly converged denoise requests
@@ -843,6 +861,12 @@ class DiffusionGemmaModelState(ModelState):
861
raise ValueError(
862
f"entropy_bound must be a positive float (got {entropy_bound})"
863
)
864
+ # The self-conditioning matmul (probs @ embed_tokens.weight) runs over a
865
+ # vocab-parallel embedding shard. Hand the sampler this rank's vocab
866
+ # slice and TP group so it can all-reduce the partial products.
867
+ embed_tokens = self.model.model.embed_tokens
868
+ shard = embed_tokens.shard_indices
869
+ tp_group = get_tp_group()
870
return DiffusionSampler(
871
sampler=sampler,
872
diffusion_config=diffusion_config,
@@ -852,8 +876,12 @@ class DiffusionGemmaModelState(ModelState):
876
t_max=gen["t_max"],
877
entropy_bound=entropy_bound,
878
confidence_threshold=gen["confidence_threshold"],
855
- embed_weight=self.model.model.embed_tokens.weight,
879
+ embed_weight=embed_tokens.weight,
880
normalizer=self.model.model.normalizer,
881
+ sc_vocab_start=shard.org_vocab_start_index,
882
+ sc_vocab_end=shard.org_vocab_end_index,
883
+ tp_size=tp_group.world_size,
884
+ tp_group_name=tp_group.unique_name,
885
), None
886
887
def apply_staged_writes(self) -> None:
@@ -1054,13 +1082,24 @@ class DiffusionSampler:
1082
entropy_bound: float,
1083
embed_weight: torch.Tensor,
1084
normalizer: torch.Tensor,
1085
+ sc_vocab_start: int = 0,
1086
+ sc_vocab_end: int | None = None,
1087
+ tp_size: int = 1,
1088
+ tp_group_name: str = "",
1089
):
1090
self.sampling_states = sampler.sampling_states
1091
self.req_states = sampler.req_states
1092
# Self-conditioning soft embed = probs @ embed_weight * normalizer,
1061
- # computed in the sampler (see _compiled_sample_step).
1093
+ # computed in the sampler (see _compiled_sample_step). ``embed_weight``
1094
+ # is the vocab-parallel shard; [sc_vocab_start, sc_vocab_end) is this
1095
+ # rank's slice of the full vocab and tp_* drive the cross-rank
1096
+ # all-reduce.
1097
self.embed_weight = embed_weight
1098
self.normalizer = normalizer
1099
+ self.sc_vocab_start = sc_vocab_start
1100
+ self.sc_vocab_end = sc_vocab_end if sc_vocab_end is not None else vocab_size
1101
+ self.tp_size = tp_size
1102
+ self.tp_group_name = tp_group_name
1103
self.canvas_length = (
1104
diffusion_config.canvas_length if diffusion_config is not None else 32
1105
)
@@ -1299,6 +1338,10 @@ class DiffusionSampler:
1338
CL=self.canvas_length,
1339
ST=states.stability_threshold,
1340
entropy_bound=self.entropy_bound,
1341
+ sc_vocab_start=self.sc_vocab_start,
1342
+ sc_vocab_end=self.sc_vocab_end,
1343
+ tp_size=self.tp_size,
1344
+ tp_group_name=self.tp_group_name,
1345
)
1346
1347
# --- Logprobs: stash on convergence, return on commit ---