[Auto Sync] Update linear.py to assert shapes (20260130) (#17966)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Byron Hsu <byronhsu1230@gmail.com>
This commit is contained in:
Lianmin Zheng
2026-01-31 01:01:55 -08:00
committed by GitHub
parent 0e184609d3
commit 7412ceb4eb

View File

@@ -1222,7 +1222,9 @@ class QKVParallelLinear(ColumnParallelLinear):
"for all partitions."
)
assert param_data.shape == loaded_weight.shape
assert (
param_data.shape == loaded_weight.shape
), f"{param_data.shape=} {loaded_weight.shape=}"
param_data.copy_(loaded_weight)
@@ -1371,7 +1373,9 @@ class RowParallelLinear(LinearBase):
if len(loaded_weight.shape) == 0:
loaded_weight = loaded_weight.reshape(1)
assert param_data.shape == loaded_weight.shape
assert (
param_data.shape == loaded_weight.shape
), f"{param_data.shape=} {loaded_weight.shape=}"
param_data.copy_(loaded_weight)
def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor):