Skip to content
Merged
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion vllm/model_executor/models/nemotron_h.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,10 +148,12 @@ def __init__(

self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe

router_logits_dtype = torch.float32
self.gate = ReplicatedLinear(
config.hidden_size,
config.n_routed_experts,
bias=False,
params_dtype=router_logits_dtype,
quant_config=None,
prefix=f"{prefix}.gate",
)
Expand Down Expand Up @@ -230,6 +232,7 @@ def __init__(
enable_eplb=self.enable_eplb,
num_redundant_experts=self.n_redundant_experts,
is_sequence_parallel=self.is_sequence_parallel,
router_logits_dtype=router_logits_dtype,
routed_input_transform=self.fc1_latent_proj,
)

Expand All @@ -241,7 +244,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states = sequence_parallel_chunk(hidden_states)

# router_logits: (num_tokens, n_experts)
router_logits, _ = self.gate(hidden_states)
router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))

# SharedFusedMoE handles:
# - shared experts (with original hidden_states)
Expand Down
Loading