mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-01-23 10:24:28 +08:00
[Bugfix] Fix _synced_weight_loader (#24565)
Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
This commit is contained in:
parent
0fc36463e0
commit
ed5ae4aace
@ -52,10 +52,11 @@ def set_weight_attrs(
|
||||
def _make_synced_weight_loader(original_weight_loader):
|
||||
|
||||
def _synced_weight_loader(param, *args, **kwargs):
|
||||
original_weight_loader(param, *args, **kwargs)
|
||||
out = original_weight_loader(param, *args, **kwargs)
|
||||
# torch._sync doesn't support, is not needed for CPU tensors.
|
||||
if param.device != torch.device("cpu"):
|
||||
torch._sync(param)
|
||||
return out
|
||||
|
||||
return _synced_weight_loader
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user