View source: R/tensor_parallel.R
| ggml_pp_forward | R Documentation |
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.
ggml_pp_forward(stages, x, out_shape, mem_per_stage = 16L * 1024L * 1024L)
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 |
out_shape |
Integer vector: the ggml |
mem_per_stage |
Bytes of ggml context metadata to reserve per stage (default 16 MiB) — raise it for stages with very many ops. |
Each stage is a list with:
devicePhysical GPU index (0-based) the stage runs on.
buildA 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_shapeInteger 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.
A numeric vector of the final stage's output activations (ggml order).
ggml_tp_dp_forward for the tensor-parallel counterpart.
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.
}
Add the following code to your website.
For more information on customizing the embed code, read Embedding Snippets.