Struct FullyShardedDataParallel 

pub struct FullyShardedDataParallel<M: Module> { /* private fields */ }

Fully Sharded Data Parallel wrapper for memory-efficient distributed training.

FSDP shards model parameters across devices, gathering them only when needed for computation and sharding them again afterward.

Implementations

impl<M: Module> FullyShardedDataParallel<M>

pub fn new(module: M, process_group: ProcessGroup) -> Self

Creates a new FSDP wrapper.

pub fn sharding_strategy(self, strategy: ShardingStrategy) -> Self

Builder: set sharding strategy.

pub fn cpu_offload(self, offload: CPUOffload) -> Self

Builder: set CPU offload configuration.

pub fn mixed_precision(self, enabled: bool) -> Self

Builder: enable mixed precision.

pub fn module(&self) -> &M

Returns reference to wrapped module.

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

Returns mutable reference to wrapped module.

pub fn process_group(&self) -> &ProcessGroup

Returns the process group.

pub fn strategy(&self) -> ShardingStrategy

Returns the sharding strategy.

pub fn gather_parameters(&mut self)

Gathers all parameter shards before forward pass.

pub fn reshard_parameters(&mut self)

Shards parameters after forward/backward pass.

pub fn sync_gradients(&self)

Synchronizes gradients across all ranks.

pub fn clip_grad_norm(&self, max_norm: f32) -> f32

Clips gradients by global norm.

pub fn memory_estimate(&self) -> FSDPMemoryStats

Estimates memory usage with different sharding strategies.

Trait Implementations

impl<M: Module> Module for FullyShardedDataParallel<M>

fn forward(&self, input: &Variable) -> Variable

Performs the forward pass.

fn parameters(&self) -> Vec<Parameter>

Returns all parameters of this module.

fn train(&mut self)

Sets the module to training mode.

fn eval(&mut self)

Sets the module to evaluation mode.

fn is_training(&self) -> bool

Returns whether the module is in training mode.

fn named_parameters(&self) -> HashMap<String, Parameter>

Returns named parameters of this module.

fn num_parameters(&self) -> usize

Returns the number of trainable parameters.

fn set_training(&mut self, _training: bool)

Sets the training mode.

fn zero_grad(&self)

Zeros all gradients of parameters.

fn name(&self) -> &'static str

Returns the module name for debugging.