Enum ShardedSafeTensors

pub enum ShardedSafeTensors<'a> {
    Sharded {
        b: MmapedSafetensors,
        make_dummy_regexes: Option<Arc<Vec<Regex>>>,
    },
    SimpleBackend(Box<dyn SimpleBackend + 'a>),
}

Variants

Sharded

Fields

b: MmapedSafetensors

make_dummy_regexes: Option<Arc<Vec<Regex>>>

SimpleBackend(Box<dyn SimpleBackend + 'a>)

Implementations

impl<'a> ShardedSafeTensors<'a>

pub unsafe fn sharded<P: AsRef<Path>>( paths: &[P], dtype: DType, dev: &Device, make_dummy_regexes: Option<Arc<Vec<Regex>>>, ) -> Result<ShardedVarBuilder<'a>>

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].

pub fn wrap( backend: Box<dyn SimpleBackend + 'a>, dtype: DType, dev: Device, ) -> ShardedVarBuilder<'a>

Trait Implementations

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

fn get( &self, target_shape: Shape, path: &str, h: Self::Hints, dtype: DType, dev: &Device, ) -> Result<Tensor>

Retrieve a tensor with some target shape.

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