Method pp_layers

tensorrt_llm/mapping.py:315–329

(self, num_layers: int)

        return self.moe_ep_size > 1

    def pp_layers(self, num_layers: int) -> List[int]:
        if self.pp_partition is not None:
            if len(self.pp_partition) != self.pp_size:
                raise ValueError(
                    f"{len(self.pp_partition)=} does not match {self.pp_size=}."
                )
            if sum(self.pp_partition) != num_layers:
                raise ValueError(
                    f"{sum(self.pp_partition)=} does not match {num_layers=}.")
            return torch.arange(num_layers).split(
                self.pp_partition)[self.pp_rank].tolist()
        else:
            # If num_layers % pp_size = n != 0, first n ranks get one extra layer
            return torch.tensor_split(torch.arange(num_layers),
                                      self.pp_size)[self.pp_rank].tolist()

    def pp_rank_of_layer(self, layer_idx: int, num_layers: int) -> int:
        """Return pipeline-parallel rank that owns `layer_idx` for a model with `num_layers` layers.