Struct DistributedCudaTrainer 

pub struct DistributedCudaTrainer { /* private fields */ }

Distributed CUDA trainer for data-parallel pretraining.

Wraps a single-GPU CudaTransformerTrainer with communication and gradient averaging logic. The actual CUDA operations remain in the underlying trainer — this layer only handles:

1. Running forward+backward in accumulate-only mode
2. Pre-averaging local gradients
3. AllReducing gradients across workers
4. Applying averaged gradients via the underlying trainer

Safety

C-STREAMSYNC-001 applies: stream.synchronize() before all D2H transfers is handled by the underlying CudaTransformerTrainer.

Implementations

impl DistributedCudaTrainer

pub fn new( trainer: CudaTransformerTrainer, comm: DistributedComm, dist_config: DistributedTrainConfig, ) -> Self

Create a new distributed trainer.

Ensures the underlying trainer has gradient accumulation buffers (required for DDP even with accumulation_steps=1, since gradients must be downloaded to CPU for AllReduce).

Arguments

• trainer - Pre-initialized single-GPU trainer (with ensure_grad_accum called)
• comm - Communication backend (local channels or TCP)
• dist_config - Distributed training configuration

pub fn train_batch(&mut self, batch: &LMBatch) -> f32

DDP training step: forward+backward → AllReduce → optimizer.

1. Local forward+backward (accumulate_only=true → grads to CPU accum)
2. Pre-average local gradients
3. AllReduce all gradient components across workers
4. Apply averaged gradients (upload to GPU + optimizer step)

Returns average loss for this worker’s batch.

pub fn dist_config(&self) -> &DistributedTrainConfig

Get the distributed configuration.

pub fn step(&self) -> usize

Get the current step.

pub fn trainer(&self) -> &CudaTransformerTrainer

Get a reference to the underlying trainer.

pub fn trainer_mut(&mut self) -> &mut CudaTransformerTrainer

Get a mutable reference to the underlying trainer.

pub fn is_coordinator(&self) -> bool

Check if this worker is the coordinator (rank 0).

pub fn world_size(&self) -> usize

Get world size.

pub fn rank(&self) -> usize

Get rank.

pub fn reached_max_steps(&self) -> bool

Check if max_steps has been reached.