Struct FSDP 

pub struct FSDP<M: Module<T>, T: Float> { /* private fields */ }

Fully Sharded Data Parallel module wrapper.

Wraps an inner Module and shards each parameter across ranks so that each rank only stores 1 / world_size of the full parameter tensor.

Forward pass

Before calling the inner module’s forward(), FSDP all-gathers each shard to reconstruct the full parameter tensor and installs it into the module. The full-parameter tensors are stored in [full_params] so that backward can accumulate gradients on them.

Gradient synchronization

After backward(), call [sync_gradients] to:

1. Read gradients from the full-parameter tensors stored during forward.
2. Reduce-scatter the full gradients so each rank gets only its shard portion of the gradient.
3. Set each shard parameter’s gradient from the reduce-scattered result.

Example

let mut fsdp = FSDP::new(model, backend)?;

loop {
    let output = fsdp.forward(&input)?;
    let loss = criterion.forward(&output, &target)?;
    ferrotorch_core::backward(&loss)?;
    fsdp.sync_gradients()?;
    optimizer.step()?;
    optimizer.zero_grad()?;
}

Implementations

impl<M: Module<T>, T: Float> FSDP<M, T>

pub fn new(module: M, backend: Arc<dyn Backend>) -> FerrotorchResult<Self>

Wrap a module for fully-sharded data-parallel training.

Each parameter is split evenly across world_size ranks. This rank keeps only its shard (the rank-th chunk). The original parameter shapes are recorded for reconstruction during forward.

Panics

Panics if any parameter’s element count is not evenly divisible by world_size.

pub fn module(&self) -> &M

Immutable access to the inner module.

pub fn module_mut(&mut self) -> &mut M

Mutable access to the inner module.

pub fn into_inner(self) -> M

Consume the wrapper and return the inner module.

pub fn backend(&self) -> &Arc<dyn Backend>

The backend used for communication.

pub fn forward(&mut self, input: &Tensor<T>) -> FerrotorchResult<Tensor<T>>

Reconstruct full parameters from shards across all ranks and run the inner module’s forward pass.

The all-gathered full-parameter tensors are stored in self.full_params so their gradients can be read after backward.

pub fn sync_gradients(&mut self) -> FerrotorchResult<()>

Reduce-scatter gradients from the full-parameter tensors stored during forward, then set each shard parameter’s gradient.

Call this after backward() and before optimizer.step().

How it works

1. For each parameter, read the gradient from the full-param tensor that was used during forward (stored in self.full_params).
2. Reduce-scatter the full gradient across ranks (mean reduction) so each rank gets only its shard portion.
3. Set the shard parameter’s .grad() to the reduce-scattered result.

Using reduce-scatter (not allreduce) is correct for FSDP because each rank only needs its own shard of the gradient to update its shard of the parameter.

pub fn update_shards(&mut self, flat_data: &[T]) -> FerrotorchResult<()>

Update shard parameters from a flat data slice.

This is used by optimizers that produce a flat parameter buffer. The slice must have exactly the number of elements expected for this rank’s shards.