ggmlR (version 0.8.1)

ggml_pp_forward: Pipeline-parallel forward pass across per-device layer stages

Description

Runs a forward pass split BY LAYERS across devices (pipeline parallelism), the complement of ggml_vulkan_split_mul_mat's split-by-matrix tensor parallelism. Each stage owns a contiguous block of the model's layers on one GPU; the activation tensor is handed from one stage to the next exactly once per pass (a single cross-device copy), versus TP's per-layer gather. This suits models too large for one card's VRAM: at ~1 GB/s host staging the single handoff costs only ~10-20 ms per pass.

Usage

ggml_pp_forward(stages, x, out_shape, mem_per_stage = 16L * 1024L * 1024L)

Value

A numeric vector of the final stage's output activations (ggml order).

Arguments

stages

A list of stage descriptors (see Details), in pipeline order.

x

A numeric vector/array of input activations for the first stage, laid out in ggml (column-major) order matching stages[[1]]$in_shape.

out_shape

Integer vector: the ggml ne shape of the final stage's output, used to size the returned vector.

mem_per_stage

Bytes of ggml context metadata to reserve per stage (default 16 MiB) — raise it for stages with very many ops.

Details

Each stage is a list with:

device

Physical GPU index (0-based) the stage runs on.

build

A function function(ctx, input) that builds this stage's sub-graph from the input tensor using the ggml_* ops and returns either the output tensor, or a list list(output = <tensor>, set_weights = function() ...). The set_weights closure (if given) is called AFTER the stage's tensors are allocated on the device, and is where weight tensors created inside build get their values (via ggml_backend_tensor_set_data).

in_shape

Integer vector: the ggml ne shape of this stage's input tensor (fastest dim first). For stage 1 this must match the supplied x; for later stages it must match the previous stage's output shape.

See Also

ggml_tp_dp_forward for the tensor-parallel counterpart.

Examples

# \donttest{
if (ggml_vulkan_available() && ggml_vulkan_device_count() >= 2) {
  K <- 32L; M <- 4L
  W1 <- matrix(rnorm(K * K), K); W2 <- matrix(rnorm(K * K), K)
  stage <- function(dev, Wt) list(
    device = dev, in_shape = c(K, M),
    build = function(ctx, input) {
      w <- ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, K)
      ggml_mul_mat(ctx, w, input)   # (weights set inside build via attr, see vignette)
    })
  # See inst/examples/pp_pipeline.R for a complete runnable stage definition.
}
# }