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 4, 08:45 UTC
[Bugfix] Fix FP8 Bias Loading (#41424)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Alex Brooks committed
May 3, 2026 at 14:30 UTC
db9a84e0cd0e17ab693467ff4a71103abd4b77bf
2 files changed
+29
-1
tests/model_executor/model_loader/test_reload.py
+28
@@ -59,6 +59,34 @@ def test_reload_lifecycle():
59
assert tensor.__dict__ == materialized_tensor.__dict__
60
61
62
+def test_materialize_layer_preserves_non_meta_tensors():
63
+ """Ensure that materialize_layer does not overwrite non meta tensors."""
64
+ layer = torch.nn.Linear(2, 3, bias=True)
65
+
66
+ # Create a non meta bias tensor and meta weight, which can happen with FP8
67
+ bias_values = torch.ones(3)
68
+ layer.bias.data.copy_(bias_values)
69
+ layer.weight = torch.nn.Parameter(layer.weight.data.to("meta"))
70
+
71
+ assert layer.weight.is_meta
72
+ assert not layer.bias.is_meta
73
+
74
+ # materialize the layer weights after the bias is initialized
75
+ info = LayerReloadingInfo(
76
+ restore_metadata=({}, {}),
77
+ restore_device=torch.device("cpu"),
78
+ )
79
+ materialize_layer(layer, info)
80
+
81
+ # Ensure the weight materialized off meta
82
+ assert not layer.weight.is_meta
83
+ assert layer.weight.device.type == "cpu"
84
+
85
+ # Ensure that the bias is (still) not meta and values are unchanged
86
+ assert not layer.bias.is_meta
87
+ assert torch.equal(layer.bias.data, bias_values)
88
+
89
+
90
def test_model_cleanup(dist_init, default_vllm_config):
91
layer = QKVParallelLinear(2, 3, 4)
92
assert layer.weight.weight_loader.__self__ is layer
vllm/model_executor/model_loader/reload/meta.py
+1
-1
@@ -102,7 +102,7 @@ def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo):
102
103
with info.restore_device:
104
for name, tensor in get_layer_tensors(layer).items():
105
- if name not in SKIP_TENSORS:
105
+ if name not in SKIP_TENSORS and tensor.is_meta:
106
setattr(layer, name, materialize_meta_tensor(tensor))
107
108