ffee741626 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -77,6 +77,7 @@ from vllm.transformers_utils.configs.AXK1 import AXK1Config

77

78 from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP

79 from .utils import (

80 + AutoWeightsLoader,

81 PPMissingLayer,

82 is_pp_missing_parameter,

83 make_empty_intermediate_tensors_factory,

@@ -727,6 +728,13 @@ class AXK1Model(nn.Module):

728 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

729 ["hidden_states", "residual"], config.hidden_size

730 )

731 + self.use_mha = all(

732 + dim == 0 for dim in (config.qk_nope_head_dim, config.qk_rope_head_dim)

733 + )

734 + self.fuse_qkv_a_proj = config.q_lora_rank is not None

735 + self.num_redundant_experts = (

736 + vllm_config.parallel_config.eplb_config.num_redundant_experts

737 + )

738

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

740 return self.embed_tokens(input_ids)

@@ -776,158 +784,6 @@ class AXK1Model(nn.Module):

784 hidden_states, _ = self.norm(hidden_states, residual)

785 return hidden_states

786

779 -

780 -class AXK1MixtureOfExperts(MixtureOfExperts):

781 - moe_mlp_layers: list[AXK1MoE]

782 - """

783 - List of MoE MLP layers in the model.

784 - """

785 -

786 - def extract_moe_parameters(self, example_moe: AXK1MoE | None):

787 - if example_moe is None:

788 - self.num_moe_layers = 0

789 - self.num_expert_groups = 0

790 - self.num_logical_experts = 0

791 - self.num_physical_experts = 0

792 - self.num_local_physical_experts = 0

793 - self.num_routed_experts = 0

794 - self.num_shared_experts = 0

795 - self.num_redundant_experts = 0

796 - logger.warning("AXK1: No AXK1MoE layer found in model.layers.")

797 - else:

798 - self.num_logical_experts = example_moe.n_logical_experts

799 - self.num_physical_experts = example_moe.n_physical_experts

800 - self.num_local_physical_experts = example_moe.n_local_physical_experts

801 - self.num_routed_experts = example_moe.n_routed_experts

802 - self.num_shared_experts = example_moe.n_shared_experts

803 - self.num_redundant_experts = example_moe.n_redundant_experts

804 -

805 - def update_physical_experts_metadata(

806 - self,

807 - num_physical_experts: int,

808 - num_local_physical_experts: int,

809 - ) -> None:

810 - assert self.num_local_physical_experts == num_local_physical_experts

811 - self.num_physical_experts = num_physical_experts

812 - self.num_local_physical_experts = num_local_physical_experts

813 - self.num_redundant_experts = num_physical_experts - self.num_logical_experts

814 - for moe in self.moe_mlp_layers:

815 - moe.n_local_physical_experts = num_local_physical_experts

816 - moe.n_physical_experts = num_physical_experts

817 - moe.n_redundant_experts = self.num_redundant_experts

818 - moe.experts.update_expert_map()

819 -

820 -

821 -class AXK1ForCausalLM(

822 - nn.Module, SupportsPP, AXK1MixtureOfExperts, SupportsLoRA, SupportsEagle

823 -):

824 - packed_modules_mapping = {

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

826 - }

827 - model_cls = AXK1Model

828 -

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

830 - super().__init__()

831 - config: AXK1Config = vllm_config.model_config.hf_config

832 - quant_config = vllm_config.quant_config

833 - self.config = config

834 - self.quant_config = quant_config

835 -

836 - qk_nope_head_dim = config.qk_nope_head_dim

837 - qk_rope_head_dim = config.qk_rope_head_dim

838 - self.use_mha = all(dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim))

839 -

840 - if self.use_mha:

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

842 -

843 - # `packed_modules_mapping` needs to be modified before

844 - # initializing AXK1Model, as it is passed inplace to

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

846 - # quant_method for relevant layers during initialization.

847 - self.fuse_qkv_a_proj = config.q_lora_rank is not None

848 - if self.fuse_qkv_a_proj:

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

850 - "q_a_proj",

851 - "kv_a_proj_with_mqa",

852 - ]

853 -

854 - self.model = self.model_cls(

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

856 - )

857 - if get_pp_group().is_last_rank:

