Support DiLoCo training.

[ZacharyGarrett]

Description

This is an initial implementation of Distributed Low-Communication training (DiLoCo) as described in https://arxiv.org/abs/2311.08105. DiLoCo is an inner-outer bi-level optimization training strategy that significantly reduces the amount of bandwidth used compared to data-parallel training by syncing between the replicas less frequently.

This implementation adds the drjax package to the pip requirements for bookkeeping and subtle configuration of the jax.vmap's spmd_axis_name argument.

Going forward, one can specify ici_diloco_parallelism or dcn_dilooco_parallelism greater than 1 (the default, which disables) to enable DiLoCo training.

Next steps would include implementing the streaming DiLoCo variant (https://arxiv.org/abs/2501.18512)

Tests

This PR introduces a new tests/diloco_test.py that has a test for numerical correctness of a simple two parameter model.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.