From cf22255d52ae2452bc54b0df66451df1bff6a3ff Mon Sep 17 00:00:00 2001 From: Taksh Date: Sat, 30 May 2026 15:47:00 +0530 Subject: [PATCH 1/2] fix: build RWKV strategy from visible CUDA device count Use multi-GPU RWKV strategy when more than one GPU is visible via --gpus / --num-gpus instead of hardcoding single-device cuda fp16. --- fastchat/model/rwkv_model.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/fastchat/model/rwkv_model.py b/fastchat/model/rwkv_model.py index bdbc14584..0a0c65bd7 100644 --- a/fastchat/model/rwkv_model.py +++ b/fastchat/model/rwkv_model.py @@ -17,9 +17,13 @@ def __init__(self, model_path): "Experimental support. Please use ChatRWKV if you want to chat with RWKV" ) self.config = SimpleNamespace(is_encoder_decoder=False) - self.model = RWKV(model=model_path, strategy="cuda fp16") - # two GPUs - # self.model = RWKV(model=model_path, strategy="cuda:0 fp16 *20 -> cuda:1 fp16") + n = torch.cuda.device_count() + strategy = ( + " -> ".join(f"cuda:{i} fp16" for i in range(n)) + if n > 1 + else "cuda fp16" + ) + self.model = RWKV(model=model_path, strategy=strategy) self.tokenizer = None self.model_path = model_path From 4a5651613cc7e2f5f514b7376b93f0bfd18e22df Mon Sep 17 00:00:00 2001 From: Taksh Date: Wed, 29 Jul 2026 18:42:54 +0300 Subject: [PATCH 2/2] style: format rwkv_model.py for black CI --- fastchat/model/rwkv_model.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/fastchat/model/rwkv_model.py b/fastchat/model/rwkv_model.py index 0a0c65bd7..4245b4de1 100644 --- a/fastchat/model/rwkv_model.py +++ b/fastchat/model/rwkv_model.py @@ -19,9 +19,7 @@ def __init__(self, model_path): self.config = SimpleNamespace(is_encoder_decoder=False) n = torch.cuda.device_count() strategy = ( - " -> ".join(f"cuda:{i} fp16" for i in range(n)) - if n > 1 - else "cuda fp16" + " -> ".join(f"cuda:{i} fp16" for i in range(n)) if n > 1 else "cuda fp16" ) self.model = RWKV(model=model_path, strategy=strategy)