ggml_pp_dp_forward: Pipeline + data parallel forward pass (PP x DP)

Description

Runs data parallelism over the batch across replicas of a pipeline-parallel model: the batch is split into contiguous shards, one per replica, and each replica runs the whole model as a device-by-layer pipeline on its own set of GPUs (via ggml_pp_forward). This is the pipeline counterpart of ggml_tp_dp_forward — use it when the model is too large for one card (split it by layers) and you also want to raise throughput by replicating the pipeline across independent GPU sets.

Usage

ggml_pp_dp_forward(
  make_stages,
  x,
  replicas,
  out_ncol,
  mem_per_stage = 16L * 1024L * 1024L
)

Value

The M x N result matrix, the batch-concatenation of each replica's pipeline output.

Arguments

make_stages: A function function(devices, m_shard) returning the list of pipeline stage descriptors (see ggml_pp_forward) for a replica whose pipeline runs on devices (an integer GPU-index vector) and whose batch shard has m_shard samples. The stages' in_shape / weight setup must be built for that m_shard.
x: Input activation matrix, M x K (M samples, K features). Its rows (the batch) are split across the replicas.
replicas: A list of integer vectors, each the GPU-index pipeline for one replica, e.g. list(c(0, 1), c(2, 3)) for PP=2 x DP=2.
out_ncol: Number of columns N of the per-sample output (the final stage produces an N x m_shard result per replica; the gathered result is M x N).
mem_per_stage: Bytes of ggml context metadata per stage (see ggml_pp_forward).

Details

On a 4-GPU box the natural layout is PP=2 x DP=2: replica A pipelines its layers across GPUs c(0, 1), replica B across c(2, 3), and the two replicas each handle half the batch with no cross-replica traffic.

Examples

Run this code

# \donttest{
if (ggml_vulkan_available() && ggml_vulkan_device_count() >= 4) {
  K <- 64L
  W1 <- matrix(rnorm(K * K), K); W2 <- matrix(rnorm(K * K), K)
  make_stages <- function(devices, m) {
    mk <- function(dev, Wt, relu) list(
      device = dev, in_shape = c(K, m),
      build = function(ctx, input) {
        w <- ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, K)
        z <- ggml_mul_mat(ctx, w, input)
        list(output = if (relu) ggml_relu(ctx, z) else z,
             set_weights = function() ggml_backend_tensor_set_data(w, as.numeric(Wt)))
      })
    list(mk(devices[1], W1, TRUE), mk(devices[2], W2, FALSE))
  }
  X <- matrix(rnorm(8 * K), nrow = 8)   # batch of 8
  Y <- ggml_pp_dp_forward(make_stages, X, replicas = list(c(0, 1), c(2, 3)),
                          out_ncol = K)
}
# }