cd74911d92 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -83,6 +83,7 @@ from vllm.model_executor.model_loader.weight_utils import (

83 maybe_remap_kv_scale_name,

84 )

85 from vllm.model_executor.models.utils import (

86 + AutoWeightsLoader,

87 extract_layer_index,

88 sequence_parallel_chunk,

89 )

@@ -1254,6 +1255,16 @@ class DeepseekV2Model(nn.Module):

1255

1256 self.aux_hidden_state_layers = tuple[int, ...]()

1257

1258 + # Needed by load_weights

1259 + qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)

1260 + qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)

1261 + self.use_mha = config.model_type == "deepseek" or all(

1262 + dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)

1263 + )

1264 + self.num_redundant_experts = (

1265 + vllm_config.parallel_config.eplb_config.num_redundant_experts

1266 + )

1267 +

1268 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1269 return self.embed_tokens(input_ids)

1270

@@ -1315,174 +1326,6 @@ class DeepseekV2Model(nn.Module):

1326 return hidden_states, aux_hidden_states

1327 return hidden_states

1328

1318 -

1319 -class DeepseekV2MixtureOfExperts(MixtureOfExperts):

1320 - moe_mlp_layers: list[DeepseekV2MoE]

1321 - """

1322 - List of MoE MLP layers in the model.

1323 - """

1324 -

1325 - def extract_moe_parameters(self, example_moe: DeepseekV2MoE | None):

1326 - if example_moe is None:

1327 - self.num_moe_layers = 0

1328 - self.num_expert_groups = 0

1329 - self.num_logical_experts = 0

1330 - self.num_physical_experts = 0

1331 - self.num_local_physical_experts = 0

1332 - self.num_routed_experts = 0

1333 - self.num_shared_experts = 0

1334 - self.num_redundant_experts = 0

1335 - logger.warning("DeepSeekV2: No DeepseekV2MoE layer found in model.layers.")

1336 - else:

1337 - self.num_logical_experts = example_moe.n_logical_experts

1338 - self.num_physical_experts = example_moe.n_physical_experts

1339 - self.num_local_physical_experts = example_moe.n_local_physical_experts

1340 - self.num_routed_experts = example_moe.n_routed_experts

1341 - self.num_shared_experts = example_moe.n_shared_experts

1342 - self.num_redundant_experts = example_moe.n_redundant_experts

1343 -

1344 - def update_physical_experts_metadata(

1345 - self,

1346 - num_physical_experts: int,

1347 - num_local_physical_experts: int,

1348 - ) -> None:

1349 - assert self.num_local_physical_experts == num_local_physical_experts

1350 - self.num_physical_experts = num_physical_experts

1351 - self.num_local_physical_experts = num_local_physical_experts

1352 - self.num_redundant_experts = num_physical_experts - self.num_logical_experts

1353 - for moe in self.moe_mlp_layers:

1354 - moe.n_local_physical_experts = num_local_physical_experts

1355 - moe.n_physical_experts = num_physical_experts

1356 - moe.n_redundant_experts = self.num_redundant_experts

1357 - moe.experts.update_expert_map()

1358 -

1359 -

1360 -class DeepseekV2ForCausalLM(

1361 - nn.Module,

1362 - SupportsPP,

1363 - DeepseekV2MixtureOfExperts,

1364 - SupportsLoRA,

1365 - SupportsEagle,

1366 - SupportsEagle3,

1367 -):

1368 - packed_modules_mapping = {

1369 - "gate_up_proj": ["gate_proj", "up_proj"],

1370 - }

1371 - model_cls = DeepseekV2Model

1372 -

1373 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1374 - super().__init__()

1375 - config = vllm_config.model_config.hf_config

1376 - quant_config = vllm_config.quant_config

1377 - self.config = config

1378 - self.quant_config = quant_config

1379 -

1380 - qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)

1381 - qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)

1382 - self.use_mha = config.model_type == "deepseek" or all(

1383 - dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)

1384 - )

1385 -

1386 - if self.use_mha:

