gistfile1.txt

    def pipeline(self, group):
        self.pipeline_rank = group.rank()
        self.pipeline_size = group.size()
        total_layers = len(self.layers)

        # Rank 0 gets a fixed number of layers (from the end)
        static_count = 20

        if self.pipeline_rank == 0:
            self.start_idx = total_layers - static_count
            self.end_idx = total_layers
        else:
            leftover_layers = total_layers - static_count
            leftover_ranks = self.pipeline_size - 1
            leftover_rank = self.pipeline_rank - 1

            layers_per_rank = leftover_layers // leftover_ranks
            extra = leftover_layers - layers_per_rank * leftover_ranks

            if leftover_rank < extra:
                layers_per_rank += 1

            # Distribute the remainder in reverse (same logic as original)
            self.start_idx = (leftover_ranks - leftover_rank - 1) * layers_per_rank
            self.end_idx = self.start_idx + layers_per_rank

        self.layers = self.layers[: self.end_idx]
        self.layers[: self.start_idx] = [None] * self.start_idx
        self.num_layers = len(self.layers) - self.start_idx