pub struct MicrobatchManager { /* private fields */ }Expand description
Manages microbatches across pipeline stages
Implementations§
Source§impl MicrobatchManager
impl MicrobatchManager
pub fn new(num_microbatches: usize, checkpoint_activations: bool) -> Self
Sourcepub fn get(&self, id: usize) -> Result<&Microbatch>
pub fn get(&self, id: usize) -> Result<&Microbatch>
Get microbatch by ID
Sourcepub fn get_mut(&mut self, id: usize) -> Result<&mut Microbatch>
pub fn get_mut(&mut self, id: usize) -> Result<&mut Microbatch>
Get mutable microbatch
Sourcepub fn enqueue_forward(&mut self, mb_id: usize)
pub fn enqueue_forward(&mut self, mb_id: usize)
Add microbatch to forward queue
Sourcepub fn enqueue_backward(&mut self, mb_id: usize)
pub fn enqueue_backward(&mut self, mb_id: usize)
Add microbatch to backward queue
Sourcepub fn dequeue_forward(&mut self) -> Option<usize>
pub fn dequeue_forward(&mut self) -> Option<usize>
Get next forward microbatch
Sourcepub fn dequeue_backward(&mut self) -> Option<usize>
pub fn dequeue_backward(&mut self) -> Option<usize>
Get next backward microbatch
Sourcepub fn maybe_clear_activation(&mut self, mb_id: usize) -> Result<()>
pub fn maybe_clear_activation(&mut self, mb_id: usize) -> Result<()>
Clear activation if checkpointing is enabled
Sourcepub fn maybe_recompute_activation(
&mut self,
mb_id: usize,
stage: &PipelineStage,
) -> Result<()>
pub fn maybe_recompute_activation( &mut self, mb_id: usize, stage: &PipelineStage, ) -> Result<()>
Recompute activation if needed
Auto Trait Implementations§
impl Freeze for MicrobatchManager
impl RefUnwindSafe for MicrobatchManager
impl Send for MicrobatchManager
impl Sync for MicrobatchManager
impl Unpin for MicrobatchManager
impl UnsafeUnpin for MicrobatchManager
impl UnwindSafe for MicrobatchManager
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