f07a673eb2 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/models/test_utils.py

+70

@@ -77,3 +77,73 @@ def test_module_with_child_containing_batchnorm_can_autoload():

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

80 +

81 +

82 +def test_module_skip_prefix():

83 + """Ensure the auto weight loader can skip prefix."""

84 + mod = ModuleWithNestedBatchNorm()

85 + # Run some data through the module with batchnorm

86 + mod(torch.Tensor([[1, 2], [3, 4]]))

87 +

88 + # Try to load the weights to a new instance

89 + def weight_generator():

90 + # weights needed to be filtered out

91 + redundant_weights = {

92 + "prefix.bn.weight": torch.Tensor([1, 2]),

93 + "prefix.bn.bias": torch.Tensor([3, 4]),

94 + }

95 + yield from (mod.state_dict() | redundant_weights).items()

96 +

97 + new_mod = ModuleWithNestedBatchNorm()

98 +

99 + assert not torch.all(

100 + new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean)

101 + assert not torch.all(

102 + new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)

103 + assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0

104 +

105 + loader = AutoWeightsLoader(new_mod, skip_prefixes=["prefix."])

106 + loader.load_weights(weight_generator())

107 +

108 + # Ensure the stats are updated

109 + assert torch.all(

110 + new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean)

111 + assert torch.all(

112 + new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)

113 + assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1

114 +

115 +

116 +def test_module_skip_substr():

117 + """Ensure the auto weight loader can skip prefix."""

118 + mod = ModuleWithNestedBatchNorm()

119 + # Run some data through the module with batchnorm

120 + mod(torch.Tensor([[1, 2], [3, 4]]))

121 +

122 + # Try to load the weights to a new instance

123 + def weight_generator():

124 + # weights needed to be filtered out

125 + redundant_weights = {

126 + "nested_mod.0.substr.weight": torch.Tensor([1, 2]),

127 + "nested_mod.0.substr.bias": torch.Tensor([3, 4]),

128 + "nested_mod.substr.weight": torch.Tensor([1, 2]),

129 + "nested_mod.substr.bias": torch.Tensor([3, 4]),

130 + }

131 + yield from (mod.state_dict() | redundant_weights).items()

132 +

133 + new_mod = ModuleWithNestedBatchNorm()

134 +

135 + assert not torch.all(

136 + new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean)

137 + assert not torch.all(

138 + new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)

139 + assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0

140 +

141 + loader = AutoWeightsLoader(new_mod, skip_substrs=["substr."])

142 + loader.load_weights(weight_generator())

143 +

144 + # Ensure the stats are updated

145 + assert torch.all(

146 + new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean)

147 + assert torch.all(

148 + new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)

149 + assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1

vllm/model_executor/models/granite.py

+6 -10

@@ -478,18 +478,14 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

478

479 def load_weights(self, weights: Iterable[tuple[str,

480 torch.Tensor]]) -> set[str]:

481 - skip_prefixes = [

482 - "rotary_emb.inv_freq",

483 - # Models trained using ColossalAI may include these tensors in

484 - # the checkpoint. Skip them.

485 - "rotary_emb.cos_cached",

486 - "rotary_emb.sin_cached",

487 - ]

481 # With tie_word_embeddings, we can skip lm_head.weight

482 # The weight might appear unnecessarily in the files if the model is

483 # processed with quantization, LoRA, fine-tuning, etc.

491 - if self.config.tie_word_embeddings:

492 - skip_prefixes.append("lm_head.weight")

484 + skip_prefixes = (["lm_head."]

485 + if self.config.tie_word_embeddings else None)

486

494 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)

487 + loader = AutoWeightsLoader(

488 + self,

489 + skip_prefixes=skip_prefixes,

490 + )

491 return loader.load_weights(weights)

vllm/model_executor/models/grok1.py

+6 -4

@@ -550,10 +550,12 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

550

