A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 17:15 UTC
[Bugfix] Fix Crashing When Loading Modules With Batchnorm Stats (#15813)
Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com>
Alex Brooks committed
Mar 31, 2025 at 07:23 UTC
c2e7507ad4a022befe418acefccc3bf883a0aad5
2 files changed
+103
tests/models/test_utils.py
new
+79
@@ -0,0 +1,79 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import torch
4
+
5
+from vllm.model_executor.models.utils import AutoWeightsLoader
6
+
7
+
8
+class ModuleWithBatchNorm(torch.nn.Module):
9
+
10
+ def __init__(self):
11
+ super().__init__()
12
+ self.bn = torch.nn.BatchNorm1d(2)
13
+
14
+ def forward(self, x):
15
+ return self.bn(x)
16
+
17
+
18
+class ModuleWithNestedBatchNorm(torch.nn.Module):
19
+
20
+ def __init__(self):
21
+ super().__init__()
22
+ self.nested_mod = ModuleWithBatchNorm()
23
+
24
+ def forward(self, x):
25
+ return self.nested_mod(x)
26
+
27
+
28
+def test_module_with_batchnorm_can_load():
29
+ """Ensure the auto weight loader can load batchnorm stats."""
30
+ mod = ModuleWithBatchNorm()
31
+ # Run some data through the module with batchnorm
32
+ mod(torch.Tensor([[1, 2], [3, 4]]))
33
+
34
+ # Try to load the weights to a new instance
35
+ def weight_generator():
36
+ yield from mod.state_dict().items()
37
+
38
+ new_mod = ModuleWithBatchNorm()
39
+
40
+ assert not torch.all(new_mod.bn.running_mean == mod.bn.running_mean)
41
+ assert not torch.all(new_mod.bn.running_var == mod.bn.running_var)
42
+ assert new_mod.bn.num_batches_tracked.item() == 0
43
+
44
+ loader = AutoWeightsLoader(new_mod)
45
+ loader.load_weights(weight_generator())
46
+
47
+ # Ensure the stats are updated
48
+ assert torch.all(new_mod.bn.running_mean == mod.bn.running_mean)
49
+ assert torch.all(new_mod.bn.running_var == mod.bn.running_var)
50
+ assert new_mod.bn.num_batches_tracked.item() == 1
51
+
52
+
53
+def test_module_with_child_containing_batchnorm_can_autoload():
54
+ """Ensure the auto weight loader can load nested modules batchnorm stats."""
55
+ mod = ModuleWithNestedBatchNorm()
56
+ # Run some data through the module with batchnorm
57
+ mod(torch.Tensor([[1, 2], [3, 4]]))
58
+
59
+ # Try to load the weights to a new instance
60
+ def weight_generator():
61
+ yield from mod.state_dict().items()
62
+
63
+ new_mod = ModuleWithNestedBatchNorm()
64
+
65
+ assert not torch.all(
66
+ new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean)
67
+ assert not torch.all(
68
+ new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)
69
+ assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0
70
+
71
+ loader = AutoWeightsLoader(new_mod)
72
+ loader.load_weights(weight_generator())
73
+
74
+ # Ensure the stats are updated
75
+ assert torch.all(
76
+ new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean)
77
+ assert torch.all(
78
+ new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)
79
+ assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
vllm/model_executor/models/utils.py
+24
@@ -158,6 +158,26 @@ class AutoWeightsLoader:
158
159
yield weight_qualname
160
161
+ def _add_loadable_non_param_tensors(self, module: nn.Module,
162
+ child_params: Dict[str, torch.Tensor]):
163
+ """
164
+ Add tensor names that are not in the model params that may be in the
165
+ safetensors, e.g., batch normalization stats.
166
+ """
167
+ if isinstance(module, (
168
+ nn.BatchNorm1d,
169
+ nn.BatchNorm2d,
170
+ nn.BatchNorm3d,
171
+ nn.LazyBatchNorm1d,
172
+ nn.LazyBatchNorm2d,
173
+ nn.LazyBatchNorm3d,
174
+ nn.SyncBatchNorm,
175
+ )):
176
+ module_state_dict = module.state_dict()
177
+ for stat_name in ("running_mean", "running_var",
178
+ "num_batches_tracked"):
179
+ child_params[stat_name] = module_state_dict[stat_name]
180
+
181
def _load_module(
182
self,
183
base_prefix: str,
@@ -186,6 +206,10 @@ class AutoWeightsLoader:
206
child_modules = dict(module.named_children())
207
child_params = dict(module.named_parameters(recurse=False))
208
209
+ # Add missing tensors the weight loader needs to be able to load
210
+ # that aren't registered as params, e.g., batchnorm statistics.
211
+ self._add_loadable_non_param_tensors(module, child_params)
212
+
213
for child_prefix, child_weights in self._groupby_prefix(weights):
214
prefix = self._get_qualname(base_prefix, child_prefix)
215