Struct Zero3CpuOffloadManager 

pub struct Zero3CpuOffloadManager { /* private fields */ }

Main ZeRO-3 CPU offload manager that orchestrates all components

This is the primary interface for ZeRO-3 operations, providing a unified API that coordinates between all the specialized modules. It maintains the same interface as the original monolithic implementation for backward compatibility.

Implementations

impl Zero3CpuOffloadManager

pub fn new( config: Zero3CpuOffloadConfig, process_group: Arc<ProcessGroup>, model_parameters: &ConfigModelParameters, ) -> TorshResult<Self>

Create a new ZeRO-3 CPU offload manager

Initializes all component systems and establishes distributed coordination. The manager will automatically partition parameters, gradients, and optimizer states according to the ZeRO-3 algorithm.

Arguments

• config - Configuration for ZeRO-3 behavior and memory management
• process_group - Distributed process group for coordination
• model_parameters - Description of model parameters to be managed

Returns

Returns a configured ZeRO-3 manager ready for training operations.

pub async fn forward_pass( &mut self, input: &Tensor<f32>, layer_names: &[String], ) -> TorshResult<Tensor<f32>>

Execute forward pass with ZeRO-3 CPU offloading

Processes each layer with intelligent parameter management:

1. Prefetches parameters for upcoming layers
2. Ensures current layer parameters are on GPU
3. Executes layer computation
4. Optionally offloads parameters back to CPU
5. Performs memory optimization as needed

Arguments

• input - Input tensor for the forward pass
• layer_names - Ordered list of layer names to execute

Returns

Returns the output tensor after processing all layers.

pub async fn backward_pass( &mut self, grad_output: &Tensor<f32>, layer_names: &[String], ) -> TorshResult<()>

Execute backward pass with ZeRO-3 CPU offloading

Processes layers in reverse order for gradient computation:

1. Ensures parameters are available for gradient computation
2. Computes gradients for each layer
3. Partitions and manages gradients according to ZeRO-3
4. Performs all-reduce synchronization across ranks

Arguments

• grad_output - Gradient tensor from the loss function
• layer_names - Ordered list of layer names (processed in reverse)

Returns

Returns Ok(()) when backward pass completes successfully.

pub async fn optimizer_step(&mut self, learning_rate: f32) -> TorshResult<()>

Update optimizer states and parameters with ZeRO-3 partitioning

Performs optimizer step with intelligent state management:

1. Gathers partitioned gradients for owned parameters
2. Fetches optimizer states from CPU if needed
3. Computes parameter updates using optimizer algorithm
4. Updates parameters and stores back to appropriate location
5. Broadcasts updates to all ranks that need them

Arguments

• learning_rate - Learning rate for parameter updates

Returns

Returns Ok(()) when optimizer step completes successfully.

pub fn get_performance_stats(&self) -> Zero3PerformanceStats

Get comprehensive performance statistics

Returns detailed performance metrics including timing, throughput, memory usage, and efficiency measurements.

pub fn get_memory_stats(&self) -> Zero3MemoryStats

Get memory usage statistics

Returns current memory usage across CPU and GPU, including parameter distribution and compression effectiveness.

pub async fn force_memory_optimization(&self) -> TorshResult<()>

Force immediate memory optimization

Triggers aggressive memory optimization regardless of current pressure. Useful for cleaning up before checkpointing or when memory is critically low.

pub fn get_prefetch_status(&self) -> PrefetchQueueStatus

Get prefetch scheduler status

Returns information about current prefetch operations and queue status.

pub async fn adapt_performance(&self) -> TorshResult<()>

Adapt system performance based on runtime metrics

Analyzes recent performance and adjusts prefetch strategies, memory management policies, and other adaptive parameters.

pub async fn reset_state(&self) -> TorshResult<()>

Clear all caches and reset state

Useful for testing or when switching between different models.