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, 14:45 UTC
[TPU] Fix dummy loading OOM (#16372)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Apr 9, 2025 at 21:06 UTC
1621b252888887cc6601c5df2dfa8841ae305d8e
1 file changed
+15
-2
vllm/model_executor/model_loader/weight_utils.py
+15
-2
@@ -658,8 +658,21 @@ def initialize_dummy_weights(
658
for param in model.state_dict().values():
659
if torch.is_floating_point(param):
660
if current_platform.is_tpu():
661
- # XLA device does not support torch.Generator()
662
- param.uniform_(low, high)
661
+ generator = torch.Generator(device="cpu")
662
+ generator.manual_seed(seed)
663
+ # Note: The param.uniform_ function cannot be used in this
664
+ # context because it demands more TPU HBM than directly copying
665
+ # from a CPU tensor.
666
+ # Note: We avoid using torch.rank_like as it doesn't currently
667
+ # support the generator argument.
668
+ param.copy_((high - low) *
669
+ torch.rand(*param.shape,
670
+ generator=generator,
671
+ dtype=param.dtype,
672
+ layout=param.layout,
673
+ requires_grad=param.requires_grad,
674
+ device="cpu") + low)
675
+ torch._sync(param)
676
continue
677
678
generator = torch.Generator(device=param.data.device)