1387 - self.packed_modules_mapping["qkv_proj"] = ["q_proj", "k_proj", "v_proj"]

1388 -

1389 - # `packed_modules_mapping` needs to be modified before

1390 - # initializing DeepseekV2Model, as it is passed inplace to

1391 - # quantization config init and may be used to select the

1392 - # quant_method for relevant layers during initialization.

1393 - self.fuse_qkv_a_proj = (

1394 - hasattr(config, "q_lora_rank") and config.q_lora_rank is not None

1395 - )

1396 - if self.fuse_qkv_a_proj:

1397 - self.packed_modules_mapping["fused_qkv_a_proj"] = [

1398 - "q_a_proj",

1399 - "kv_a_proj_with_mqa",

1400 - ]

1401 -

1402 - self.model = self.model_cls(

1403 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

1404 - )

1405 - if get_pp_group().is_last_rank:

1406 - self.lm_head = ParallelLMHead(

1407 - config.vocab_size,

1408 - config.hidden_size,

1409 - quant_config=quant_config,

1410 - prefix=maybe_prefix(prefix, "lm_head"),

1411 - )

1412 - else:

1413 - self.lm_head = PPMissingLayer()

1414 - self.logits_processor = LogitsProcessor(config.vocab_size)

1415 - self.make_empty_intermediate_tensors = (

1416 - self.model.make_empty_intermediate_tensors

1417 - )

1418 - # Set MoE hyperparameters

1419 - self.num_moe_layers = (

1420 - self.config.num_hidden_layers - self.config.first_k_dense_replace

1421 - )

1422 - self.set_moe_parameters()

1423 -

1424 - def set_moe_parameters(self):

1425 - self.expert_weights = []

1426 -

1427 - self.num_expert_groups = getattr(self.config, "n_group", 1)

1428 -

1429 - self.moe_layers = []

1430 - self.moe_mlp_layers = []

1431 - example_moe = None

1432 - for layer in self.model.layers:

1433 - if isinstance(layer, PPMissingLayer):

1434 - continue

1435 -

1436 - assert isinstance(layer, DeepseekV2DecoderLayer)

1437 - if isinstance(layer.mlp, DeepseekV2MoE):

1438 - # Pick last one layer since the first ones may be dense layers.

1439 - example_moe = layer.mlp

1440 - self.moe_mlp_layers.append(layer.mlp)

1441 - self.moe_layers.append(layer.mlp.experts)

1442 -

1443 - self.extract_moe_parameters(example_moe)

1444 -

1445 - def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:

1446 - self.model.aux_hidden_state_layers = layers

1447 -

1448 - def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:

1449 - num_layers = len(self.model.layers)

