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 11, 05:00 UTC
fix fp8 online quantization streaming with tp > 1 (#30900)
Signed-off-by: vasiliy <vasiliy@fb.com>
Vasiliy Kuznetsov committed
Dec 18, 2025 at 11:45 UTC
f4ee2c3d908ade890bc2e42945d307bc2c341b59
1 file changed
+33
-8
vllm/model_executor/layers/quantization/fp8.py
+33
-8
@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Optional
8
import torch
9
from torch.nn import Module
10
from torch.nn.parameter import Parameter
11
+from torch.utils._python_dispatch import TorchDispatchMode
12
13
import vllm.envs as envs
14
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -363,6 +364,26 @@ class Fp8Config(QuantizationConfig):
364
return None
365
366
367
+class CopyNumelCounter(TorchDispatchMode):
368
+ """
369
+ Tracks total number of elements modified with `copy_`. Useful for keeping
370
+ track of weight loading where underlying weights can be arbitrarily
371
+ transformed (such as with `narrow`) before calling copy.
372
+ """
373
+
374
+ def __init__(self):
375
+ super().__init__()
376
+ self.copied_numel = 0
377
+
378
+ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
379
+ if kwargs is None:
380
+ kwargs = {}
381
+ out = func(*args, **kwargs)
382
+ if func == torch.ops.aten.copy_.default:
383
+ self.copied_numel += args[0].numel()
384
+ return out
385
+
386
+
387
class Fp8LinearMethod(LinearMethodBase):
388
"""Linear method for FP8.
389
Supports loading FP8 checkpoints with static weight scale and
@@ -469,13 +490,15 @@ class Fp8LinearMethod(LinearMethodBase):
490
else:
491
492
def patched_weight_loader(param, loaded_weight, *args, **kwargs):
472
- # load the current weight chunk
473
- res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]
474
-
493
# track how many elements we have updated
494
if not hasattr(layer, "_loaded_numel"):
495
layer._loaded_numel = 0
478
- layer._loaded_numel += loaded_weight.numel()
496
+
497
+ # load the current weight chunk
498
+ copy_numel_counter = CopyNumelCounter()
499
+ with copy_numel_counter:
500
+ res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]
501
+ layer._loaded_numel += copy_numel_counter.copied_numel
502
503
# if we have loaded all of the elements, call
504
# process_weights_after_loading
@@ -1348,13 +1371,15 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1371
new_extra_weight_attrs = extra_weight_attrs
1372
1373
def patched_weight_loader(param, loaded_weight, *args, **kwargs):
1351
- # load the current weight chunk
1352
- res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]
1353
-
1374
# add a counter to track how many elements we have updated
1375
if not hasattr(layer, "_loaded_numel"):
1376
layer._loaded_numel = 0
1357
- layer._loaded_numel += loaded_weight.numel()
1377
+
1378
+ # load the current weight chunk
1379
+ copy_numel_counter = CopyNumelCounter()
1380
+ with copy_numel_counter:
1381
+ res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]
1382
+ layer._loaded_numel += copy_numel_counter.copied_numel
1383
1384
# if we have loaded all of the elements, call
1385
# process_weights_after_loading