Skip to main content

Module distributed

Module distributed 

Source
Expand description

Multi-rank gradient averaging for Tier-2 training.

Each process/rank:

  1. writes local gradients into a shared sync_dir
  2. rank 0 averages them and writes the result
  3. every rank loads the averaged gradients before optimizer.step

This is slower than GPU collectives, but it is a working Python-free multi-process path without pulling in MPI/NCCL.

Structs§

FileGradSyncConfig
Settings for averaging gradients across processes through files in a shared directory (a prototype for multi-process data parallelism).

Functions§

average_grads_file_sync
Average grads across ranks and overwrite them in place.
average_host_grads
Average host-side parameter vectors (unit-testable without GPU).
shard_indices
Select sample indices owned by rank in a world of world_size.