1450 - return (2, num_layers // 2, num_layers - 3)

1451 -

1452 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1453 - return self.model.embed_input_ids(input_ids)

1454 -

1455 - def forward(

1456 - self,

1457 - input_ids: torch.Tensor | None,

1458 - positions: torch.Tensor,

1459 - intermediate_tensors: IntermediateTensors | None = None,

1460 - inputs_embeds: torch.Tensor | None = None,

1461 - ) -> torch.Tensor | IntermediateTensors:

1462 - hidden_states = self.model(

1463 - input_ids, positions, intermediate_tensors, inputs_embeds

1464 - )

1465 - return hidden_states

1466 -

1467 - def compute_logits(

1468 - self,

1469 - hidden_states: torch.Tensor,

1470 - ) -> torch.Tensor | None:

1471 - logits = self.logits_processor(self.lm_head, hidden_states)

1472 - return logits

1473 -

1474 - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

1475 - # Params for weights, fp8 weight scales, fp8 activation scales

1476 - # (param_name, weight_name, expert_id, shard_id)

1477 - return fused_moe_make_expert_params_mapping(

1478 - self,

1479 - ckpt_gate_proj_name="gate_proj",

1480 - ckpt_down_proj_name="down_proj",

1481 - ckpt_up_proj_name="up_proj",

1482 - num_experts=self.config.n_routed_experts,

1483 - num_redundant_experts=0,

1484 - )

1485 -

1329 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1330 rocm_aiter_moe_shared_expert_enabled = (

1331 rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()

@@ -1703,6 +1546,178 @@ class DeepseekV2ForCausalLM(

1546 return loaded_params

1547

1548

1549 +class DeepseekV2MixtureOfExperts(MixtureOfExperts):

1550 + moe_mlp_layers: list[DeepseekV2MoE]

1551 + """

1552 + List of MoE MLP layers in the model.

1553 + """

1554 +

1555 + def extract_moe_parameters(self, example_moe: DeepseekV2MoE | None):

1556 + if example_moe is None:

1557 + self.num_moe_layers = 0

1558 + self.num_expert_groups = 0

1559 + self.num_logical_experts = 0

1560 + self.num_physical_experts = 0

1561 + self.num_local_physical_experts = 0

1562 + self.num_routed_experts = 0

1563 + self.num_shared_experts = 0

1564 + self.num_redundant_experts = 0

1565 + logger.warning("DeepSeekV2: No DeepseekV2MoE layer found in model.layers.")

1566 + else:

1567 + self.num_logical_experts = example_moe.n_logical_experts

1568 + self.num_physical_experts = example_moe.n_physical_experts

1569 + self.num_local_physical_experts = example_moe.n_local_physical_experts

1570 + self.num_routed_experts = example_moe.n_routed_experts

1571 + self.num_shared_experts = example_moe.n_shared_experts

1572 + self.num_redundant_experts = example_moe.n_redundant_experts

1573 +

1574 + def update_physical_experts_metadata(

1575 + self,

1576 + num_physical_experts: int,

1577 + num_local_physical_experts: int,

1578 + ) -> None:

1579 + assert self.num_local_physical_experts == num_local_physical_experts

1580 + self.num_physical_experts = num_physical_experts

1581 + self.num_local_physical_experts = num_local_physical_experts

1582 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts

1583 + for moe in self.moe_mlp_layers:

1584 + moe.n_local_physical_experts = num_local_physical_experts

1585 + moe.n_physical_experts = num_physical_experts

1586 + moe.n_redundant_experts = self.num_redundant_experts

1587 + moe.experts.update_expert_map()

1588 +

1589 +

1590 +class DeepseekV2ForCausalLM(

1591 + nn.Module,

1592 + SupportsPP,

1593 + DeepseekV2MixtureOfExperts,

1594 + SupportsLoRA,

1595 + SupportsEagle,

1596 + SupportsEagle3,

1597 +):

1598 + packed_modules_mapping = {

1599 + "gate_up_proj": ["gate_proj", "up_proj"],

1600 + }

1601 + model_cls = DeepseekV2Model

1602 +

1603 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1604 + super().__init__()

1605 + config = vllm_config.model_config.hf_config

1606 + quant_config = vllm_config.quant_config

1607 + self.config = config

1608 + self.quant_config = quant_config

1609 +

1610 + qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)

1611 + qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)

1612 + self.use_mha = config.model_type == "deepseek" or all(

1613 + dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)

1614 + )

1615 +

1616 + if self.use_mha:

1617 + self.packed_modules_mapping["qkv_proj"] = ["q_proj", "k_proj", "v_proj"]

1618 +

1619 + # `packed_modules_mapping` needs to be modified before

1620 + # initializing DeepseekV2Model, as it is passed inplace to

1621 + # quantization config init and may be used to select the

1622 + # quant_method for relevant layers during initialization.

1623 + self.fuse_qkv_a_proj = (

1624 + hasattr(config, "q_lora_rank") and config.q_lora_rank is not None

1625 + )

1626 + if self.fuse_qkv_a_proj:

1627 + self.packed_modules_mapping["fused_qkv_a_proj"] = [

1628 + "q_a_proj",

1629 + "kv_a_proj_with_mqa",

1630 + ]

1631 +

1632 + self.model = self.model_cls(

1633 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

1634 + )

1635 + if get_pp_group().is_last_rank:

