Tensor parallelism: split individual tensors across devices

Column-parallel: shard W by columns, y_i = x W_i, concatenate (no communication). Row-parallel: shard W by rows and x by columns to match, y = sum_i x_i W_i (all-reduce).

Usage

morie_geron_tensor_parallelism(
  model,
  n_devices = 2,
  x = NULL,
  scheme = "column"
)

Arguments

model

A weight matrix or list of weight matrices.

n_devices

Number of devices.

x

Optional input activations; default ones.

scheme

"column" or "row".

Value

list with output, reference, max_diff, shards, params_per_device, params_total, comm_elements, all_reduce, scheme, estimate, n, method.

Examples

W <- diag(2)
morie_geron_tensor_parallelism(W, 2, x = matrix(c(1, 1), 1, 2))
#> $output
#>      [,1] [,2]
#> [1,]    1    1
#> 
#> $reference
#>      [,1] [,2]
#> [1,]    1    1
#> 
#> $max_diff
#> [1] 0
#> 
#> $shards
#> $shards[[1]]
#> $shards[[1]][[1]]
#>      [,1]
#> [1,]    1
#> [2,]    0
#> 
#> $shards[[1]][[2]]
#>      [,1]
#> [1,]    0
#> [2,]    1
#> 
#> 
#> 
#> $params_per_device
#> [1] 2
#> 
#> $params_total
#> [1] 4
#> 
#> $comm_elements
#> [1] 0
#> 
#> $all_reduce
#> [1] FALSE
#> 
#> $scheme
#> [1] "column"
#> 
#> $estimate
#> [1] 0
#> 
#> $n
#> [1] 1
#> 
#> $method
#> [1] "column-parallel sharding across 2 devices, verified against the unsharded computation"
#>