Enum ShardedSafeTensors

pub enum ShardedSafeTensors {
    Sharded {
        b: MmapedSafetensors,
        make_dummy_regexes: Option<Arc<Vec<Regex>>>,
        predicate: Arc<dyn Fn(String) -> bool + Send + Sync + 'static>,
    },
    SimpleBackend(Box<dyn SimpleBackend + 'static>),
}

Variants

Sharded

Fields

b: MmapedSafetensors

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

predicate: Arc<dyn Fn(String) -> bool + Send + Sync + 'static>

SimpleBackend(Box<dyn SimpleBackend + 'static>)

Implementations

impl ShardedSafeTensors

pub unsafe fn sharded<P: AsRef<Path>>( paths: &[P], dtype: DType, dev: &Device, make_dummy_regexes: Option<Arc<Vec<Regex>>>, predicate: Arc<dyn Fn(String) -> bool + Send + Sync + 'static>, ) -> Result<ShardedVarBuilder>

Initializes a VarBuilder that retrieves tensors stored in a collection of safetensors files and make them usable in a sharded way.

• If regexes is specified, this will be used in make_dummy_predicate based on .any
• Only include keys for which predicate evaluates to true.

Safety

The unsafe is inherited from [memmap2::MmapOptions].

impl ShardedSafeTensors

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

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