pub struct ShardedSafeTensors(/* private fields */);Implementations§
source§impl ShardedSafeTensors
impl ShardedSafeTensors
sourcepub unsafe fn var_builder<P: AsRef<Path>>(
paths: &[P],
dtype: DType,
dev: &Device,
) -> Result<ShardedVarBuilder<'static>>
pub unsafe fn var_builder<P: AsRef<Path>>( paths: &[P], dtype: DType, dev: &Device, ) -> Result<ShardedVarBuilder<'static>>
Initializes a VarBuilder that retrieves tensors stored in a collection of safetensors
files and make them usable in a sharded way.
§Safety
The unsafe is inherited from [memmap2::MmapOptions].
Trait Implementations§
source§impl Backend for ShardedSafeTensors
impl Backend for ShardedSafeTensors
Get part of a tensor, typically used to do Tensor Parallelism sharding.
If the tensor is of size (1024, 1024).
dim corresponds to the dimension to slice into
rank is the rank of the current process
world_size is the total number of ranks in the process group
get_sharded("tensor", 0, 0, 2) means tensor.i((..512))
get_sharded("tensor", 0, 1, 2) means tensor.i((512..))
get_sharded("tensor", 1, 0, 2) means tensor.i((.., ..512))
type Hints = Shard
source§fn get(
&self,
target_shape: Shape,
path: &str,
h: Self::Hints,
dtype: DType,
dev: &Device,
) -> Result<Tensor>
fn get( &self, target_shape: Shape, path: &str, h: Self::Hints, dtype: DType, dev: &Device, ) -> Result<Tensor>
Retrieve a tensor with some target shape.
source§fn get_unchecked(
&self,
_name: &str,
_dtype: DType,
_dev: &Device,
) -> Result<Tensor>
fn get_unchecked( &self, _name: &str, _dtype: DType, _dev: &Device, ) -> Result<Tensor>
Retrieve a tensor based on the name.
fn contains_tensor(&self, name: &str) -> bool
Auto Trait Implementations§
impl Freeze for ShardedSafeTensors
impl RefUnwindSafe for ShardedSafeTensors
impl Send for ShardedSafeTensors
impl Sync for ShardedSafeTensors
impl Unpin for ShardedSafeTensors
impl UnwindSafe for ShardedSafeTensors
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> 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