Struct TensorParallel 

pub struct TensorParallel { /* private fields */ }

Tensor parallel wrapper for modules

Implementations

impl TensorParallel

pub fn new( module: Box<dyn Module>, tp_group: Arc<ProcessGroup>, config: TensorParallelConfig, layer_info: TensorParallelLayer, ) -> TorshResult<Self>

Create a new tensor parallel wrapper

pub fn tp_rank(&self) -> usize

Get tensor parallel rank

pub fn tp_world_size(&self) -> usize

Get tensor parallel world size

pub fn get_shard_info(&self, param_name: &str) -> Option<&ShardInfo>

Get sharding information for a parameter

pub fn uses_sequence_parallel(&self) -> bool

Check if layer uses sequence parallelism

pub fn memory_stats(&self) -> TensorParallelStats

Get memory usage statistics

Trait Implementations

impl Module for TensorParallel

fn forward(&self, input: &Tensor) -> Result<Tensor>

Forward pass through the module

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

Get all parameters in the module (non-recursive)

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

Get named parameters (non-recursive)

fn training(&self) -> bool

Check if in training mode

fn train(&mut self)

Set training mode

fn eval(&mut self)

Set evaluation mode

fn to_device(&mut self, device: DeviceType) -> Result<()>

Move module to device

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

Get all parameters recursively including submodules

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

Get all named parameters recursively with module prefixes

fn set_training(&mut self, _training: bool)

Set training mode (internal implementation)

fn load_state_dict( &mut self, state_dict: &HashMap<String, Tensor>, strict: bool, ) -> Result<(), TorshError>

Load state dictionary into the module

fn load_state_dict_strict( &mut self, state_dict: &HashMap<String, Tensor>, ) -> Result<(), TorshError>

Load state dictionary with default strict=true

fn state_dict(&self) -> HashMap<String, Tensor>

Save state dictionary from the module

fn name(&self) -> Option<&str>

Get the module name (optional, for debugging and serialization)

fn buffers(&self) -> Vec<Arc<RwLock<RawRwLock, Tensor>>>

Get all buffers (non-trainable parameters)

fn named_buffers(&self) -> HashMap<String, Arc<RwLock<RawRwLock, Tensor>>>

Get named buffers

fn children(&self) -> Vec<&dyn Module>

Get all direct child modules

fn named_children(&self) -> Vec<(String, &dyn Module)>

Get all direct child modules with names

fn modules(&self) -> Vec<&dyn Module>

where Self: Sized,

Get all modules recursively (depth-first traversal)

fn named_modules(&self) -> Vec<(String, &dyn Module)>

where Self: Sized,

Get all modules recursively with hierarchical names

fn zero_grad(&mut self)

Zero all gradients recursively

fn num_parameters(&self) -> usize

Count total number of parameters

fn num_trainable_parameters(&self) -> usize

Count trainable parameters

fn memory_usage(&self) -> usize

Get memory usage estimate in bytes

fn freeze(&mut self)

Freeze all parameters (set requires_grad = false)

fn unfreeze(&mut self)

Unfreeze all parameters (set requires_grad = true)

fn extra_repr(&self) -> String

Get string representation

fn register_hook( &mut self, _hook_type: HookType, _callback: Box<dyn Fn(&dyn Module, &Tensor, Option<&Tensor>) -> Result<(), TorshError> + Send + Sync>, ) -> Option<HookHandle>

Register a hook for this module (default implementation does nothing)

fn remove_hook(&mut self, _hook_type: HookType, _handle: HookHandle) -> bool

Remove a hook by handle (default implementation does nothing)

fn execute_hooks( &self, _hook_type: HookType, _input: &Tensor, _output: Option<&Tensor>, ) -> Result<(), TorshError>

Execute hooks of a specific type (default implementation does nothing)

fn forward_with_hooks(&self, input: &Tensor) -> Result<Tensor, TorshError>

Forward pass with hooks support

fn has_hooks(&self, _hook_type: HookType) -> bool

Check if module has hooks registered

fn call(&self, input: &Tensor) -> Result<Tensor, TorshError>

Convenient method to call forward and handle common patterns

fn apply(&self, input: &Tensor) -> Result<Tensor, TorshError>

Apply the module to input (alias for forward)

fn has_parameters(&self) -> bool

Check if the module has any parameters

fn has_children(&self) -> bool

Check if the module has any child modules

fn parameter_count(&self) -> usize

Get parameter count (convenience method)

fn trainable_parameter_count(&self) -> usize

Get trainable parameter count (convenience method)

fn memory_usage_mb(&self) -> f64

Get memory usage in MB (convenience method)

fn toggle_training(&mut self)

Toggle training mode (convenience method)

fn eval_mode(&self) -> bool

Check if module is in evaluation mode

fn sequential_forward( modules: &[&dyn Module], input: Tensor, ) -> Result<Tensor, TorshError>

where Self: Sized,

Sequential forward pass through multiple modules

fn batch_forward(&self, inputs: &[Tensor]) -> Result<Vec<Tensor>, TorshError>

Apply module multiple times with different inputs (batch processing)

fn conditional_forward( &self, input: &Tensor, condition: bool, ) -> Result<Tensor, TorshError>

Forward with condition - only apply if condition is true

fn residual_forward(&self, input: &Tensor) -> Result<Tensor, TorshError>

Forward with residual connection

fn module_info(&self) -> ModuleInfo

Get detailed module information for debugging

fn check_training_readiness(&self) -> Result<(), TorshError>

Check if module is ready for training

fn parameter_names_matching(&self, pattern: &str) -> Vec<String>

Get parameter names matching a pattern

fn parameters_by_type(&self, param_type: &str) -> HashMap<String, Parameter>

Get parameters by layer type (e.g., “weight”, “bias”)

fn clone_parameters(&self) -> HashMap<String, Tensor>

Clone module parameters (for creating copies or checkpoints)

fn diagnose(&self) -> ModuleDiagnostics

Quick diagnostic check of module health