feat(server): Use safetensors

Co-authored-by: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>

Cargo.toml

@@ -4,6 +4,9 @@ members = [
     "router/client",
     "launcher"
 ]
+exclude = [
+    "server/safetensors",
+]
 
 [profile.release]
 debug = 1

Dockerfile

@@ -36,10 +36,10 @@ ENV LANG=C.UTF-8 \
 
 SHELL ["/bin/bash", "-c"]
 
-RUN apt-get update && apt-get install -y unzip wget libssl-dev && rm -rf /var/lib/apt/lists/*
+RUN apt-get update && apt-get install -y unzip curl libssl-dev && rm -rf /var/lib/apt/lists/*
 
 RUN cd ~ && \
-    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    curl -L -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
     chmod +x Miniconda3-latest-Linux-x86_64.sh && \
     bash ./Miniconda3-latest-Linux-x86_64.sh -bf -p /opt/miniconda && \
     conda create -n text-generation python=3.9 -y
@@ -54,13 +54,22 @@ RUN cd server && make install-torch
 # Install specific version of transformers
 RUN cd server && make install-transformers
 
+# Install specific version of safetensors
+# FIXME: This is a temporary fix while we wait for a new release
+RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
+ENV PATH="/root/.cargo/bin:${PATH}"
+RUN cd server && make install-safetensors
+
 # Install server
+COPY proto proto
 COPY server server
 RUN cd server && \
+    make gen-server && \
     /opt/miniconda/envs/text-generation/bin/pip install . --no-cache-dir
 
 # Install router
 COPY --from=router-builder /usr/local/cargo/bin/text-generation-router /usr/local/bin/text-generation-router
+# Install launcher
 COPY --from=launcher-builder /usr/local/cargo/bin/text-generation-launcher /usr/local/bin/text-generation-launcher
 
-CMD text-generation-launcher --model-name $MODEL_NAME --num-shard $NUM_GPUS --shard-directory $MODEL_BASE_PATH
+CMD HUGGINGFACE_HUB_CACHE=$MODEL_BASE_PATH text-generation-launcher --model-name $MODEL_NAME --num-shard $NUM_GPUS
\ No newline at end of file

Makefile

 install-server:
-	cd server && make pip-install
+	cd server && make install
 
 install-router:
 	cd router && cargo install --path .
@@ -7,13 +7,16 @@ install-router:
 install-launcher:
 	cd launcher && cargo install --path .
 
-install:
-	make install-server
-	make install-router
-	make install-launcher
+install: install-server install-router install-launcher
+
+server-dev:
+	cd server && make run-dev
+
+router-dev:
+	cd router && cargo run
 
 run-bloom-560m:
-	text-generation-launcher --model-name bigscience/bloom-560m --shard-directory /tmp/models --num-shard 2
+	text-generation-launcher --model-name bigscience/bloom-560m --num-shard 2
 
 run-bloom:
-	text-generation-launcher --model-name bigscience/bloom --shard-directory /tmp/models --num-shard 8
+	text-generation-launcher --model-name bigscience/bloom --num-shard 8

README.md

@@ -41,11 +41,21 @@ make run-bloom-560m
 
 ```shell
 curl 127.0.0.1:3000/generate \
+    -v \
     -X POST \
     -d '{"inputs":"Testing API","parameters":{"max_new_tokens":9}}' \
     -H 'Content-Type: application/json'
 ```
 
+## Develop
+
+```shell
+make server-dev
+make router-dev
+```
+
 ## TODO:
 
 - [ ] Add tests for the `server/model` logic
+- [ ] Backport custom CUDA kernels to Transformers
+- [ ] Install safetensors with pip
\ No newline at end of file

launcher/src/main.rs

 use clap::Parser;
+use std::env;
 use std::io::{BufRead, BufReader, Read};
 use std::path::Path;
 use std::process::ExitCode;
@@ -20,8 +21,6 @@ struct Args {
     model_name: String,
     #[clap(long, env)]
     num_shard: Option<usize>,
-    #[clap(long, env)]
-    shard_directory: Option<String>,
     #[clap(default_value = "128", long, env)]
     max_concurrent_requests: usize,
     #[clap(default_value = "1000", long, env)]
@@ -47,7 +46,6 @@ fn main() -> ExitCode {
     let Args {
         model_name,
         num_shard,
-        shard_directory,
         max_concurrent_requests,
         max_input_length,
         max_batch_size,
@@ -82,7 +80,6 @@ fn main() -> ExitCode {
     for rank in 0..num_shard {
         let model_name = model_name.clone();
         let uds_path = shard_uds_path.clone();
-        let shard_directory = shard_directory.clone();
         let master_addr = master_addr.clone();
         let status_sender = status_sender.clone();
         let shutdown = shutdown.clone();
@@ -91,7 +88,6 @@ fn main() -> ExitCode {
             shard_manager(
                 model_name,
                 uds_path,
-                shard_directory,
                 rank,
                 num_shard,
                 master_addr,
@@ -237,7 +233,6 @@ enum ShardStatus {
 fn shard_manager(
     model_name: String,
     uds_path: String,
-    shard_directory: Option<String>,
     rank: usize,
     world_size: usize,
     master_addr: String,
@@ -265,10 +260,27 @@ fn shard_manager(
         shard_argv.push("--sharded".to_string());
     }
 
-    if let Some(shard_directory) = shard_directory {
-        shard_argv.push("--shard-directory".to_string());
-        shard_argv.push(shard_directory);
-    }
+    let mut env = vec![
+        ("RANK".parse().unwrap(), rank.to_string().parse().unwrap()),
+        (
+            "WORLD_SIZE".parse().unwrap(),
+            world_size.to_string().parse().unwrap(),
+        ),
+        ("MASTER_ADDR".parse().unwrap(), master_addr.parse().unwrap()),
+        (
+            "MASTER_PORT".parse().unwrap(),
+            master_port.to_string().parse().unwrap(),
+        ),
+    ];
+
+    // If the HUGGINGFACE_HUB_CACHE env var is set, pass it to the shard
+    // Useful when running inside a docker container
+    if let Ok(huggingface_hub_cache) = env::var("HUGGINGFACE_HUB_CACHE") {
+        env.push((
+            "HUGGINGFACE_HUB_CACHE".parse().unwrap(),
+            huggingface_hub_cache.parse().unwrap(),
+        ));
+    };
 
     // Start process
     tracing::info!("Starting shard {}", rank);
@@ -280,18 +292,7 @@ fn shard_manager(
             // Needed for the shutdown procedure
             setpgid: true,
             // NCCL env vars
-            env: Some(vec![
-                ("RANK".parse().unwrap(), rank.to_string().parse().unwrap()),
-                (
-                    "WORLD_SIZE".parse().unwrap(),
-                    world_size.to_string().parse().unwrap(),
-                ),
-                ("MASTER_ADDR".parse().unwrap(), master_addr.parse().unwrap()),
-                (
-                    "MASTER_PORT".parse().unwrap(),
-                    master_port.to_string().parse().unwrap(),
-                ),
-            ]),
+            env: Some(env),
             ..Default::default()
         },
     ) {
@@ -312,6 +313,7 @@ fn shard_manager(
 
     let mut ready = false;
     let start_time = Instant::now();
+    let mut wait_time = Instant::now();
     loop {
         // Process exited
         if p.poll().is_some() {
@@ -337,9 +339,12 @@ fn shard_manager(
             status_sender.send(ShardStatus::Ready).unwrap();
             ready = true;
         } else if !ready {
+            if wait_time.elapsed() > Duration::from_secs(10) {
                 tracing::info!("Waiting for shard {} to be ready...", rank);
+                wait_time = Instant::now();
+            }
         }
-        sleep(Duration::from_secs(5));
+        sleep(Duration::from_millis(100));
     }
 }
 

server/.gitignore

@@ -153,3 +153,6 @@ dmypy.json
 
 # Cython debug symbols
 cython_debug/
+
+transformers
+safetensors
\ No newline at end of file

server/Makefile

 gen-server:
+	# Compile protos
+	pip install grpcio-tools==1.49.1 --no-cache-dir
 	mkdir bloom_inference/pb || true
 	python -m grpc_tools.protoc -I../proto --python_out=bloom_inference/pb --grpc_python_out=bloom_inference/pb ../proto/generate.proto
 	find bloom_inference/pb/ -type f -name "*.py" -print0 -exec sed -i -e 's/^\(import.*pb2\)/from . \1/g' {} \;
@@ -7,25 +9,31 @@ gen-server:
 install-transformers:
 	# Install specific version of transformers
 	rm transformers || true
-	wget https://github.com/huggingface/transformers/archive/46d37bece7d3ffdef97b1ee4a3170c0a0627d921.zip
+	rm transformers-46d37bece7d3ffdef97b1ee4a3170c0a0627d921 || true
+	curl -L -O https://github.com/huggingface/transformers/archive/46d37bece7d3ffdef97b1ee4a3170c0a0627d921.zip
 	unzip 46d37bece7d3ffdef97b1ee4a3170c0a0627d921.zip
 	rm 46d37bece7d3ffdef97b1ee4a3170c0a0627d921.zip
 	mv transformers-46d37bece7d3ffdef97b1ee4a3170c0a0627d921 transformers
 	cd transformers && python setup.py install
 
+install-safetensors:
+	# Install specific version of safetensors
+	pip install setuptools_rust
+	rm safetensors || true
+	rm safetensors-634deccbcbad5eaf417935281f8b3be7ebca69c5 || true
+	curl -L -O https://github.com/huggingface/safetensors/archive/634deccbcbad5eaf417935281f8b3be7ebca69c5.zip
+	unzip 634deccbcbad5eaf417935281f8b3be7ebca69c5.zip
+	rm 634deccbcbad5eaf417935281f8b3be7ebca69c5.zip
+	mv safetensors-634deccbcbad5eaf417935281f8b3be7ebca69c5 safetensors
+	cd safetensors/bindings/python && python setup.py develop
+
 install-torch:
 	# Install specific version of torch
 	pip install torch --extra-index-url https://download.pytorch.org/whl/cu116 --no-cache-dir
 
-pip-install:
-	pip install grpcio-tools
-	make gen-server
-	make install-torch
-	make install-transformers
-	pip install .
+install: gen-server install-torch install-transformers install-safetensors
+	pip install pip --upgrade
+	pip install -e . --no-cache-dir
 
-install:
-	poetry install
-	make gen-server
+run-dev:
+	python -m torch.distributed.run --nproc_per_node=2 bloom_inference/cli.py serve bigscience/bloom-560m --sharded
\ No newline at end of file
-	make install-torch
-	make install-transformers

server/README.md

@@ -2,14 +2,14 @@
 
 A Python gRPC server for BLOOM Inference
 
-## Local Install (with poetry)
+## Install
 
 ```shell
 make install
 ```
 
-## Local Install (with pip)
+## Run
 
 ```shell
-make pip-install
+make run-dev
 ```
\ No newline at end of file

server/bloom_inference/cli.py

@@ -2,9 +2,8 @@ import os
 import typer
 
 from pathlib import Path
-from typing import Optional
 
-from bloom_inference import prepare_weights, server
+from bloom_inference import server, utils
 
 app = typer.Typer()
 
@@ -13,13 +12,9 @@ app = typer.Typer()
 def serve(
     model_name: str,
     sharded: bool = False,
-    shard_directory: Optional[Path] = None,
     uds_path: Path = "/tmp/bloom-inference",
 ):
     if sharded:
-        assert (
-            shard_directory is not None
-        ), "shard_directory must be set when sharded is True"
         assert (
             os.getenv("RANK", None) is not None
         ), "RANK must be set when sharded is True"
@@ -33,19 +28,14 @@ def serve(
             os.getenv("MASTER_PORT", None) is not None
         ), "MASTER_PORT must be set when sharded is True"
 
-    server.serve(model_name, sharded, uds_path, shard_directory)
+    server.serve(model_name, sharded, uds_path)
 
 
 @app.command()
-def prepare_weights(
+def download_weights(
     model_name: str,
-    shard_directory: Path,
-    cache_directory: Path,
-    num_shard: int = 1,
 ):
-    prepare_weights.prepare_weights(
-        model_name, cache_directory, shard_directory, num_shard
-    )
+    utils.download_weights(model_name)
 
 
 if __name__ == "__main__":

server/bloom_inference/model.py

@@ -2,19 +2,24 @@ import torch
 import torch.distributed
 
 from dataclasses import dataclass
-from pathlib import Path
 from typing import List, Tuple, Optional, Dict
 
+from accelerate import init_empty_weights
+from safetensors import safe_open
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
-from transformers.modeling_utils import no_init_weights
+from transformers.models.bloom.parallel_layers import (
+    TensorParallelColumnLinear,
+    TensorParallelEmbedding,
+    TensorParallelRowLinear,
+)
 
 from bloom_inference.pb import generate_pb2
-from bloom_inference.prepare_weights import prepare_weights, match_suffix
 from bloom_inference.utils import (
     StoppingCriteria,
     NextTokenChooser,
     initialize_torch_distributed,
-    set_default_dtype,
+    weight_files,
+    download_weights
 )
 
 torch.manual_seed(0)
@@ -63,7 +68,7 @@ class Batch:
             )
             stopping_criterias.append(StoppingCriteria(max_new_tokens=r.max_new_tokens))
 
-        input_ids = tokenizer(inputs, return_tensors="pt", padding=True).to(device)
+        input_ids = tokenizer(inputs, return_tensors="pt", padding=True, pad_to_multiple_of=8).to(device)
         all_input_ids = input_ids["input_ids"].unsqueeze(-1)
 
         return cls(
@@ -369,7 +374,7 @@ class BLOOM:
 
 
 class BLOOMSharded(BLOOM):
-    def __init__(self, model_name: str, shard_directory: Path):
+    def __init__(self, model_name: str):
         super(BLOOM, self).__init__()
         self.process_group, self.rank, self.world_size = initialize_torch_distributed()
         self.master = self.rank == 0
@@ -382,30 +387,11 @@ class BLOOMSharded(BLOOM):
 
         self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
 
-        # shard state_dict
-        if self.master:
-            # TODO @thomasw21 do some caching
-            shard_state_dict_paths = prepare_weights(
-                model_name,
-                shard_directory / "cache",
-                shard_directory,
-                tp_world_size=self.world_size,
-            )
-            shard_state_dict_paths = [
-                str(path.absolute()) for path in shard_state_dict_paths
-            ]
-        else:
-            shard_state_dict_paths = [None] * self.world_size
-
-        torch.distributed.broadcast_object_list(
-            shard_state_dict_paths, src=0, group=self.process_group
-        )
-        shard_state_dict_path = shard_state_dict_paths[self.rank]
-
         config = AutoConfig.from_pretrained(
             model_name, slow_but_exact=False, tp_parallel=True
         )
         config.pad_token_id = 3
+        self.num_heads = config.n_head // self.process_group.size()
 
         # The flag below controls whether to allow TF32 on matmul. This flag defaults to False
         # in PyTorch 1.12 and later.
@@ -414,38 +400,89 @@ class BLOOMSharded(BLOOM):
         # The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
         torch.backends.cudnn.allow_tf32 = True
 
-        with set_default_dtype(dtype):
-            with no_init_weights():
-                # we can probably set the device to `meta` here?
-                model = AutoModelForCausalLM.from_config(config).to(dtype)
+        # Only download weights for small models
+        if self.master and model_name == "bigscience/bloom-560m":
+            download_weights(model_name)
 
         torch.distributed.barrier(group=self.process_group)
-        # print_rank_0(f"Initialized model")
-        state_dict = torch.load(shard_state_dict_path)
-        # TODO @thomasw21: HACK in order to transpose all weight prior
-        for key in state_dict.keys():
-            do_transpose = False
-            if not match_suffix(key, "weight"):
-                continue
-
-            for potential_suffix in [
-                "self_attention.query_key_value.weight",
-                "self_attention.dense.weight",
-                "dense_h_to_4h.weight",
-                "dense_4h_to_h.weight",
-            ]:
-                if match_suffix(key, potential_suffix):
-                    do_transpose = True
-
-            if do_transpose:
-                state_dict[key] = state_dict[key].transpose(1, 0).contiguous()
-
-        model.load_state_dict(state_dict, strict=False)
-        model.tie_weights()
-        self.model = model.to(self.device).eval()
-        self.num_heads = config.n_head // self.process_group.size()
+        filenames = weight_files(model_name)
+
+        with init_empty_weights():
+            model = AutoModelForCausalLM.from_config(config)
+
+        torch.distributed.barrier(group=self.process_group)
+        self.load_weights(
+            model,
+            filenames,
+            device=self.device,
+            rank=self.rank,
+            world_size=self.world_size,
+        )
+        self.model = model.eval().to(dtype)
         torch.distributed.barrier(group=self.process_group)
 
+    @staticmethod
+    def load_weights(
+        model, filenames: List[str], device: torch.device, rank: int, world_size: int
+    ):
+        parameters = dict(model.named_parameters())
+        for file in filenames:
+            with safe_open(file, framework="pt", device=str(device)) as f:
+                for name in f.keys():
+                    full_name = f"transformer.{name}"
+
+                    module_name, param_name = full_name.rsplit(".", 1)
+                    module = model.get_submodule(module_name)
+                    current_tensor = parameters[full_name]
+
+                    slice_ = f.get_slice(name)
+
+                    if isinstance(module, TensorParallelColumnLinear):
+                        if param_name == "weight":
+                            size = slice_.get_shape()[0]
+                            block_size = size // world_size
+                            start = rank * block_size
+                            stop = (rank + 1) * block_size
+                            tensor = slice_[start:stop]
+                            tensor = tensor.transpose(1, 0)
+                        else:
+                            size = slice_.get_shape()[0]
+                            block_size = size // world_size
+                            start = rank * block_size
+                            stop = (rank + 1) * block_size
+                            tensor = slice_[start:stop]
+                    elif isinstance(module, TensorParallelRowLinear):
+                        if param_name == "weight":
+                            size = slice_.get_shape()[1]
+                            block_size = size // world_size
+                            start = rank * block_size
+                            stop = (rank + 1) * block_size
+                            tensor = slice_[:, start:stop]
+                            tensor = tensor.transpose(1, 0)
+                        else:
+                            tensor = slice_[:]
+                            # XXX: Hack for Rowlinear to add the bias only once.
+                            if rank != 0:
+                                tensor = torch.zeros_like(tensor)
+                    elif isinstance(module, TensorParallelEmbedding):
+                        size = slice_.get_shape()[0]
+                        block_size = size // world_size
+                        start = rank * block_size
+                        stop = (rank + 1) * block_size
+                        tensor = slice_[start:stop]
+                    else:
+                        tensor = slice_[:]
+
+                    if current_tensor.shape != tensor.shape:
+                        raise ValueError(
+                            f"Name {name} -- Current {current_tensor.shape} and got {tensor.shape}"
+                        )
+
+                    tensor = tensor.contiguous()
+                    module._parameters[param_name] = tensor
+                    if name == "word_embeddings.weight":
+                        model.lm_head._parameters["weight"] = tensor
+
     def forward(self, input_ids, attention_mask, past_key_values: Optional = None):
         outputs = self.model.forward(
             input_ids=input_ids,

server/bloom_inference/prepare_weights.py

-import torch
-import os
-import tempfile
-import json
-
-from typing import BinaryIO
-from joblib import Parallel, delayed
-from functools import partial
-from pathlib import Path
-from tqdm import tqdm
-
-from huggingface_hub import hf_hub_url
-from huggingface_hub.file_download import _request_wrapper, hf_raise_for_status
-
-
-def match_suffix(text, suffix):
-    return text[-len(suffix) :] == suffix
-
-
-def http_get(
-    url: str,
-    temp_file: BinaryIO,
-    *,
-    timeout=10.0,
-    max_retries=0,
-):
-    """
-    Download a remote file. Do not gobble up errors, and will return errors tailored to the Hugging Face Hub.
-    """
-    r = _request_wrapper(
-        method="GET",
-        url=url,
-        stream=True,
-        timeout=timeout,
-        max_retries=max_retries,
-    )
-    hf_raise_for_status(r)
-    for chunk in r.iter_content(chunk_size=1024):
-        if chunk:  # filter out keep-alive new chunks
-            temp_file.write(chunk)
-
-
-def cache_download_url(url: str, root_dir: Path):
-    filename = root_dir / url.split("/")[-1]
-
-    if not filename.exists():
-        temp_file_manager = partial(
-            tempfile.NamedTemporaryFile, mode="wb", dir=root_dir, delete=False
-        )
-        with temp_file_manager() as temp_file:
-            http_get(url, temp_file)
-
-        os.replace(temp_file.name, filename)
-    return filename
-
-
-def prepare_weights(
-    model_name: str, cache_path: Path, save_path: Path, tp_world_size: int
-):
-    save_paths = [
-        save_path / f"{model_name}_tp-rank-{tp_rank}-of-{tp_world_size}.pty"
-        for tp_rank in range(tp_world_size)
-    ]
-
-    if all(save_path.exists() for save_path in save_paths):
-        print("Weights are already prepared")
-        return save_paths
-
-    cache_path.mkdir(parents=True, exist_ok=True)
-    if model_name == "bigscience/bloom-560m":
-        url = hf_hub_url(model_name, filename="pytorch_model.bin")
-        cache_download_url(url, cache_path)
-
-    elif model_name == "bigscience/bloom":
-        url = hf_hub_url(model_name, filename="pytorch_model.bin.index.json")
-        index_path = cache_download_url(url, cache_path)
-        with index_path.open("r") as f:
-            index = json.load(f)
-
-        # Get unique file names
-        weight_files = list(
-            set([filename for filename in index["weight_map"].values()])
-        )
-        urls = [hf_hub_url(model_name, filename=filename) for filename in weight_files]
-
-        Parallel(n_jobs=5)(
-            delayed(cache_download_url)(url, cache_path) for url in tqdm(urls)
-        )
-    else:
-        raise ValueError(f"Unknown model name: {model_name}")
-
-    shards_state_dicts = [{} for _ in range(tp_world_size)]
-
-    for weight_path in tqdm(Path(cache_path).glob("*.bin")):
-        state_dict = torch.load(weight_path, map_location="cpu")
-
-        keys = list(state_dict.keys())
-        for state_name in keys:
-            state = state_dict[state_name]
-            if any(
-                match_suffix(state_name, candidate)
-                for candidate in [
-                    "self_attention.query_key_value.weight",
-                    "self_attention.query_key_value.bias",
-                    "mlp.dense_h_to_4h.weight",
-                    "mlp.dense_h_to_4h.bias",
-                    "word_embeddings.weight",
-                ]
-            ):
-                output_size = state.shape[0]
-                assert output_size % tp_world_size == 0
-                block_size = output_size // tp_world_size
-                sharded_weights = torch.split(state, block_size, dim=0)
-                assert len(sharded_weights) == tp_world_size
-
-                for tp_rank, shard in enumerate(sharded_weights):
-                    shards_state_dicts[tp_rank][
-                        "transformer." + state_name
-                    ] = shard.detach().clone()
-
-            elif match_suffix(state_name, "lm_head.weight"):
-                output_size = state.shape[0]
-                assert output_size % tp_world_size == 0
-                block_size = output_size // tp_world_size
-                sharded_weights = torch.split(state, block_size, dim=0)
-                assert len(sharded_weights) == tp_world_size
-
-                for tp_rank, shard in enumerate(sharded_weights):
-                    shards_state_dicts[tp_rank][state_name] = shard.detach().clone()
-
-            elif any(
-                match_suffix(state_name, candidate)
-                for candidate in [
-                    "self_attention.dense.weight",
-                    "mlp.dense_4h_to_h.weight",
-                ]
-            ):
-                input_size = state.shape[1]
-                assert input_size % tp_world_size == 0
-                block_size = input_size // tp_world_size
-                sharded_weights = torch.split(state, block_size, dim=1)
-                assert len(sharded_weights) == tp_world_size
-                for tp_rank, shard in enumerate(sharded_weights):
-                    shards_state_dicts[tp_rank][
-                        "transformer." + state_name
-                    ] = shard.detach().clone()
-
-            elif any(
-                match_suffix(state_name, candidate)
-                for candidate in [
-                    "self_attention.dense.bias",
-                    "mlp.dense_4h_to_h.bias",
-                ]
-            ):
-                shards_state_dicts[0][
-                    "transformer." + state_name
-                ] = state.detach().clone()
-                for tp_rank in range(1, tp_world_size):
-                    shards_state_dicts[tp_rank][
-                        "transformer." + state_name
-                    ] = torch.zeros_like(state)
-
-            else:
-                # We duplicate parameters across tp ranks
-                for tp_rank in range(tp_world_size):
-                    shards_state_dicts[tp_rank][
-                        "transformer." + state_name
-                    ] = state.detach().clone()
-
-            del state_dict[state_name]  # delete key from state_dict
-            del state  # delete tensor
-        del state_dict
-
-    # we save state_dict
-    for tp_rank, (save_path, shard_state_dict) in enumerate(
-        zip(save_paths, shards_state_dicts)
-    ):
-        save_paths.append(save_path)
-        save_path.parent.mkdir(parents=True, exist_ok=True)
-        if save_path.exists():
-            print(f"Skipping {save_path} as it already exists")
-        else:
-            torch.save(shard_state_dict, save_path)
-
-    return save_paths

server/bloom_inference/server.py

@@ -69,18 +69,14 @@ def serve(
     model_name: str,
     sharded: bool,
     uds_path: Path,
-    shard_directory: Optional[Path] = None,
 ):
     async def serve_inner(
         model_name: str,
         sharded: bool = False,
-        shard_directory: Optional[Path] = None,
     ):
         unix_socket_template = "unix://{}-{}"
         if sharded:
-            if shard_directory is None:
-                raise ValueError("shard_directory must be set when sharded is True")
-            model = BLOOMSharded(model_name, shard_directory)
+            model = BLOOMSharded(model_name)
             server_urls = [
                 unix_socket_template.format(uds_path, rank)
                 for rank in range(model.world_size)
@@ -109,4 +105,4 @@ def serve(
             print("Signal received. Shutting down")
             await server.stop(0)
 
-    asyncio.run(serve_inner(model_name, sharded, shard_directory))
+    asyncio.run(serve_inner(model_name, sharded))

server/bloom_inference/utils.py

 import os
-import contextlib
 import torch
 import torch.distributed
 
 from datetime import timedelta
+
+from functools import partial
+from joblib import Parallel, delayed
+from huggingface_hub import HfApi, hf_hub_download, try_to_load_from_cache
+from huggingface_hub.utils import LocalEntryNotFoundError
+from tqdm import tqdm
 from transformers.generation_logits_process import (
     LogitsProcessorList,
     TemperatureLogitsWarper,
@@ -87,11 +92,42 @@ def initialize_torch_distributed():
     return torch.distributed.distributed_c10d._get_default_group(), rank, world_size
 
 
-@contextlib.contextmanager
-def set_default_dtype(dtype):
-    saved_dtype = torch.get_default_dtype()
-    torch.set_default_dtype(dtype)
-    try:
-        yield
-    finally:
-        torch.set_default_dtype(saved_dtype)
+def weight_hub_files(model_name):
+    """Get the safetensors filenames on the hub"""
+    api = HfApi()
+    info = api.model_info(model_name)
+    filenames = [
+        s.rfilename for s in info.siblings if s.rfilename.endswith(".safetensors")
+    ]
+    return filenames
+
+
+def weight_files(model_name):
+    """Get the local safetensors filenames"""
+    filenames = weight_hub_files(model_name)
+    files = []
+    for filename in filenames:
+        cache_file = try_to_load_from_cache(model_name, filename=filename)
+        if cache_file is None:
+            raise LocalEntryNotFoundError(
+                f"File {filename} of model {model_name} not found in "
+                f"{os.getenv('HUGGINGFACE_HUB_CACHE', 'the local cache')}. "
+                f"Please run `bloom-inference-server download-weights {model_name}` first."
+            )
+        files.append(cache_file)
+
+    return files
+
+
+def download_weights(model_name):
+    """Download the safetensors files from the hub"""
+    filenames = weight_hub_files(model_name)
+
+    download_function = partial(
+        hf_hub_download, repo_id=model_name, local_files_only=False
+    )
+    # FIXME: fix the overlapping progress bars
+    files = Parallel(n_jobs=5)(
+        delayed(download_function)(filename=filename) for filename in tqdm(filenames)
+    )
+    return files