551 def load_weights(self, weights: Iterable[tuple[str,

552 torch.Tensor]]) -> set[str]:

553 - skip_prefixes = ["rotary_emb.inv_freq"]

553 # Skip lm_head when tie_word_embeddings is True

555 - if self.config.tie_word_embeddings:

556 - skip_prefixes.append("lm_head")

554 + skip_prefixes = (["lm_head"]

555 + if self.config.tie_word_embeddings else None)

556

558 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)

557 + loader = AutoWeightsLoader(

558 + self,

559 + skip_prefixes=skip_prefixes,

560 + )

561 return loader.load_weights(weights)

vllm/model_executor/models/mixtral.py

+1 -1

@@ -482,5 +482,5 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

482

483 def load_weights(self, weights: Iterable[tuple[str,

484 torch.Tensor]]) -> set[str]:

485 - loader = AutoWeightsLoader(self, skip_prefixes=["rotary_emb.inv_freq"])

485 + loader = AutoWeightsLoader(self)

486 return loader.load_weights(weights)

vllm/model_executor/models/mixtral_quant.py

+1 -4

@@ -447,8 +447,5 @@ class MixtralForCausalLM(nn.Module, SupportsPP):

447

448 def load_weights(self, weights: Iterable[tuple[str,

449 torch.Tensor]]) -> set[str]:

450 - loader = AutoWeightsLoader(

451 - self,

452 - skip_prefixes=(["rotary_emb.inv_freq"]),

453 - )

450 + loader = AutoWeightsLoader(self)

451 return loader.load_weights(weights)

vllm/model_executor/models/nemotron.py

+1 -10

@@ -502,14 +502,5 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

502

503 def load_weights(self, weights: Iterable[tuple[str,

504 torch.Tensor]]) -> set[str]:

505 - loader = AutoWeightsLoader(

506 - self,

507 - skip_prefixes=([

508 - "rotary_emb.inv_freq",

509 - # Models trained using ColossalAI may include these tensors in

510 - # the checkpoint. Skip them.

511 - "rotary_emb.cos_cached",

512 - "rotary_emb.sin_cached"

513 - ]),

514 - )

505 + loader = AutoWeightsLoader(self)

506 return loader.load_weights(weights)

vllm/model_executor/models/olmo.py

+2 -14

@@ -382,19 +382,7 @@ class OlmoForCausalLM(nn.Module, SupportsPP):

382 torch.Tensor]]) -> set[str]:

383 loader = AutoWeightsLoader(

384 self,

385 - skip_prefixes=([

386 - "rotary_emb.inv_freq",

387 - # Models trained using ColossalAI may include these tensors in

388 - # the checkpoint. Skip them.

389 - "rotary_emb.cos_cached",

390 - "rotary_emb.sin_cached",

391 - "lm_head.weight"

392 - ] if self.config.tie_word_embeddings else [

393 - "rotary_emb.inv_freq",

394 - # Models trained using ColossalAI may include these tensors in

395 - # the checkpoint. Skip them.

396 - "rotary_emb.cos_cached",

397 - "rotary_emb.sin_cached"

398 - ]),

385 + skip_prefixes=(["lm_head.weight"]

386 + if self.config.tie_word_embeddings else None),

387 )

388 return loader.load_weights(weights)

vllm/model_executor/models/olmo2.py

+2 -14

@@ -403,19 +403,7 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP):

403 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

404 loader = AutoWeightsLoader(

405 self,

406 - skip_prefixes=([

407 - "rotary_emb.inv_freq",

408 - # Models trained using ColossalAI may include these tensors in

409 - # the checkpoint. Skip them.

410 - "rotary_emb.cos_cached",

411 - "rotary_emb.sin_cached",

412 - "lm_head.weight"

413 - ] if self.config.tie_word_embeddings else [

414 - "rotary_emb.inv_freq",

415 - # Models trained using ColossalAI may include these tensors in

416 - # the checkpoint. Skip them.

417 - "rotary_emb.cos_cached",

418 - "rotary_emb.sin_cached"

419 - ]),

406 + skip_prefixes=(["lm_head.weight"]

407 + if self.config.tie_word_embeddings else None),

408 )