1636 + self.lm_head = ParallelLMHead(

1637 + config.vocab_size,

1638 + config.hidden_size,

1639 + quant_config=quant_config,

1640 + prefix=maybe_prefix(prefix, "lm_head"),

1641 + )

1642 + else:

1643 + self.lm_head = PPMissingLayer()

1644 + self.logits_processor = LogitsProcessor(config.vocab_size)

1645 + self.make_empty_intermediate_tensors = (

1646 + self.model.make_empty_intermediate_tensors

1647 + )

1648 + # Set MoE hyperparameters

1649 + self.num_moe_layers = (

1650 + self.config.num_hidden_layers - self.config.first_k_dense_replace

1651 + )

1652 + self.set_moe_parameters()

1653 +

1654 + def set_moe_parameters(self):

1655 + self.expert_weights = []

1656 +

1657 + self.num_expert_groups = getattr(self.config, "n_group", 1)

1658 +

1659 + self.moe_layers = []

1660 + self.moe_mlp_layers = []

1661 + example_moe = None

1662 + for layer in self.model.layers:

1663 + if isinstance(layer, PPMissingLayer):

1664 + continue

1665 +

1666 + assert isinstance(layer, DeepseekV2DecoderLayer)

1667 + if isinstance(layer.mlp, DeepseekV2MoE):

1668 + # Pick last one layer since the first ones may be dense layers.

1669 + example_moe = layer.mlp

1670 + self.moe_mlp_layers.append(layer.mlp)

1671 + self.moe_layers.append(layer.mlp.experts)

1672 +

1673 + self.extract_moe_parameters(example_moe)

1674 +

1675 + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:

1676 + self.model.aux_hidden_state_layers = layers

1677 +

1678 + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:

1679 + num_layers = len(self.model.layers)

1680 + return (2, num_layers // 2, num_layers - 3)

1681 +

1682 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1683 + return self.model.embed_input_ids(input_ids)

1684 +

1685 + def forward(

1686 + self,

1687 + input_ids: torch.Tensor | None,

1688 + positions: torch.Tensor,

1689 + intermediate_tensors: IntermediateTensors | None = None,

1690 + inputs_embeds: torch.Tensor | None = None,

1691 + ) -> torch.Tensor | IntermediateTensors:

1692 + hidden_states = self.model(

1693 + input_ids, positions, intermediate_tensors, inputs_embeds

1694 + )

1695 + return hidden_states

1696 +

1697 + def compute_logits(

1698 + self,

1699 + hidden_states: torch.Tensor,

1700 + ) -> torch.Tensor | None:

1701 + logits = self.logits_processor(self.lm_head, hidden_states)

1702 + return logits

1703 +

1704 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

1705 + # Params for weights, fp8 weight scales, fp8 activation scales

1706 + # (param_name, weight_name, expert_id, shard_id)

1707 + return fused_moe_make_expert_params_mapping(

1708 + self,

1709 + ckpt_gate_proj_name="gate_proj",

1710 + ckpt_down_proj_name="down_proj",

1711 + ckpt_up_proj_name="up_proj",

1712 + num_experts=self.config.n_routed_experts,

1713 + num_redundant_experts=0,

1714 + )

1715 +

1716 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1717 + loader = AutoWeightsLoader(self)

1718 + return loader.load_weights(weights)

1719 +

1720 +

1721 class DeepseekForCausalLM(DeepseekV2ForCausalLM):

1722 pass

1723

@@ -1726,6 +1741,8 @@ def get_spec_layer_idx_from_weight_name(

1741 ):

1742 layer_idx = config.num_hidden_layers

1743 for i in range(config.num_nextn_predict_layers):

1729 - if weight_name.startswith(f"model.layers.{layer_idx + i}."):

1744 + if weight_name.startswith(

1745 + f"model.layers.{layer_idx + i}."

1746 + ) or weight_name.startswith(f"layers.{layer_idx + i}."):

1747 return layer_idx + i

1748 return None

[Model] use AutoWeightsLoader for DeepSeekV2 (#41706)