pub struct SequenceParallelism { /* private fields */ }Expand description
Sequence parallelism coordinator
Implementations§
Source§impl SequenceParallelism
impl SequenceParallelism
Sourcepub fn new(
config: SequenceParallelismConfig,
global_rank: usize,
world_size: usize,
sequence_group: Arc<dyn ProcessGroup>,
) -> Result<Self>
pub fn new( config: SequenceParallelismConfig, global_rank: usize, world_size: usize, sequence_group: Arc<dyn ProcessGroup>, ) -> Result<Self>
Create a new sequence parallelism coordinator
Sourcepub fn split_sequence(
&mut self,
total_sequence_length: usize,
) -> Result<Vec<SequenceChunk>>
pub fn split_sequence( &mut self, total_sequence_length: usize, ) -> Result<Vec<SequenceChunk>>
Split a sequence across multiple devices
Sourcepub fn forward_chunk(
&self,
chunk_id: usize,
input: &Tensor,
_attention_mask: Option<&Tensor>,
) -> Result<Tensor>
pub fn forward_chunk( &self, chunk_id: usize, input: &Tensor, _attention_mask: Option<&Tensor>, ) -> Result<Tensor>
Process forward pass for a sequence chunk
Sourcepub fn synchronize_gradients(
&self,
gradients: &mut HashMap<String, Tensor>,
) -> Result<()>
pub fn synchronize_gradients( &self, gradients: &mut HashMap<String, Tensor>, ) -> Result<()>
Synchronize gradients across sequence chunks
Sourcepub fn update_memory_usage(
&self,
chunk_id: usize,
memory_usage: u64,
) -> Result<()>
pub fn update_memory_usage( &self, chunk_id: usize, memory_usage: u64, ) -> Result<()>
Update memory usage statistics
Sourcepub fn get_statistics(&self) -> SequenceParallelismStats
pub fn get_statistics(&self) -> SequenceParallelismStats
Get sequence parallelism statistics
Sourcepub fn local_chunks(&self) -> &[usize]
pub fn local_chunks(&self) -> &[usize]
Get local chunk IDs
Sourcepub fn get_chunk(&self, chunk_id: usize) -> Option<&SequenceChunk>
pub fn get_chunk(&self, chunk_id: usize) -> Option<&SequenceChunk>
Get chunk information
Sourcepub fn config(&self) -> &SequenceParallelismConfig
pub fn config(&self) -> &SequenceParallelismConfig
Get configuration
Auto Trait Implementations§
impl Freeze for SequenceParallelism
impl !RefUnwindSafe for SequenceParallelism
impl Send for SequenceParallelism
impl Sync for SequenceParallelism
impl Unpin for SequenceParallelism
impl UnsafeUnpin for SequenceParallelism
impl !UnwindSafe for SequenceParallelism
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more
Source§impl<T> Instrument for T
impl<T> Instrument for T
Source§fn instrument(self, span: Span) -> Instrumented<Self>
fn instrument(self, span: Span) -> Instrumented<Self>
Source§fn in_current_span(self) -> Instrumented<Self>
fn in_current_span(self) -> Instrumented<Self>
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
Converts
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
Converts
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more