fix examples

Signed-off-by: Takeshi Yoshimura <tyos@jp.ibm.com>

Author: takeshi-yoshimura

.github/workflows/test-torch.yaml

@@ -58,8 +58,9 @@ jobs:
           mkdir -p /tmp/pytest-log
           export TEST_FASTSAFETENSORS_FRAMEWORK=pytorch
           COVERAGE_FILE=.coverage_0 pytest -s --cov=${LIBDIR} test_fastsafetensors.py > /tmp/pytest-log/0.log 2>&1
-          COVERAGE_FILE=.coverage_1 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 tests/test_multi.py --cov=${LIBDIR} -s tests/test_multi.py > /tmp/pytest-log/1.log 2>&1 &
-          COVERAGE_FILE=.coverage_2 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 tests/test_multi.py --cov=${LIBDIR} -s tests/test_multi.py > /tmp/pytest-log/2.log 2>&1
+          COVERAGE_FILE=.coverage_1 pytest -s --cov=${LIBDIR} test_vllm.py > /tmp/pytest-log/1.log 2>&1
+          COVERAGE_FILE=.coverage_2 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 tests/test_multi.py --cov=${LIBDIR} -s tests/test_multi.py > /tmp/pytest-log/2.log 2>&1 &
+          COVERAGE_FILE=.coverage_3 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 tests/test_multi.py --cov=${LIBDIR} -s tests/test_multi.py > /tmp/pytest-log/3.log 2>&1
           coverage combine .coverage_*
           coverage html
           mv htmlcov /tmp/pytest-log

Makefile

@@ -19,20 +19,21 @@ unittest:
 	@FST_DIR=$(FST_DIR); \
 	TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_0 pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py && \
 	TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_1 CUDA_VISIBLE_DEVICES="" pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py && \
-	TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_2 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/2.log 2>&1 & \
-	TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_3 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/3.log 2>&1 && \
+	TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_2 pytest -s --cov=$(FST_DIR) -s tests/test_vllm.py && \
+	TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_3 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/3.log 2>&1 & \
+	TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_4 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/4.log 2>&1 && \
 	wait
 
 unittest-paddle:
 	@FST_DIR=$(FST_DIR); \
-	TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_4 CUDA_VISIBLE_DEVICES="" pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py && \
-	TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_5 CUDA_VISIBLE_DEVICES="" WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 0 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/2.log 2>&1 & \
-	TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_6 CUDA_VISIBLE_DEVICES="" WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 1 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/3.log 2>&1 && \
+	TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_5 CUDA_VISIBLE_DEVICES="" pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py && \
+	TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_6 CUDA_VISIBLE_DEVICES="" WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 0 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/6.log 2>&1 & \
+	TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_7 CUDA_VISIBLE_DEVICES="" WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 1 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/7.log 2>&1 && \
 	wait
 
 unittest-paddle-gpu:
 	@FST_DIR=$(FST_DIR); \
-	TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_7 pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py
+	TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_8 pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py
 
 htmlcov:
 	coverage combine .coverage_* && \

examples/a_paddle.safetensors

@@ -2,7 +2,7 @@
 import sys
 from typing import Dict, List
 
-from safetensors import safe_open
+from safetensors.torch import load_file
 
 from fastsafetensors import SafeTensorsFileLoader, SingleGroup
 
@@ -23,9 +23,7 @@
         ".safetensors"
     ):
         src_files[0].append(input_file_or_dir)
-        with safe_open(input_file_or_dir, framework="pytorch") as f:
-            for key in f.keys():
-                orig_keys[key] = f.get_tensor(key)
+        orig_keys = load_file(input_file_or_dir)
     loader.add_filenames(src_files)
     fb = loader.copy_files_to_device()
     if len(orig_keys) > 0:

examples/b_paddle.safetensors

@@ -19,7 +19,7 @@ def gen_paddle():
     from safetensors.paddle import save_file
 
     t0 = paddle.concat(
-        [paddle.full((1, 8), i, dtype=paddle.float16) for i in range(0, 16)], dim=0
+        [paddle.full((1, 8), i, dtype=paddle.float16) for i in range(0, 16)], axis=0
     )
 
     for file_prefix in ["a", "b"]:

examples/extract_keys.py

@@ -2,12 +2,12 @@
 
 import torch
 
-from fastsafetensors import SafeTensorsFileLoader
+from fastsafetensors import SafeTensorsFileLoader, SingleGroup
 
 sys.path.insert(0, "/nvme/manish/repos/fastsafetensors/fastsafetensors")
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
-loader = SafeTensorsFileLoader(device, nogds=True, debug_log=True)
+loader = SafeTensorsFileLoader(SingleGroup(), device, nogds=True, debug_log=True)
 
 loader.add_filenames({0: ["a.safetensors"]})  # {rank: files}
 fb = loader.copy_files_to_device()

examples/gen.py

@@ -1,38 +1,29 @@
 #!/usr/bin/env python3
 
-
-def run_torch():
-    import torch
-
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    return device
-
-
-def run_paddle():
-    import paddle
-
-    device = "gpu" if paddle.device.cuda.device_count() else "cpu"
-    return device
-
-
-runs = {
-    "torch": run_torch,
-    "paddle": run_paddle,
-}
-
 if __name__ == "__main__":
     import sys
 
     from fastsafetensors import cpp as fstcpp
     from fastsafetensors import fastsafe_open
 
     framework = "torch"
+    filenames = ["a.safetensors", "b.safetensors"]
     if len(sys.argv) > 1:
         framework = sys.argv[1]
+    if framework == "torch":
+        import torch
+
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    elif framework == "paddle":
+        import paddle
+
+        device = "gpu" if paddle.device.cuda.device_count() else "cpu"
+        filenames = ["a_paddle.safetensors", "b_paddle.safetensors"]
+    else:
+        raise Exception(f"unknown framework: {framework}")
 
-    device = runs[framework]()
     with fastsafe_open(
-        ["a.safetensors", "b.safetensors"],
+        filenames,
         device=device,
         nogds=not fstcpp.is_cufile_found(),
         framework=framework,

examples/run_reuse_loader.py

@@ -4,12 +4,3 @@
 from .common import SafeTensorsMetadata, SingleGroup, TensorFrame, get_device_numa_node
 from .file_buffer import FilesBufferOnDevice
 from .loader import SafeTensorsFileLoader, fastsafe_open
-
-loaded_nvidia: bool = False
-if not loaded_nvidia:
-    from . import cpp
-
-    cpp.load_nvidia_functions()
-    if cpp.init_gds() != 0:
-        raise Exception(f"[FAIL] init_gds()")
-    loaded_nvidia = True