858 - self.lm_head = ParallelLMHead(

859 - config.vocab_size,

860 - config.hidden_size,

861 - quant_config=quant_config,

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

863 - )

864 - else:

865 - self.lm_head = PPMissingLayer()

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

867 - self.make_empty_intermediate_tensors = (

868 - self.model.make_empty_intermediate_tensors

869 - )

870 - # Set MoE hyperparameters

871 - self.num_moe_layers = (

872 - self.config.num_hidden_layers - self.config.first_k_dense_replace

873 - )

874 - self.set_moe_parameters()

875 -

876 - def set_moe_parameters(self):

877 - self.expert_weights = []

878 -

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

880 -

881 - self.moe_layers = []

882 - self.moe_mlp_layers = []

883 - example_moe = None

884 - for layer in self.model.layers:

885 - if isinstance(layer, PPMissingLayer):

886 - continue

887 -

888 - assert isinstance(layer, AXK1DecoderLayer)

889 - if isinstance(layer.mlp, AXK1MoE):

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

891 - example_moe = layer.mlp

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

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

894 -

895 - self.extract_moe_parameters(example_moe)

896 -

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

898 - return self.model.embed_input_ids(input_ids)

899 -

900 - def forward(

901 - self,

902 - input_ids: torch.Tensor | None,

903 - positions: torch.Tensor,

904 - intermediate_tensors: IntermediateTensors | None = None,

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

906 - ) -> torch.Tensor | IntermediateTensors:

907 - hidden_states = self.model(

908 - input_ids, positions, intermediate_tensors, inputs_embeds

909 - )

910 - return hidden_states

911 -

912 - def compute_logits(

913 - self,

914 - hidden_states: torch.Tensor,

915 - ) -> torch.Tensor | None:

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

917 - return logits

918 -

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

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

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

922 - return fused_moe_make_expert_params_mapping(

923 - self,

924 - ckpt_gate_proj_name="gate_proj",

925 - ckpt_down_proj_name="down_proj",

926 - ckpt_up_proj_name="up_proj",

927 - num_experts=self.config.n_routed_experts,

928 - num_redundant_experts=0,

929 - )

930 -

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

788 rocm_aiter_moe_shared_expert_enabled = (

789 rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()

@@ -1135,12 +991,170 @@ class AXK1ForCausalLM(

991 return loaded_params

992

993

994 +class AXK1MixtureOfExperts(MixtureOfExperts):

995 + moe_mlp_layers: list[AXK1MoE]

996 + """

997 + List of MoE MLP layers in the model.

998 + """

999 +

1000 + def extract_moe_parameters(self, example_moe: AXK1MoE | None):

1001 + if example_moe is None:

1002 + self.num_moe_layers = 0

1003 + self.num_expert_groups = 0

1004 + self.num_logical_experts = 0

1005 + self.num_physical_experts = 0

1006 + self.num_local_physical_experts = 0

1007 + self.num_routed_experts = 0

1008 + self.num_shared_experts = 0

1009 + self.num_redundant_experts = 0

1010 + logger.warning("AXK1: No AXK1MoE layer found in model.layers.")

1011 + else:

1012 + self.num_logical_experts = example_moe.n_logical_experts

1013 + self.num_physical_experts = example_moe.n_physical_experts

1014 + self.num_local_physical_experts = example_moe.n_local_physical_experts

1015 + self.num_routed_experts = example_moe.n_routed_experts

1016 + self.num_shared_experts = example_moe.n_shared_experts

1017 + self.num_redundant_experts = example_moe.n_redundant_experts

1018 +

1019 + def update_physical_experts_metadata(

1020 + self,

1021 + num_physical_experts: int,

1022 + num_local_physical_experts: int,

1023 + ) -> None:

1024 + assert self.num_local_physical_experts == num_local_physical_experts

1025 + self.num_physical_experts = num_physical_experts

1026 + self.num_local_physical_experts = num_local_physical_experts

1027 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts

1028 + for moe in self.moe_mlp_layers:

1029 + moe.n_local_physical_experts = num_local_physical_experts

1030 + moe.n_physical_experts = num_physical_experts

1031 + moe.n_redundant_experts = self.num_redundant_experts

1032 + moe.experts.update_expert_map()

1033 +

1034 +

1035 +class AXK1ForCausalLM(

1036 + nn.Module, SupportsPP, AXK1MixtureOfExperts, SupportsLoRA, SupportsEagle

1037 +):

1038 + packed_modules_mapping = {

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

1040 + }

1041 + model_cls = AXK1Model

1042 +

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

1044 + super().__init__()

1045 + config: AXK1Config = vllm_config.model_config.hf_config

1046 + quant_config = vllm_config.quant_config

1047 + self.config = config

1048 + self.quant_config = quant_config

1049 +

1050 + qk_nope_head_dim = config.qk_nope_head_dim

1051 + qk_rope_head_dim = config.qk_rope_head_dim

1052 + self.use_mha = all(dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim))

1053 +

1054 + if self.use_mha:

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

1056 +

1057 + # `packed_modules_mapping` needs to be modified before

1058 + # initializing AXK1Model, as it is passed inplace to

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

1060 + # quant_method for relevant layers during initialization.

1061 + self.fuse_qkv_a_proj = config.q_lora_rank is not None

1062 + if self.fuse_qkv_a_proj:

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

1064 + "q_a_proj",

1065 + "kv_a_proj_with_mqa",

1066 + ]

