"...git@developer.sourcefind.cn:kecinstone/2024-pra-vllm.git" did not exist on "55fe8a81ec5deab319c7c6b02913c21273b764ca"
Unverified Commit dd9433e7 authored by Peter St. John's avatar Peter St. John Committed by GitHub
Browse files

Don't pickle an empty dict in LayerNorm and pt base modules (#2253)



Don't pickle an empty dict in LayerNorm and BasicOperation layers
Signed-off-by: default avatarPeter St. John <pstjohn@nvidia.com>
Co-authored-by: default avatarTim Moon <4406448+timmoon10@users.noreply.github.com>
parent 8a7ab3dd
...@@ -595,6 +595,9 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): ...@@ -595,6 +595,9 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta):
extra[key] = val extra[key] = val
state[mode]["extra_fp8_variables"] = extra state[mode]["extra_fp8_variables"] = extra
if not state:
return torch.empty(0, dtype=torch.uint8)
# Serialize state into byte tensor # Serialize state into byte tensor
torch.cuda.synchronize() torch.cuda.synchronize()
state_serialized = bytearray(pickle.dumps(state)) state_serialized = bytearray(pickle.dumps(state))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment