pub struct RowParallelLinear { /* private fields */ }Expand description
Row-parallel linear layer
Splits the weight matrix by rows (input dimension) across devices. This is typically used for the second linear layer in MLP blocks.
Implementations§
Source§impl RowParallelLinear
impl RowParallelLinear
pub fn new( in_features: usize, out_features: usize, bias: bool, mp_context: Arc<ModelParallelContext>, ) -> Result<Self>
pub fn forward(&self, input: &DistributedTensor) -> Result<Tensor>
pub fn backward( &mut self, grad_output: &Tensor, input: &DistributedTensor, ) -> Result<DistributedTensor>
Auto Trait Implementations§
impl Freeze for RowParallelLinear
impl !RefUnwindSafe for RowParallelLinear
impl Send for RowParallelLinear
impl Sync for RowParallelLinear
impl Unpin for RowParallelLinear
impl UnsafeUnpin for RowParallelLinear
impl !UnwindSafe for RowParallelLinear
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