xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-03-16 14:27:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Model][Bugfix] fix pipeline parallelism support for NemotronH (#27968)

Browse Source 
Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com>
This commit is contained in:
tomeras91 2025-11-04 14:28:36 +02:00 committed by GitHub
parent 300a265978
commit 77f8001f53
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
vllm/model_executor/models/nemotron_h.py
View File

@ -20,6 +20,7 @@
import typing import typing
from collections.abc import Callable, Iterable from collections.abc import Callable, Iterable
from itertools import islice
import torch import torch
from torch import nn from torch import nn
@ -549,7 +550,7 @@ class NemotronHModel(nn.Module):
self.start_layer, self.end_layer, self.layers = make_layers( self.start_layer, self.end_layer, self.layers = make_layers(
len(config.hybrid_override_pattern), get_layer, prefix=f"{prefix}.layers" len(config.hybrid_override_pattern), get_layer, prefix=f"{prefix}.layers"
) )
self.make_empty_intmd_tensors = make_empty_intermediate_tensors_factory( self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
["hidden_states", "residual"], config.hidden_size ["hidden_states", "residual"], config.hidden_size
) )
@ -564,7 +565,7 @@ class NemotronHModel(nn.Module):
positions: torch.Tensor, positions: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None, intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor: ) -> torch.Tensor | IntermediateTensors:
if get_pp_group().is_first_rank: if get_pp_group().is_first_rank:
if inputs_embeds is not None: if inputs_embeds is not None:
hidden_states = inputs_embeds hidden_states = inputs_embeds
@ -576,8 +577,7 @@ class NemotronHModel(nn.Module):
hidden_states = intermediate_tensors["hidden_states"] hidden_states = intermediate_tensors["hidden_states"]
residual = intermediate_tensors["residual"] residual = intermediate_tensors["residual"]
residual = None for layer in islice(self.layers, self.start_layer, self.end_layer):
for i, layer in enumerate(self.layers):
hidden_states, residual = layer( hidden_states, residual = layer(
positions=positions, positions=positions,
hidden_states=hidden_states, hidden_states=hidden_states,
@ -633,6 +633,9 @@ class NemotronHModel(nn.Module):
if name.endswith(".bias") and name not in params_dict: if name.endswith(".bias") and name not in params_dict:
continue continue
if is_pp_missing_parameter(name, self):
continue
param = params_dict[name] param = params_dict[name]
weight_loader = param.weight_loader weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id) weight_loader(param, loaded_weight, shard_id)
@ -678,6 +681,9 @@ class NemotronHModel(nn.Module):
if is_expert_weight: if is_expert_weight:
continue continue
if is_pp_missing_parameter(name, self):
continue
param = params_dict[name] param = params_dict[name]
weight_loader = getattr( weight_loader = getattr(
param, "weight_loader", default_weight_loader param, "weight_loader", default_weight_loader
@ -792,7 +798,9 @@ class NemotronHForCausalLM(
self.unpadded_vocab_size, config.vocab_size self.unpadded_vocab_size, config.vocab_size
) )
self.make_empty_intmd_tensors = self.model.make_empty_intmd_tensors self.make_empty_intermediate_tensors = (
self.model.make_empty_intermediate_tensors
)
# Set MoE hyperparameters # Set MoE hyperparameters
if self.model.has_moe: if self.model.has_moe: