mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-05-16 23:35:44 +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 _make_synced_weight_loader(original_weight_loader):
|
||||||
|
|
||||||
def _synced_weight_loader(param, *args, **kwargs):
|
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.
|
# torch._sync doesn't support, is not needed for CPU tensors.
|
||||||
if param.device != torch.device("cpu"):
|
if param.device != torch.device("cpu"):
|
||||||
torch._sync(param)
|
torch._sync(param)
|
||||||
|
return out
|
||||||
|
|
||||||
return _synced_weight_loader
|
return _synced_weight_loader
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user