409 return loader.load_weights(weights)

vllm/model_executor/models/olmoe.py

+1 -4

@@ -442,8 +442,5 @@ class OlmoeForCausalLM(nn.Module, SupportsPP):

442

443 def load_weights(self, weights: Iterable[tuple[str,

444 torch.Tensor]]) -> set[str]:

445 - loader = AutoWeightsLoader(

446 - self,

447 - skip_prefixes=["rotary_emb.inv_freq"],

448 - )

445 + loader = AutoWeightsLoader(self)

446 return loader.load_weights(weights)

vllm/model_executor/models/orion.py

+1 -10

@@ -344,14 +344,5 @@ class OrionForCausalLM(nn.Module, SupportsPP):

344

345 def load_weights(self, weights: Iterable[tuple[str,

346 torch.Tensor]]) -> set[str]:

347 - loader = AutoWeightsLoader(

348 - self,

349 - skip_prefixes=([

350 - "rotary_emb.inv_freq",

351 - # Models trained using ColossalAI may include these tensors in

352 - # the checkpoint. Skip them.

353 - "rotary_emb.cos_cached",

354 - "rotary_emb.sin_cached"

355 - ]),

356 - )

347 + loader = AutoWeightsLoader(self)

348 return loader.load_weights(weights)

vllm/model_executor/models/phi4mm.py

+1 -3

@@ -1228,9 +1228,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1228

1229 def load_weights(self, weights: Iterable[tuple[str,

1230 torch.Tensor]]) -> None:

1231 - weights = ((name, data) for name, data in weights

1232 - if "lora" not in name)

1233 - loader = AutoWeightsLoader(self)

1231 + loader = AutoWeightsLoader(self, skip_substrs=["lora"])

1232 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

1233

1234 def get_mm_mapping(self) -> MultiModelKeys:

vllm/model_executor/models/phimoe.py

+1 -4

@@ -660,8 +660,5 @@ class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

660

661 def load_weights(self, weights: Iterable[tuple[str,

662 torch.Tensor]]) -> set[str]:

663 - loader = AutoWeightsLoader(

664 - self,

665 - skip_prefixes=(["rotary_emb.inv_freq"]),

666 - )

663 + loader = AutoWeightsLoader(self)

664 return loader.load_weights(weights)

vllm/model_executor/models/qwen2_moe.py

+1 -4

@@ -535,8 +535,5 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP):

535

536 def load_weights(self, weights: Iterable[tuple[str,

537 torch.Tensor]]) -> set[str]:

538 - loader = AutoWeightsLoader(

539 - self,

540 - skip_prefixes=(["rotary_emb.inv_freq"]),

541 - )

538 + loader = AutoWeightsLoader(self)

539 return loader.load_weights(weights)

vllm/model_executor/models/qwen3_moe.py

+1 -4

@@ -530,8 +530,5 @@ class Qwen3MoeForCausalLM(nn.Module, SupportsPP):

530

531 def load_weights(self, weights: Iterable[tuple[str,

532 torch.Tensor]]) -> set[str]:

533 - loader = AutoWeightsLoader(

534 - self,

535 - skip_prefixes=(["rotary_emb.inv_freq"]),

536 - )

533 + loader = AutoWeightsLoader(self)

534 return loader.load_weights(weights)

vllm/model_executor/models/solar.py

+1 -10

@@ -500,14 +500,5 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

500

501 def load_weights(self, weights: Iterable[tuple[str,

502 torch.Tensor]]) -> set[str]:

503 - loader = AutoWeightsLoader(

504 - self,

505 - skip_prefixes=([

506 - "rotary_emb.inv_freq",

507 - # Models trained using ColossalAI may include these tensors in

508 - # the checkpoint. Skip them.

509 - "rotary_emb.cos_cached",

510 - "rotary_emb.sin_cached"

511 - ]),

512 - )

503 + loader = AutoWeightsLoader(self)

504 return loader.load_weights(weights)

vllm/model_executor/models/stablelm.py

+1 -9

@@ -338,13 +338,5 @@ class StablelmForCausalLM(nn.Module, SupportsPP):

338

339 def load_weights(self, weights: Iterable[tuple[str,

340 torch.Tensor]]) -> set[str]:

341 - loader = AutoWeightsLoader(

342 - self,

343 - # Models trained using ColossalAI may include these tensors in

344 - # the checkpoint. Skip them.

345 - skip_prefixes=[

346 - "rotary_emb.inv_freq", "rotary_emb.cos_cached",

347 - "rotary_emb.sin_cached"

348 - ],

349 - )

341 + loader = AutoWeightsLoader(self)

342 return loader.load_weights(weights)

vllm/model_executor/models/starcoder2.py

+2 -3

@@ -349,8 +349,7 @@ class Starcoder2ForCausalLM(nn.Module, SupportsPP):

349 self,

350 # Models trained using ColossalAI may include these tensors in

351 # the checkpoint. Skip them.

352 - skip_prefixes=([

353 - "rotary_emb.inv_freq", "lm_head.weight"

354 - ] if self.config.tie_word_embeddings else ["rotary_emb.inv_freq"]),

352 + skip_prefixes=(["lm_head.weight"]

353 + if self.config.tie_word_embeddings else None),

354 )

355 return loader.load_weights(weights)

vllm/model_executor/models/utils.py

+17 -1

@@ -80,18 +80,30 @@ class AutoWeightsLoader:

80 environment variable ``VLLM_LOGGING_LEVEL=DEBUG``.

81 """

82

83 + # Models trained using early version ColossalAI

84 + # may include these tensors in checkpoint. Skip them.

85 + ROTARY_EMBEDS_UNUSED_WEIGHTS = [

86 + "rotary_emb.inv_freq",

87 + "rotary_emb.cos_cached",

88 + "rotary_emb.sin_cached",

89 + ]

90 +

91 def __init__(

92 self,

93 module: nn.Module,

94 *,

95 skip_prefixes: Optional[list[str]] = None,

96 + skip_substrs: Optional[list[str]] = None,

97 ignore_unexpected_prefixes: Optional[list[str]] = None,

98 ) -> None:

99 super().__init__()

100

101 self.module = module

102 self.skip_prefixes = skip_prefixes or []

103 + self.skip_substrs = skip_substrs or []

104 self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or []

105 + # update default skip_substrs

106 + self.skip_substrs += self.ROTARY_EMBEDS_UNUSED_WEIGHTS

107

108 def _groupby_prefix(

109 self,

@@ -119,7 +131,8 @@ class AutoWeightsLoader:

131 return ".".join((prefix, rest))

132

133 def _can_skip(self, qualname: str) -> bool:

122 - return any(qualname.startswith(p) for p in self.skip_prefixes)

134 + return (any(qualname.startswith(p) for p in self.skip_prefixes)

135 + or any(substr in qualname for substr in self.skip_substrs))

136

137 def _can_ignore_unexpected(self, qualname: str) -> bool:

138 return any(

@@ -257,6 +270,9 @@ class AutoWeightsLoader:

270 ) -> set[str]:

271 if mapper is not None:

272 weights = mapper.apply(weights)

273 + # filter out weights with first-prefix/substr to skip in name

274 + weights = ((name, weight) for name, weight in weights

275 + if not self._can_skip(name))

276

277 autoloaded_weights = set(self._load_module("", self.module, weights))

278 return autoloaded_weights

[Misc] Allow `AutoWeightsLoader` to skip loading weights with specific substr in name (#18358)