1067 +

1068 + self.model = self.model_cls(

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

1070 + )

1071 + if get_pp_group().is_last_rank:

1072 + self.lm_head = ParallelLMHead(

1073 + config.vocab_size,

1074 + config.hidden_size,

1075 + quant_config=quant_config,

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

1077 + )

1078 + else:

1079 + self.lm_head = PPMissingLayer()

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

1081 + self.make_empty_intermediate_tensors = (

1082 + self.model.make_empty_intermediate_tensors

1083 + )

1084 + # Set MoE hyperparameters

1085 + self.num_moe_layers = (

1086 + self.config.num_hidden_layers - self.config.first_k_dense_replace

1087 + )

1088 + self.set_moe_parameters()

1089 +

1090 + def set_moe_parameters(self):

1091 + self.expert_weights = []

1092 +

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

1094 +

1095 + self.moe_layers = []

1096 + self.moe_mlp_layers = []

1097 + example_moe = None

1098 + for layer in self.model.layers:

1099 + if isinstance(layer, PPMissingLayer):

1100 + continue

1101 +

1102 + assert isinstance(layer, AXK1DecoderLayer)

1103 + if isinstance(layer.mlp, AXK1MoE):

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

1105 + example_moe = layer.mlp

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

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

1108 +

1109 + self.extract_moe_parameters(example_moe)

1110 +

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

1112 + return self.model.embed_input_ids(input_ids)

1113 +

1114 + def forward(

1115 + self,

1116 + input_ids: torch.Tensor | None,

1117 + positions: torch.Tensor,

1118 + intermediate_tensors: IntermediateTensors | None = None,

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

1120 + ) -> torch.Tensor | IntermediateTensors:

1121 + hidden_states = self.model(

1122 + input_ids, positions, intermediate_tensors, inputs_embeds

1123 + )

1124 + return hidden_states

1125 +

1126 + def compute_logits(

1127 + self,

1128 + hidden_states: torch.Tensor,

1129 + ) -> torch.Tensor | None:

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

1131 + return logits

1132 +

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

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

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

1136 + return fused_moe_make_expert_params_mapping(

1137 + self,

1138 + ckpt_gate_proj_name="gate_proj",

1139 + ckpt_down_proj_name="down_proj",

1140 + ckpt_up_proj_name="up_proj",

1141 + num_experts=self.config.n_routed_experts,

1142 + num_redundant_experts=0,

1143 + )

1144 +

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

1146 + loader = AutoWeightsLoader(self)

1147 + return loader.load_weights(weights)

1148 +

1149 +

1150 def get_spec_layer_idx_from_weight_name(

1151 config: AXK1Config, weight_name: str

1152 ) -> int | None:

1153 if config.num_nextn_predict_layers and config.num_nextn_predict_layers > 0:

1154 layer_idx = config.num_hidden_layers

1155 for i in range(config.num_nextn_predict_layers):

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

1156 + if weight_name.startswith(

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

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

1159 return layer_idx + i

1160 return None

[Model] Use AutoWeightsLoader for AXK1 (#41901)