A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 16:30 UTC
[BugFix] Fix `fastsafetensors` TP all procs using all GPUs (#34070)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Nick Hill committed
Feb 8, 2026 at 23:15 UTC
d9bede0314ba19a3f8336dcaeeeaf9e2c5487053
1 file changed
+9
-6
vllm/model_executor/model_loader/weight_utils.py
+9
-6
@@ -801,8 +801,8 @@ def runai_safetensors_weights_iterator(
801
yield from tensor_iter
802
803
804
-def _init_loader(
805
- pg: torch.distributed.ProcessGroup,
804
+def _init_fastsafetensors_loader(
805
+ pg: "torch.distributed.ProcessGroup",
806
device: torch.device,
807
f_list: list[str],
808
*,
@@ -825,13 +825,16 @@ def fastsafetensors_weights_iterator(
825
else:
826
pg = SingleGroup()
827
828
- device = torch.device(f"cuda:{pg.rank()}")
828
+ device = torch.device(f"cuda:{current_platform.current_device()}")
829
weight_files_sub_lists = [
830
hf_weights_files[i : i + pg.size()]
831
for i in range(0, len(hf_weights_files), pg.size())
832
]
833
834
- nogds = False
834
+ # Use nogds=True for TP > 1 to avoid cuFileDriverOpen() which
835
+ # initializes the GDS DMA subsystem for all visible GPUs, creating
836
+ # unwanted CUDA contexts on every device.
837
+ nogds = pg.size() > 1
838
839
for f_list in tqdm(
840
weight_files_sub_lists,
@@ -839,7 +842,7 @@ def fastsafetensors_weights_iterator(
842
disable=not enable_tqdm(use_tqdm_on_load),
843
bar_format=_BAR_FORMAT,
844
):
842
- loader = _init_loader(pg, device, f_list, nogds=nogds)
845
+ loader = _init_fastsafetensors_loader(pg, device, f_list, nogds=nogds)
846
try:
847
try:
848
fb = loader.copy_files_to_device()
@@ -853,7 +856,7 @@ def fastsafetensors_weights_iterator(
856
"GDS not enabled, setting `nogds=True`.\n"
857
"For more information, see: https://github.com/foundation-model-stack/fastsafetensors?tab=readme-ov-file#basic-api-usages"
858
)
856
- loader = _init_loader(pg, device, f_list, nogds=nogds)
859
+ loader = _init_fastsafetensors_loader(pg, device, f_list, nogds=nogds)
860
fb = loader.copy_files_to_device()
861
862
try: