Add MultiGroupDistributedDataParallel

[AkshitaB]

Summary

Adds a data-parallel module wrapper, MultiGroupDistributedDataParallel, under olmo_core.nn.parallel. It all-reduces gradients across one or more process groups.

The distinguishing feature vs. torch.nn.parallel.DistributedDataParallel: different parameters can be reduced over different process groups, via a param_process_group_fn(name, param) -> process_group mapping. This is needed when subsets of a model are replicated over different device meshes — e.g. expert weights under expert parallelism, which must be data-parallel-reduced only over the ranks holding the same experts, while the rest of the model reduces over the full DP group.

Mechanics

• Bucket views: gradients are accumulated into flat, pre-allocated buffers; each parameter's .grad is a view into its bucket slice. This enables one large all-reduce per bucket instead of many tiny ones.
• Per-group buckets: a bucket never mixes process groups, so each bucket is reduced over its own group with its own world size.
• Overlap: bucketed all-reduces launch during the backward pass as each bucket fills; finalize_grad_reduce() is called after loss.backward() to wait for completion.
• Optional fp32: gradients can be accumulated and/or reduced in fp32 for numerical stability with low-precision parameters.

ddp = MultiGroupDistributedDataParallel(model)
out = ddp(inputs)
out.loss.backward()
ddp.finalize_grad_reduce()  # waits for the overlapped all-reduces
optimizer.step()

Notes

• This lands the primitive ahead of its consumers; nothing in-tree wires it in yet.
• The wrapper includes opt-in, duck-typed hooks (getattr(module, ...)) for syncing fp8/mxfp8 gradient stores; these are no-ops for standard modules.

Tests

New gloo/nccl distributed test (src/test/nn/parallel/distributed_test.py):

• test_grad_parity — DDP-averaged gradients match a manual all-reduce reference.
• test_no_sync_accumulation — no_sync() defers the reduction and correctly accumulates gradients across micro-batches before a single reduce.

pytest src/test/nn/parallel/ passes on gloo (nccl auto-skips without a GPU); make checks (isort/black/ruff/mypy) clean.

🤖 Generated with Claude Code