huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 19:18:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Speedup ShardedSafeTensors to load Tensors with default hints (#1384)

Browse Source 
* Speedup ShardedSafeTensors to load Tensors with default hints

* Tweaks.

---------

Co-authored-by: Laurent <laurent.mazare@gmail.com>
This commit is contained in:
YiiSh
2023-12-14 22:08:56 +08:00
committed by GitHub
parent 7be982f6f7
commit e60f9b5dfc
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-nn/src/var_builder.rs
View File

@ -535,12 +535,18 @@ impl Backend for ShardedSafeTensors {
fn get( fn get(
&self, &self,
_target_shape: Shape, // The size is not checked for ShardedTensors target_shape: Shape, // The size is only checked when the world size is 1.
path: &str, path: &str,
h: Self::Hints, h: Self::Hints,
dtype: DType, dtype: DType,
dev: &Device, dev: &Device,
) -> Result<Tensor> { ) -> Result<Tensor> {
if h.world_size == 1 {
// There is no sharding to be applied here so we use the default backend to speed
// things up.
return SimpleBackend::get(&self.0, target_shape, path, Default::default(), dtype, dev);
}
let Shard { let Shard {
dim, dim,
rank, rank,