dimwit-sharding

Data parallelism for DimWit: split a tensor axis across a mesh of devices and keep the split visible in the type. A Mesh binds mesh axis labels (trait X derives MeshLabel) to devices, t.shard(mesh, Axis[Batch] -> MeshAxis[X]) places one shard per device and renames the axis to Batch |@| X, and unshard gathers it back. Because Batch |@| X is just another axis label, every DimWit operation (sum, vmap, dot, grad, jit, ...) works unchanged on sharded tensors, XLA turns reductions over the sharded axis into all-reduces, and the compiler rejects mixing sharded with unsharded tensors. It works with any DimWit release from 0.1.0 on and only uses DimWit's public API.

Usage

libraryDependencies ++= Seq(
  "ch.contrafactus" %% "dimwit-core" % "0.1.0", // or any newer version
  "ch.contrafactus" %% "dimwit-sharding" % "0.1.0-SNAPSHOT"
)

dimwit-sharding does not pull in dimwit-core itself (it is a provided dependency), so the DimWit version is entirely your choice.

import dimwit.*
import dimwit.sharding.*

trait Batch derives Label
trait Feature derives Label
trait X derives MeshLabel

val data: Tensor2[Batch, Feature, Float32] =
  Tensor2(Axis[Batch], Axis[Feature]).fromArray(Array.tabulate(8, 4)((i, j) => (i * 4 + j).toFloat))

val mesh = Mesh1(MeshAxis[X] -> 4)
val sharded: Tensor2[Batch |@| X, Feature, Float32] = data.shard(mesh, Axis[Batch] -> MeshAxis[X])

// An ordinary sum; XLA makes it an all-reduce across the 4 devices.
val total: Tensor1[Feature, Float32] = sharded.sum(Axis[Batch |@| X])
val gathered: Tensor2[Batch, Feature, Float32] = sharded.unshard(Axis[Batch |@| X])

A single process normally sees one CPU device. To try this on a laptop, ask XLA for more (before JAX is initialised): XLA_FLAGS=--xla_force_host_platform_device_count=8, or Jax.jax.config.update("jax_num_cpu_devices", 4) as in the examples.

Development

The build compiles against DimWit 0.1.0 by default. To build against a newer DimWit, e.g. a locally published snapshot:

sbt -Ddimwit.version=0.2-SNAPSHOT -Dscala.version=3.9.0 test

sbt docsRoot/mdoc regenerates this README from mdocs/README.md.