DRAFT: Add sharding support for data-parallelism

[benikm91]

This PR is a first draft to allow limited JAX sharding in DimWit, specifically to allow data-parallelism only (full model distributed across multiple GPUs to process different batches in parallel to speed up training).

Adding sharding for the model parameters requires some more effort (and understanding on my side) and is left to future work.

[benikm91]

Tested on virtual CPUs (mac) and on two GPUs:

[[Amebr@sanfrancisco:/usr/src/dimwit$ TMPDIR=/tmp sbt "examples/runMain dimwit.examples.basic.runSingleDevice"
[info] welcome to sbt 1.12.13 (Ubuntu Java 17.0.20)
[info] loading settings for project dimwit-build from plugins.sbt...
[info] loading project definition from /usr/src/dimwit/project
[info] loading settings for project root from build.sbt...
[info] set current project to dimwit-root (in build file:/usr/src/dimwit/)
[warn] there's a key that's not used by any other settings/tasks:
[warn]  
[warn] * ThisBuild / publishMavenStyle
[warn]   +- /usr/src/dimwit/build.sbt:19
[warn]  
[warn] note: a setting might still be used by a command; to exclude a key from this `lintUnused` check
[warn] either append it to `Global / excludeLintKeys` or call .withRank(KeyRanks.Invisible) on the key
[info] running (fork) dimwit.examples.basic.runSingleDevice 
[info] single device  200 steps in  1.409 s  ( 7.04 ms/step)
[info]                final bias +0.0007, weights[0] -1.9984
[success] Total time: 9 s, completed Sep 8, 2026, 1:45:27 PM
mebr@sanfrancisco:/usr/src/dimwit$ TMPDIR=/tmp sbt "examples/runMain dimwit.examples.basic.runSharded"
[info] welcome to sbt 1.12.13 (Ubuntu Java 17.0.20)
[info] loading settings for project dimwit-build from plugins.sbt...
[info] loading project definition from /usr/src/dimwit/project
[info] loading settings for project root from build.sbt...
[info] set current project to dimwit-root (in build file:/usr/src/dimwit/)
[warn] there's a key that's not used by any other settings/tasks:
[warn]  
[warn] * ThisBuild / publishMavenStyle
[warn]   +- /usr/src/dimwit/build.sbt:19
[warn]  
[warn] note: a setting might still be used by a command; to exclude a key from this `lintUnused` check
[warn] either append it to `Global / excludeLintKeys` or call .withRank(KeyRanks.Invisible) on the key
[info] running (fork) dimwit.examples.basic.runSharded 
[info] Mesh(X -> 2) over gpu:0, gpu:1
[info] batch Batch@X x Feature, 4096 rows per device
[info] sharded        200 steps in  2.599 s  (12.99 ms/step)
[info]                final bias +0.0007, weights[0] -1.9984
[success] Total time: 11 s, completed Sep 8, 2026, 1:45:55 PM
mebr@sanfrancisco:/usr/src/dimwit$ 

[benikm91]

yay :) Also working on 2 GPUs for the GPT-2 example:

Roughly twice the speed as we had on Spark :) I will try to finish a run on 4 GPUs this week within a day.

[benikm91]

And with 4 GPUs :)

[benikm91]

@marcelluethi

DimWit does not support multiple GPUs. JAX has a nice Sharding solution. This PR lifts a minimal version of this solution to DimWit to allow data-parallelism, the most common parallelism, where the model is small enough to fit fully on one GPU, yet the batches can be distributed across GPUs to calculate/approximate gradients.

Design decisions for this PR:

• Minimal Sharding Support: JAX Sharding is a big topic and I have little experience in it. Therefore, I decided to only lift data-parallelism, which removes many aspects of Sharding: Reduce behavior is hard-coded ReduceAll and no additional parameter. A tensor should only be distributed across on dimension (the data/batch dimension).
• Minimal Adaption: This lead to a minimal adaption. This PR only adds to DimWit's core. No function like sum is changed, e.g., to allow reduction behavior by parameter. The JAX API is also simplified, requiring no naming spec.
• MeshAxis, MeshLabel: I decided on using an own Label and Axis mechanism for device meshes as I see no overlap in concepts for this.
• Express sharding in type: If we shard a tensor it is expressed in its type Batch |@| X:

val shardedX: Tensor2[Batch |@| X, Feature, Float32] = x.shard(mesh, Axis[Batch] -> MeshAxis[X])

• No implicit sharding: We do not support Tensor2[Batch |@| X, Feature, Float32] + Tensor2[Batch, Feature, Float32]; unsure if this is the right design, but it reminds me of broadcasting. Being more restrictive maybe the safer approach with compile-time errors being quick in reporting mistakes.
• 2 Examples: This PR contains two examples. We should probably remove to a single example before merging. Also adding AGENTS.md is TODO. I tested this PR also in deepwit on GPT with distributed training: https://github.com/benikm91/deepwit/tree/distributed-gpt
• Mesh alias: Mesh mirrors shape in its implementation, also providing Mesh1 aliases, etc.
• build.sbt needed a change that we maybe should remove before merging.