A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 20:45 UTC
[Model][Bugfix] fix pipeline parallelism support for NemotronH (#27968)
Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com>
tomeras91 committed
Nov 4, 2025 at 14:28 UTC
77f8001f533021ece46779f5b7e69edc1d3b514f
1 file changed
+13
-5
vllm/model_executor/models/nemotron_h.py
+13
-5
@@ -20,6 +20,7 @@
20
21
import typing
22
from collections.abc import Callable, Iterable
23
+from itertools import islice
24
25
import torch
26
from torch import nn
@@ -549,7 +550,7 @@ class NemotronHModel(nn.Module):
550
self.start_layer, self.end_layer, self.layers = make_layers(
551
len(config.hybrid_override_pattern), get_layer, prefix=f"{prefix}.layers"
552
)
552
- self.make_empty_intmd_tensors = make_empty_intermediate_tensors_factory(
553
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
554
["hidden_states", "residual"], config.hidden_size
555
)
556
@@ -564,7 +565,7 @@ class NemotronHModel(nn.Module):
565
positions: torch.Tensor,
566
intermediate_tensors: IntermediateTensors | None = None,
567
inputs_embeds: torch.Tensor | None = None,
567
- ) -> torch.Tensor:
568
+ ) -> torch.Tensor | IntermediateTensors:
569
if get_pp_group().is_first_rank:
570
if inputs_embeds is not None:
571
hidden_states = inputs_embeds
@@ -576,8 +577,7 @@ class NemotronHModel(nn.Module):
577
hidden_states = intermediate_tensors["hidden_states"]
578
residual = intermediate_tensors["residual"]
579
579
- residual = None
580
- for i, layer in enumerate(self.layers):
580
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
581
hidden_states, residual = layer(
582
positions=positions,
583
hidden_states=hidden_states,
@@ -633,6 +633,9 @@ class NemotronHModel(nn.Module):
633
if name.endswith(".bias") and name not in params_dict:
634
continue
635
636
+ if is_pp_missing_parameter(name, self):
637
+ continue
638
+
639
param = params_dict[name]
640
weight_loader = param.weight_loader
641
weight_loader(param, loaded_weight, shard_id)
@@ -678,6 +681,9 @@ class NemotronHModel(nn.Module):
681
if is_expert_weight:
682
continue
683
684
+ if is_pp_missing_parameter(name, self):
685
+ continue
686
+
687
param = params_dict[name]
688
weight_loader = getattr(
689
param, "weight_loader", default_weight_loader
@@ -792,7 +798,9 @@ class NemotronHForCausalLM(
798
self.unpadded_vocab_size, config.vocab_size
799
)
800
795
- self.make_empty_intmd_tensors = self.model.make_empty_intmd_tensors
801
+ self.make_empty_intermediate_tensors = (
802
+ self.model.make_empty_intermediate_tensors
803
+ )
804
805
# Set MoE hyperparameters
806
if self.model.has_moe: