diff --git a/build.sbt b/build.sbt index 54fdaa0..a75c195 100644 --- a/build.sbt +++ b/build.sbt @@ -84,6 +84,8 @@ lazy val core = (project in file("core")) fork := true, javaOptions ++= scalapyJavaOptions, Test / envVars += "DIMWIT_SKIP_SYNC" -> "true", + // The sharding tests need a mesh of devices; a default JAX process reports a single CPU device. + Test / envVars += "XLA_FLAGS" -> "--xla_force_host_platform_device_count=8", coverageMinimumStmtTotal := 80, coverageFailOnMinimum := false, coverageHighlighting := true, diff --git a/core/src/main/resources/python/jax_helper.py b/core/src/main/resources/python/jax_helper.py index fe1b761..7a73731 100644 --- a/core/src/main/resources/python/jax_helper.py +++ b/core/src/main/resources/python/jax_helper.py @@ -85,3 +85,27 @@ def jit(f): def jit_fn(f, jit_kwargs=None): return wrap(jax.jit, f, kwargs=jit_kwargs) + +# --- sharding ------------------------------------------------------------- + +def device_mesh(devices, shape, axis_names): + """Builds a jax.sharding.Mesh from a flat, row-major list of devices.""" + import numpy as np + grid = np.array(list(devices), dtype=object).reshape(tuple(shape)) + return jax.sharding.Mesh(grid, axis_names=tuple(axis_names)) + +def named_sharding(mesh, rank, axis_index, mesh_axis_name): + """NamedSharding splitting axis `axis_index` of a rank-`rank` array over `mesh_axis_name`. + + Every other axis is replicated. + """ + spec = [None] * rank + spec[axis_index] = mesh_axis_name + return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) + +def replicate(x): + """Gathers a sharded array so that every device of its mesh holds the whole thing.""" + sharding = x.sharding + if not hasattr(sharding, "mesh"): + return x + return jax.device_put(x, jax.sharding.NamedSharding(sharding.mesh, jax.sharding.PartitionSpec())) diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index 555ecde..00c4c61 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/core/src/main/scala/dimwit/package.scala @@ -39,6 +39,16 @@ package object dimwit: given [A, B](using labelA: Label[A], labelB: Label[B]): Label[A |+| B] with val name: String = s"${labelA.name}+${labelB.name}" + /** Marks the data axis `A` as sharded over the mesh axis `M`, as in `Tensor2[Batch |@| X, Feature, Float32]`. + * + * `Batch |@| X` is just another axis label, so every tensor operation applies unchanged. + */ + @targetName("Sharded") + infix trait |@|[A, M] + object `|@|`: + given [A, M](using label: Label[A], meshLabel: MeshLabel[M]): Label[A |@| M] with + val name: String = s"${label.name}@${meshLabel.name}" + // Export tensor and related types export dimwit.tensor.{Tensor, Tensor0, Tensor1, Tensor2, Tensor3, Tensor4, TypedIndex} export dimwit.tensor.{Shape, Shape0, Shape1, Shape2, Shape3} @@ -88,6 +98,10 @@ package object dimwit: export dimwit.stats.{Prob, LogProb} export dimwit.stats.{Distribution, IndependentDistribution, MultivariateDistribution, UnivariateDistribution} + // Export sharding stuff + export dimwit.sharding.{Mesh, Mesh1, Mesh2, Mesh3, MeshAxis, MeshAxisExtent, MeshAxisIndex, MeshLabel, MeshLabels} + export dimwit.sharding.ShardingOps.* + /** Memory management helpfer making sure * all python objects allocated ar freed * after the function is executed. diff --git a/core/src/main/scala/dimwit/sharding/Mesh.scala b/core/src/main/scala/dimwit/sharding/Mesh.scala new file mode 100644 index 0000000..dab3ec5 --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/Mesh.scala @@ -0,0 +1,83 @@ +package dimwit.sharding + +import dimwit.hardware.Device +import dimwit.jax.Jax +import me.shadaj.scalapy.py.SeqConverters + +import scala.annotation.implicitNotFound + +/** A grid of devices, binding each mesh axis label of `M` to a number of devices. */ +final class Mesh[M <: Tuple: MeshLabels] private[sharding] ( + val axisSizes: List[Int], + val devices: Seq[Device] +): + + lazy val axisNames: List[String] = MeshLabels[M].names + + def size: Int = axisSizes.product + + def sizeOf[A](meshAxis: MeshAxis[A])(using ev: MeshAxisIndex[M, A]): Int = axisSizes(ev.index) + + private[dimwit] lazy val jaxMesh: Jax.PyDynamic = + Jax.jax_helper + .device_mesh(devices.map(_.toJaxDevice).toPythonProxy, axisSizes.toPythonProxy, axisNames.toPythonProxy) + .as[Jax.PyDynamic] + + override def toString: String = + axisNames.zip(axisSizes).map((name, size) => s"$name -> $size").mkString("Mesh(", ", ", ")") + +object Mesh: + + private[sharding] type ExtractLabel[Extent] = Extent match + case MeshAxisExtent[a] => a + + private[sharding] type ExtractLabels[Extents <: Tuple] = Tuple.Map[Extents, ExtractLabel] + + def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh1[A] = Mesh1(extent) + def apply[A: MeshLabel, B: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B]): Mesh2[A, B] = Mesh2(extent1, extent2) + def apply[A: MeshLabel, B: MeshLabel, C: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B], extent3: MeshAxisExtent[C]): Mesh3[A, B, C] = Mesh3(extent1, extent2, extent3) + def apply[Extents <: Tuple](extents: Extents)(using MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = fromTuple(extents) + + def fromTuple[Extents <: Tuple](extents: Extents)(using labels: MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = + val sizes = extents.toList.collect: + case extent: MeshAxisExtent[?] => extent.size + val required = sizes.product + val available = Jax.devices + require( + available.size >= required, + s"Mesh ${labels.names.zip(sizes).map((name, size) => s"$name -> $size").mkString("(", ", ", ")")} requires $required devices, but only ${available.size} are available" + ) + new Mesh(sizes, available.take(required)) + +type Mesh1[A] = Mesh[Tuple1[A]] +type Mesh2[A, B] = Mesh[(A, B)] +type Mesh3[A, B, C] = Mesh[(A, B, C)] + +object Mesh1: + def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh1[A] = Mesh.fromTuple(Tuple1(extent)) + +object Mesh2: + def apply[A: MeshLabel, B: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B]): Mesh2[A, B] = + Mesh.fromTuple((extent1, extent2)) + +object Mesh3: + def apply[A: MeshLabel, B: MeshLabel, C: MeshLabel]( + extent1: MeshAxisExtent[A], + extent2: MeshAxisExtent[B], + extent3: MeshAxisExtent[C] + ): Mesh3[A, B, C] = Mesh.fromTuple((extent1, extent2, extent3)) + +/** Finds the position of a mesh axis in a mesh. */ +@implicitNotFound("MeshAxis[${A}] not found in Mesh[${M}]") +trait MeshAxisIndex[M <: Tuple, A]: + def index: Int + +object MeshAxisIndex: + + def apply[M <: Tuple, A](using idx: MeshAxisIndex[M, A]): Int = idx.index + + given found[A, Tail <: Tuple]: MeshAxisIndex[A *: Tail, A] with + val index = 0 + + given search[H, T <: Tuple, A](using next: MeshAxisIndex[T, A]): MeshAxisIndex[H *: T, A] with + val index = 1 + next.index diff --git a/core/src/main/scala/dimwit/sharding/MeshAxis.scala b/core/src/main/scala/dimwit/sharding/MeshAxis.scala new file mode 100644 index 0000000..2394d3e --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/MeshAxis.scala @@ -0,0 +1,14 @@ +package dimwit.sharding + +/** Represents an axis of a device [[Mesh]]. The mesh-side counterpart of [[dimwit.tensor.Axis]]. */ +final class MeshAxis[A: MeshLabel]: + + def name: String = summon[MeshLabel[A]].name + + def extent(size: Int): MeshAxisExtent[A] = MeshAxisExtent(this, size) + def ->(size: Int): MeshAxisExtent[A] = this.extent(size) + + override def toString: String = s"MeshAxis($name)" + +/** A mesh axis together with the number of devices along it. */ +case class MeshAxisExtent[A: MeshLabel](axis: MeshAxis[A], size: Int) diff --git a/core/src/main/scala/dimwit/sharding/MeshLabels.scala b/core/src/main/scala/dimwit/sharding/MeshLabels.scala new file mode 100644 index 0000000..77639da --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/MeshLabels.scala @@ -0,0 +1,51 @@ +package dimwit.sharding + +import scala.quoted.* + +/** A label for an axis of a device [[Mesh]], declared with `trait X derives MeshLabel`. + * The mesh-side counterpart of [[dimwit.tensor.Label]]. + */ +@scala.annotation.implicitNotFound(""" +A mesh axis label ${T} was given or inferred, which does not have a MeshLabel instance. +Ensure that all mesh axis types ${T} are defined with 'derives MeshLabel' (e.g. 'trait X derives MeshLabel') +""") +trait MeshLabel[T]: + def name: String + +object MeshLabel: + + def apply[T](using meshLabel: MeshLabel[T]): MeshLabel[T] = meshLabel + + inline def derived[T]: MeshLabel[T] = ${ derivedMacro[T] } + + private def derivedMacro[T: Type](using Quotes): Expr[MeshLabel[T]] = + import quotes.reflect.* + val tpe = TypeRepr.of[T] + val simpleName = tpe.typeSymbol.name + '{ + new MeshLabel[T]: + def name: String = ${ Expr(simpleName) } + } + +@scala.annotation.implicitNotFound(""" +A tuple of mesh axis labels ${T} was given or inferred that does not have a valid MeshLabels instance. + +Ensure that all of the types in the tuple have a 'derives MeshLabel' clause. +""") +trait MeshLabels[T]: + def names: List[String] + +private class MeshLabelsImpl[T](val names: List[String]) extends MeshLabels[T] + +object MeshLabels: + + def apply[T](using labels: MeshLabels[T]): MeshLabels[T] = labels + + given emptyTuple: MeshLabels[EmptyTuple] = new MeshLabelsImpl[EmptyTuple](Nil) + + given lift[A](using v: MeshLabel[A]): MeshLabels[A] = new MeshLabelsImpl[A](List(v.name)) + + given consTuple[H, T <: Tuple](using + head: MeshLabel[H], + tail: MeshLabels[T] + ): MeshLabels[H *: T] = new MeshLabelsImpl[H *: T](head.name :: tail.names) diff --git a/core/src/main/scala/dimwit/sharding/ShardingOps.scala b/core/src/main/scala/dimwit/sharding/ShardingOps.scala new file mode 100644 index 0000000..a00d391 --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/ShardingOps.scala @@ -0,0 +1,51 @@ +package dimwit.sharding + +import dimwit.jax.Jax +import dimwit.|@| +import dimwit.tensor.Axis +import dimwit.tensor.Label +import dimwit.tensor.Labels +import dimwit.tensor.ShapeTypeHelpers.AxisReplacer +import dimwit.tensor.Tensor + +object ShardingOps: + + extension [T <: Tuple: Labels, V](t: Tensor[T, V]) + + /** Splits `t` across `mesh` along one axis, placing one shard per device and rewriting that axis from `L` to `L |@| A`. + * + * {{{ + * val mesh = Mesh1(MeshAxis[X] -> 4) + * val sharded: Tensor2[Batch |@| X, Feature, Float32] = t.shard(mesh, Axis[Batch] -> MeshAxis[X]) + * }}} + * + * @throws IllegalArgumentException if the extent of `L` is not divisible by the mesh axis size. + */ + def shard[M <: Tuple, L: Label, A: MeshLabel](mesh: Mesh[M], mapping: (Axis[L], MeshAxis[A]))(using + replacer: AxisReplacer[T, L, L |@| A], + meshIndex: MeshAxisIndex[M, A], + labels: Labels[replacer.NewShape] + ): Tensor[replacer.NewShape, V] = + val axisName = summon[Label[L]].name + val meshAxisName = mapping._2.name + val extent = t.shape.dimensions(replacer.index) + val meshSize = mesh.axisSizes(meshIndex.index) + if extent % meshSize != 0 then + throw new IllegalArgumentException( + s"Cannot shard axis $axisName of extent $extent over mesh axis $meshAxisName of size $meshSize: " + + s"$extent is not divisible by $meshSize." + ) + val sharding = Jax.jax_helper.named_sharding(mesh.jaxMesh, t.shape.rank, replacer.index, meshAxisName) + Tensor[replacer.NewShape, V](Jax.device_put(t.jaxValue, sharding.as[Jax.PyDynamic])) + + /** Gathers a sharded axis back across the mesh, rewriting it from `L |@| A` to `L` so that every device holds the whole tensor. + * + * {{{ + * val gathered: Tensor2[Batch, Feature, Float32] = sharded.unshard(Axis[Batch |@| X]) + * }}} + */ + def unshard[L, A](axis: Axis[L |@| A])(using + replacer: AxisReplacer[T, L |@| A, L], + labels: Labels[replacer.NewShape] + ): Tensor[replacer.NewShape, V] = + Tensor[replacer.NewShape, V](Jax.jax_helper.replicate(t.jaxValue).as[Jax.PyDynamic]) diff --git a/core/src/test/scala/dimwit/sharding/ShardingSuite.scala b/core/src/test/scala/dimwit/sharding/ShardingSuite.scala new file mode 100644 index 0000000..480ebfd --- /dev/null +++ b/core/src/test/scala/dimwit/sharding/ShardingSuite.scala @@ -0,0 +1,202 @@ +package dimwit.sharding + +import dimwit.* +import dimwit.hardware.DeviceBackend +import me.shadaj.scalapy.py + +import scala.compiletime.testing.typeCheckErrors + +/** Mesh axis labels: a different kind than the data axis labels `A`, `B`, ... in the test package. */ +trait X derives MeshLabel +trait Y derives MeshLabel + +class ShardingSuite extends DimwitTest: + + private val BatchExtent = 8 + private val FeatureExtent = 4 + private val MeshExtent = 4 + + private def tensorOf(batch: Int, feature: Int = FeatureExtent): Tensor2[A, B, Float32] = + Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(batch, feature)((i, j) => (i * feature + j).toFloat)) + + private val t: Tensor2[A, B, Float32] = tensorOf(BatchExtent) + + private def enoughDevices(): Unit = + assume( + DeviceBackend.CPU.devices.size >= MeshExtent, + s"needs at least $MeshExtent devices, found ${DeviceBackend.CPU.devices.size}" + ) + + private def isFullyReplicated[T <: Tuple: Labels, V](tensor: Tensor[T, V]): Boolean = + tensor.jaxValue.sharding.is_fully_replicated.as[Boolean] + + private def deviceCount[T <: Tuple: Labels, V](tensor: Tensor[T, V]): Int = + tensor.jaxValue.sharding.num_devices.as[Int] + + private def mesh1 = Mesh1(MeshAxis[X] -> MeshExtent) + + private def sharded: Tensor2[A |@| X, B, Float32] = t.shard(mesh1, Axis[A] -> MeshAxis[X]) + + describe("Mesh"): + + it("binds mesh axis labels to a grid of devices"): + enoughDevices() + val mesh = mesh1 + mesh.size shouldBe MeshExtent + mesh.axisNames shouldBe List("X") + mesh.axisSizes shouldBe List(MeshExtent) + mesh.sizeOf(MeshAxis[X]) shouldBe MeshExtent + mesh.devices.map(_.id).distinct should have size MeshExtent + + it("is built from a tuple of extents, with Mesh1/Mesh2 as special cases"): + enoughDevices() + val single = Mesh1(MeshAxis[X] -> MeshExtent) + single.axisNames shouldBe List("X") + single.axisSizes shouldBe List(MeshExtent) + + val grid = Mesh((MeshAxis[X] -> 2, MeshAxis[Y] -> 2)) + grid.axisNames shouldBe List("X", "Y") + grid.size shouldBe 4 + grid.sizeOf(MeshAxis[Y]) shouldBe 2 + + Mesh1(MeshAxis[X] -> MeshExtent).axisSizes shouldBe single.axisSizes + Mesh2(MeshAxis[X] -> 2, MeshAxis[Y] -> 2).axisNames shouldBe grid.axisNames + + it("fails with a clear error when there are not enough devices"): + val error = intercept[IllegalArgumentException](Mesh1(MeshAxis[X] -> 1000000)) + error.getMessage should include("requires 1000000 devices") + + describe("shard"): + + it("rewrites the axis type and places the shards on the mesh devices"): + enoughDevices() + val result = sharded + result.axes shouldBe List("A@X", "B") + result.shape.dimensions shouldBe List(BatchExtent, FeatureExtent) + deviceCount(result) shouldBe MeshExtent + isFullyReplicated(result) shouldBe false + result shouldEqual t + + it("rejects an axis whose extent does not divide over the mesh axis"): + enoughDevices() + val error = intercept[IllegalArgumentException](tensorOf(7).shard(mesh1, Axis[A] -> MeshAxis[X])) + error.getMessage shouldBe "Cannot shard axis A of extent 7 over mesh axis X of size 4: 7 is not divisible by 4." + + describe("unshard"): + + it("gathers the axis back and drops the mesh annotation"): + enoughDevices() + val gathered = sharded.unshard(Axis[A |@| X]) + gathered.axes shouldBe List("A", "B") + gathered shouldEqual t + isFullyReplicated(gathered) shouldBe true + + it("round-trips with shard"): + enoughDevices() + sharded.unshard(Axis[A |@| X]).shard(mesh1, Axis[A] -> MeshAxis[X]) shouldEqual t + + it("gives back a tensor that mixes with unsharded ones again"): + enoughDevices() + val gathered = sharded.unshard(Axis[A |@| X]) + (gathered + t) shouldEqual (t + t) + + it("cannot be applied to an axis that is not sharded"): + enoughDevices() + val errors = typeCheckErrors("sharded.unshard(Axis[B])") + errors should not be empty + + describe("a sharded axis is just another axis"): + + it("reducing the sharded axis all-reduces and equals the unsharded sum, bit for bit"): + enoughDevices() + val result = sharded.sum(Axis[A |@| X]) + result shouldEqual t.sum(Axis[A]) + result.axes shouldBe List("B") + isFullyReplicated(result) shouldBe true + + it("reducing every axis works on both, and is how one loss function serves both runs"): + enoughDevices() + sharded.sum shouldEqual t.sum + sharded.mean shouldEqual t.mean + + it("reducing another axis is local and keeps the mesh annotation"): + enoughDevices() + val result = sharded.sum(Axis[B]) + result shouldEqual t.sum(Axis[B]) + result.axes shouldBe List("A@X") + isFullyReplicated(result) shouldBe false + deviceCount(result) shouldBe MeshExtent + + it("other reductions over the sharded axis work too, with no extra support"): + enoughDevices() + sharded.mean(Axis[A |@| X]) shouldEqual t.mean(Axis[A]) + sharded.max(Axis[A |@| X]) shouldEqual t.max(Axis[A]) + sharded.min(Axis[A |@| X]) shouldEqual t.min(Axis[A]) + + it("vmap maps over the sharded axis, which stays sharded"): + enoughDevices() + val result = sharded.vmap(Axis[A |@| X])(row => row *! Tensor0(2.0f)) + result shouldEqual t.vmap(Axis[A])(row => row *! Tensor0(2.0f)) + result.axes shouldBe List("A@X", "B") + deviceCount(result) shouldBe MeshExtent + + it("contracts a replicated tensor against the unsharded axis, staying sharded"): + enoughDevices() + val weights = Tensor1(Axis[B]).fromArray(Array.fill(FeatureExtent)(2.0f)) + val result = sharded.dot(Axis[B])(weights) + result shouldEqual t.dot(Axis[B])(weights) + result.axes shouldBe List("A@X") + deviceCount(result) shouldBe MeshExtent + + it("broadcasts a replicated tensor over the axes it does share"): + enoughDevices() + val bias = Tensor1(Axis[B]).fromArray(Array.fill(FeatureExtent)(1.0f)) + val result = sharded -! bias + result shouldEqual (t -! bias) + result.axes shouldBe List("A@X", "B") + deviceCount(result) shouldBe MeshExtent + + it("zips with another tensor sharded the same way"): + enoughDevices() + val labels = Tensor1(Axis[A]).fromArray(Array.tabulate(BatchExtent)(_.toFloat)) + val shardedLabels = labels.shard(mesh1, Axis[A] -> MeshAxis[X]) + val result = zipvmap(Axis[A |@| X])(sharded, shardedLabels) { case (row, label) => row.sum + label } + result.axes shouldBe List("A@X") + result shouldEqual zipvmap(Axis[A])(t, labels) { case (row, label) => row.sum + label } + + it("reads back to the host, gathering from every device"): + enoughDevices() + sharded.toArray shouldEqual t.toArray + + describe("sharded and unsharded tensors do not mix"): + + it("a sharded tensor cannot be combined with an unsharded one"): + enoughDevices() + val errors = typeCheckErrors("sharded + t") + errors should not be empty + + it("the sharded axis cannot be named without its mesh axis"): + enoughDevices() + val errors = typeCheckErrors("sharded.sum(Axis[A])") + errors should not be empty + errors.head.message should include("not found in Tensor") + + describe("mesh labels are a separate kind from data axis labels"): + + it("a data axis label cannot be used as a mesh axis label"): + val errors = typeCheckErrors("MeshAxis[A]") + errors should not be empty + errors.head.message should include("MeshLabel") + + it("a data axis label cannot be used to build a mesh"): + val errors = typeCheckErrors("Mesh1(MeshAxis[A] -> 4)") + errors should not be empty + + it("a mesh axis label cannot be used as a data axis label"): + val errors = typeCheckErrors("Axis[X]") + errors should not be empty + errors.head.message should include("Label") + + it("a data axis label cannot stand on the mesh side of |@|"): + val errors = typeCheckErrors("summon[Label[A |@| B]]") + errors should not be empty diff --git a/examples/src/main/scala/dimwit/basic/ShardedLinearRegression.scala b/examples/src/main/scala/dimwit/basic/ShardedLinearRegression.scala new file mode 100644 index 0000000..801bcd5 --- /dev/null +++ b/examples/src/main/scala/dimwit/basic/ShardedLinearRegression.scala @@ -0,0 +1,85 @@ +package dimwit.examples.basic + +import dimwit.Conversions.given +import dimwit.* +import dimwit.autodiff.* +import dimwit.jax.Jax +import dimwit.optimizer.GradientDescent +import dimwit.random.Random +import dimwit.stats.Normal + +object ShardedLinearRegression: + + val Devices = 4 + val BatchSize = 8192 // must be divisible by Devices + val FeatureSize = 64 + val Steps = 200 + val LearningRate = 0.05f + + trait Batch derives Label + trait Feature derives Label + + case class Params(weights: Tensor1[Feature, Float32], bias: Tensor0[Float32]) + object Params: + def init: Params = Params(Tensor1(Axis[Feature]).fromArray(Array.fill(FeatureSize)(0f)), Tensor0(0f)) + + class LinearRegression(params: Params) extends (Tensor1[Feature, Float32] => Tensor0[Float32]): + def apply(x: Tensor1[Feature, Float32]): Tensor0[Float32] = x.dot(Axis[Feature])(params.weights) + params.bias + + @main def runSingleDevice(): Unit = + dimwit.initialize() + + val (x, y) = dataset(Random.Key(0)) + + def loss(params: Params): Tensor0[Float32] = + val model = LinearRegression(params) + val predictions = x.vmap(Axis[Batch])(model) + // `sum` names no axis, so this line is identical in the sharded run. + (predictions - y).pow(2f).sum / Tensor0(BatchSize.toFloat) + + val descent = GradientDescent(Tensor0(LearningRate)).iterate(Params.init)(grad(jit(loss))) + descent.next().bias.item // warm up JIT + val start = System.nanoTime() + val params = descent.drop(Steps - 1).next() + params.bias.item + report("single device", params, (System.nanoTime() - start) / 1e9) + + @main def runSharded(): Unit = + dimwit.initialize() + Jax.jax.config.update("jax_num_cpu_devices", Devices) + + val (x, y) = dataset(Random.Key(0)) + + trait X derives MeshLabel + val mesh = Mesh1(MeshAxis[X] -> Devices) + + // Inputs and targets are split the same way; the parameters stay replicated. + val shardedX: Tensor2[Batch |@| X, Feature, Float32] = x.shard(mesh, Axis[Batch] -> MeshAxis[X]) + val shardedY: Tensor1[Batch |@| X, Float32] = y.shard(mesh, Axis[Batch] -> MeshAxis[X]) + + println(s"$mesh over devices ${mesh.devices.map(_.id).mkString(", ")}") + println(s"batch ${shardedX.axes.mkString(" x ")}, ${BatchSize / Devices} rows per device") + + def loss(params: Params): Tensor0[Float32] = + val model = LinearRegression(params) + // The batch axis is called `Batch |@| X` now; `vmap` treats it like any other. + val predictions = shardedX.vmap(Axis[Batch |@| X])(model) + (predictions - shardedY).pow(2f).sum / Tensor0(BatchSize.toFloat) + + val descent = GradientDescent(Tensor0(LearningRate)).iterate(Params.init)(grad(jit(loss))) + descent.next().bias.item // warm up JIT + val start = System.nanoTime() + val params = descent.drop(Steps - 1).next() + params.bias.item + report("sharded", params, (System.nanoTime() - start) / 1e9) + + private def dataset(key: Random.Key): (Tensor2[Batch, Feature, Float32], Tensor1[Batch, Float32]) = + val (xKey, noiseKey) = key.split2() + val x = Normal.standardNormal(Shape(Axis[Batch] -> BatchSize, Axis[Feature] -> FeatureSize)).sample(xKey) + val trueWeights = Tensor1(Axis[Feature]).fromArray(Array.tabulate(FeatureSize)(i => (i % 5).toFloat - 2f)) + val noise = Normal.standardNormal(Shape(Axis[Batch] -> BatchSize)).sample(noiseKey) *! Tensor0(0.1f) + (x, x.dot(Axis[Feature])(trueWeights) + noise) + + private def report(setting: String, params: Params, seconds: Double): Unit = + println(f"$setting%-14s $Steps steps in $seconds%6.3f s (${seconds / Steps * 1000}%5.2f ms/step)") + println(f"${" "}%-14s final bias ${params.bias.item}%+.4f, weights[0] ${params.weights.slice(Axis[Feature].at(0)).item}%+.4f") diff --git a/examples/src/main/scala/dimwit/basic/ShardedSum.scala b/examples/src/main/scala/dimwit/basic/ShardedSum.scala new file mode 100644 index 0000000..17ae013 --- /dev/null +++ b/examples/src/main/scala/dimwit/basic/ShardedSum.scala @@ -0,0 +1,60 @@ +package dimwit.examples.basic + +import dimwit.* +import dimwit.jax.Jax + +/** Summing a tensor whose `Batch` axis is split across several devices. */ +object ShardedSum: + + trait Batch derives Label + trait Feature derives Label + trait X derives MeshLabel + + /** How many CPU devices to run on. `BatchSize` has to be divisible by this. */ + val Devices = 4 + + val BatchSize = 8 + + val FeatureSize = 4 + + @main def runShardedSum(): Unit = + dimwit.initialize() + Jax.jax.config.update("jax_num_cpu_devices", Devices) + + val data: Tensor2[Batch, Feature, Float32] = + Tensor2(Axis[Batch], Axis[Feature]).fromArray( + Array.tabulate(BatchSize, FeatureSize)((batch, feature) => (batch * FeatureSize + feature).toFloat) + ) + val reference: Tensor1[Feature, Float32] = data.sum(Axis[Batch]) + + val mesh = Mesh1(MeshAxis[X] -> Devices) + println(s"$mesh on devices ${mesh.devices.map(_.id).mkString(", ")}") + + val sharded = data.shard(mesh, Axis[Batch] -> MeshAxis[X]) + println(s"sharded axes: ${sharded.axes.mkString(", ")} (${BatchSize / Devices} rows per device)") + + // Reducing the sharded axis is the ordinary `sum`; XLA makes it an all-reduce. + val total: Tensor1[Feature, Float32] = sharded.sum(Axis[Batch |@| X]) + + println(s"unsharded sum: $reference") + println(s"sharded sum: $total") + println(s"identical: ${total == reference}") + + // Naming no axis works on both, so one function can serve a sharded and an unsharded run. + println(s"total of everything, either way: ${data.sum} / ${sharded.sum}") + + // Reducing an unsharded axis needs no communication; Batch stays spread over X. + val perRow: Tensor1[Batch |@| X, Float32] = sharded.sum(Axis[Feature]) + println(s"per-row sums, still sharded over X: $perRow") + + val doubled: Tensor2[Batch |@| X, Feature, Float32] = + sharded.vmap(Axis[Batch |@| X])(row => row *! Tensor0(2.0f)) + println(s"each row doubled, still sharded over X: ${doubled.axes.mkString(", ")}") + println(s"mean over the sharded axis: ${sharded.mean(Axis[Batch |@| X])}") + + // Things to try, each of which should fail: + // sharded.sum(Axis[Batch]) - the sharded axis is called Batch |@| X now (compile error) + // sharded + data - sharded and unsharded tensors do not mix (compile error) + // MeshAxis[Batch] - a data axis label is not a mesh label (compile error) + // Axis[X] - and a mesh label is not a data axis label (compile error) + // BatchSize = 7 - 7 rows do not divide over 4 devices (runtime error)