diff --git a/docs/supported_models/text_generation/diffusion_language_models.md b/docs/supported_models/text_generation/diffusion_language_models.md index 7dbb4828b..2faa0206e 100644 --- a/docs/supported_models/text_generation/diffusion_language_models.md +++ b/docs/supported_models/text_generation/diffusion_language_models.md @@ -107,5 +107,3 @@ Below the supported models are summarized in a table. | Model Family | Example Model | Description | | -------------------------- | ---------------------------- | ---------------------------------------------------------------------------------------------------- | | **LLaDA2.0 (mini, flash)** | `inclusionAI/LLaDA2.0-flash` | LLaDA2.0-flash is a diffusion language model featuring a 100B Mixture-of-Experts (MoE) architecture. | -| **SDAR (JetLM)** | `JetLM/SDAR-8B-Chat` | SDAR series diffusion language model (Chat), dense architecture. | -| **SDAR (JetLM)** | `JetLM/SDAR-30B-A3B-Chat` | SDAR series diffusion language model (Chat), MoE architecture. | diff --git a/python/sglang/srt/dllm/config.py b/python/sglang/srt/dllm/config.py index edd204926..ce601cbdc 100644 --- a/python/sglang/srt/dllm/config.py +++ b/python/sglang/srt/dllm/config.py @@ -31,19 +31,14 @@ class DllmConfig: model_path=server_args.model_path, model_revision=server_args.revision, ) - DLLM_PARAMS = { - "LLaDA2MoeModelLM": {"block_size": 32, "mask_id": 156895}, - "SDARForCausalLM": {"block_size": 4, "mask_id": 151669}, - "SDARMoeForCausalLM": {"block_size": 4, "mask_id": 151669}, - } - arch = model_config.hf_config.architectures[0] - if arch in DLLM_PARAMS: - params = DLLM_PARAMS[arch] - block_size = params["block_size"] - mask_id = params["mask_id"] + if model_config.hf_config.architectures[0] == "LLaDA2MoeModelLM": + block_size = 32 + mask_id = 156895 else: - raise RuntimeError(f"Unknown diffusion LLM: {arch}") + raise RuntimeError( + f"Unknown diffusion LLM: {model_config.hf_config.architectures[0]}" + ) max_running_requests = ( 1 diff --git a/python/sglang/srt/models/sdar.py b/python/sglang/srt/models/sdar.py deleted file mode 100644 index 70ab59a48..000000000 --- a/python/sglang/srt/models/sdar.py +++ /dev/null @@ -1,589 +0,0 @@ -# coding=utf-8 -""" -SGLang SDARModelLM (block diffusion / dLLM-style forward). -""" - -import logging -from typing import Iterable, Optional, Tuple, Union - -import torch -from torch import nn -from transformers import PretrainedConfig - -from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size -from sglang.srt.layers.activation import SiluAndMul -from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes -from sglang.srt.layers.dp_attention import ( - get_attention_tp_rank, - get_attention_tp_size, - is_dp_attention_enabled, -) -from sglang.srt.layers.layernorm import RMSNorm -from sglang.srt.layers.linear import ( - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from sglang.srt.layers.logits_processor import LogitsProcessor -from sglang.srt.layers.quantization.base_config import QuantizationConfig -from sglang.srt.layers.radix_attention import AttentionType, RadixAttention -from sglang.srt.layers.rotary_embedding import get_rope -from sglang.srt.layers.utils import PPMissingLayer, get_layer_id -from sglang.srt.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -from sglang.srt.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) -from sglang.srt.models.utils import ( - apply_qk_norm, - create_fused_set_kv_buffer_arg, - enable_fused_set_kv_buffer, -) -from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import add_prefix, is_cuda, make_layers - -logger = logging.getLogger(__name__) -_is_cuda = is_cuda() - - -class SDARMLP(nn.Module): - def __init__( - self, - config: PretrainedConfig, - quant_config=None, - reduce_results: bool = True, - prefix: str = "", - ): - super().__init__() - self.gate_up_proj = MergedColumnParallelLinear( - config.hidden_size, - [config.intermediate_size] * 2, - bias=False, - quant_config=quant_config, - prefix=add_prefix("gate_up_proj", prefix), - ) - self.down_proj = RowParallelLinear( - config.intermediate_size, - config.hidden_size, - bias=False, - reduce_results=reduce_results, - quant_config=quant_config, - prefix=add_prefix("down_proj", prefix), - ) - self.act_fn = SiluAndMul() - - def forward(self, hidden_states: torch.Tensor, use_reduce_scatter: bool = False): - gate_up, _ = self.gate_up_proj(hidden_states) - hidden_states = self.act_fn(gate_up) - hidden_states, _ = self.down_proj( - hidden_states, skip_all_reduce=use_reduce_scatter - ) - return hidden_states - - -class SDARAttention(nn.Module): - def __init__( - self, - config: PretrainedConfig, - layer_id: int, - quant_config=None, - reduce_results: bool = True, - prefix: str = "", - alt_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__() - self.layer_id = layer_id - self.hidden_size = config.hidden_size - self.total_num_heads = config.num_attention_heads - self.tp_size = get_tensor_model_parallel_world_size() - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() - - assert self.total_num_heads % attn_tp_size == 0 - self.num_heads = self.total_num_heads // attn_tp_size - self.total_num_kv_heads = config.num_key_value_heads - if self.total_num_kv_heads >= attn_tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % attn_tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert attn_tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) - - self.head_dim = getattr( - config, "head_dim", self.hidden_size // self.total_num_heads - ) - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scale = self.head_dim**-0.5 - - self.qkv_proj = QKVParallelLinear( - self.hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=getattr(config, "attention_bias", False), - quant_config=quant_config, - tp_rank=attn_tp_rank, - tp_size=attn_tp_size, - prefix=add_prefix("qkv_proj", prefix), - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - self.hidden_size, - bias=getattr(config, "attention_bias", False), - quant_config=quant_config, - reduce_results=reduce_results, - tp_rank=attn_tp_rank, - tp_size=attn_tp_size, - prefix=add_prefix("o_proj", prefix), - ) - - self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - - rope_theta = getattr(config, "rope_theta", 10000.0) - rope_scaling = getattr(config, "rope_scaling", None) - max_pos = getattr(config, "max_position_embeddings", 32768) - self.rotary_dim = self.head_dim - self.rotary_emb = get_rope( - self.head_dim, - rotary_dim=self.rotary_dim, - max_position=max_pos, - base=rope_theta, - rope_scaling=rope_scaling, - ) - - # RadixAttention: ENCODER_ONLY lets ForwardBatch provide non-causal / block masks (dLLM) - # NOTE: this is the key change vs AR Llama-style DECODER self-attn. - self.attn = RadixAttention( - self.num_heads, - self.head_dim, - self.scale, - num_kv_heads=self.num_kv_heads, - layer_id=layer_id, - attn_type=AttentionType.ENCODER_ONLY, - prefix=add_prefix("attn", prefix), - ) - self.alt_stream = alt_stream - - def forward_prepare_native(self, positions, hidden_states, forward_batch): - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = apply_qk_norm( - q=q, - k=k, - q_norm=self.q_norm, - k_norm=self.k_norm, - head_dim=self.head_dim, - alt_stream=self.alt_stream, - ) - q, k = self.rotary_emb( - positions, - q, - k, - fused_set_kv_buffer_arg=( - create_fused_set_kv_buffer_arg( - value=v, - layer=self.attn, - forward_batch=forward_batch, - ) - if enable_fused_set_kv_buffer(forward_batch) - else None - ), - ) - return q, k, v - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - ): - if get_global_server_args().rl_on_policy_target is not None: - hidden_states = hidden_states.bfloat16() - - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = apply_qk_norm( - q=q, - k=k, - q_norm=self.q_norm, - k_norm=self.k_norm, - head_dim=self.head_dim, - alt_stream=self.alt_stream, - ) - q, k = self.rotary_emb( - positions, - q, - k, - fused_set_kv_buffer_arg=( - create_fused_set_kv_buffer_arg( - value=v, - layer=self.attn, - forward_batch=forward_batch, - ) - if enable_fused_set_kv_buffer(forward_batch) - else None - ), - ) - - if get_global_server_args().rl_on_policy_target is not None: - q = q.to(torch.bfloat16) - k = k.to(torch.bfloat16) - - context_layer = self.attn( - q, - k, - v, - forward_batch, - save_kv_cache=not enable_fused_set_kv_buffer(forward_batch), - ) - out, _ = self.o_proj(context_layer) - return out - - -class SDARBlock(nn.Module): - def __init__( - self, - config: PretrainedConfig, - layer_id: int, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - alt_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__() - self.hidden_size = config.hidden_size - self.layer_id = layer_id - - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_global_server_args().rl_on_policy_target is not None - else {} - ) - self.input_layernorm = RMSNorm( - self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs - ) - self.post_attention_layernorm = RMSNorm( - self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs - ) - - self.self_attn = SDARAttention( - layer_id=layer_id, - config=config, - quant_config=quant_config, - reduce_results=False, - prefix=add_prefix("self_attn", prefix), - alt_stream=alt_stream, - ) - - self.mlp = SDARMLP( - config=config, - quant_config=quant_config, - reduce_results=True, - prefix=add_prefix("mlp", prefix), - ) - - self.layer_scatter_modes = LayerScatterModes.init_new( - layer_id=layer_id, - num_layers=config.num_hidden_layers, - is_layer_sparse=False, - is_previous_layer_sparse=False, - is_next_layer_sparse=False, - ) - self.layer_communicator = LayerCommunicator( - layer_scatter_modes=self.layer_scatter_modes, - input_layernorm=self.input_layernorm, - post_attention_layernorm=self.post_attention_layernorm, - allow_reduce_scatter=True, - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - residual: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor]: - hidden_states, residual = self.layer_communicator.prepare_attn( - hidden_states, - residual, - forward_batch, - ) - - if hidden_states.shape[0] != 0: - hidden_states = self.self_attn( - positions=positions, - hidden_states=hidden_states, - forward_batch=forward_batch, - ) - - hidden_states, residual = self.layer_communicator.prepare_mlp( - hidden_states, - residual, - forward_batch, - ) - - use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( - forward_batch - ) - hidden_states = self.mlp(hidden_states, use_reduce_scatter=use_reduce_scatter) - - hidden_states, residual = self.layer_communicator.postprocess_layer( - hidden_states, residual, forward_batch - ) - - return hidden_states, residual - - -class SDARModel(nn.Module): - def __init__( - self, - config: PretrainedConfig, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - alt_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__() - self.config = config - self.vocab_size = config.vocab_size - self.embed_dim = config.hidden_size - self.pp_group = get_pp_group() - - if self.pp_group.is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - self.embed_dim, - quant_config=quant_config, - use_attn_tp_group=is_dp_attention_enabled(), - prefix=add_prefix("embed_tokens", prefix), - ) - else: - self.embed_tokens = PPMissingLayer() - - self.layers, self.start_layer, self.end_layer = make_layers( - config.num_hidden_layers, - lambda idx, prefix: SDARBlock( - layer_id=idx, - config=config, - quant_config=quant_config, - prefix=prefix, - alt_stream=alt_stream, - ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, - prefix=add_prefix("layers", prefix), - ) - if self.pp_group.is_last_rank: - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_global_server_args().rl_on_policy_target is not None - else {} - ) - self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps, **norm_kwargs) - else: - self.norm = PPMissingLayer(return_tuple=True) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - input_embeds: torch.Tensor = None, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - ) -> Union[torch.Tensor, PPProxyTensors]: - if self.pp_group.is_first_rank: - hidden_states = ( - self.embed_tokens(input_ids) if input_embeds is None else input_embeds - ) - residual = None - else: - assert pp_proxy_tensors is not None - hidden_states = pp_proxy_tensors["hidden_states"] - residual = pp_proxy_tensors.get("residual", None) - - for i in range(self.start_layer, self.end_layer): - layer = self.layers[i] - hidden_states, residual = layer( - positions, hidden_states, forward_batch, residual - ) - - if not self.pp_group.is_last_rank: - return PPProxyTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - else: - if not forward_batch.forward_mode.is_idle(): - hidden_states, residual = self.norm(hidden_states, residual) - return hidden_states - - -class SDARForCausalLM(nn.Module): - def __init__( - self, - config: PretrainedConfig, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - ): - super().__init__() - self.pp_group = get_pp_group() - assert self.pp_group.world_size == 1, ( - f"SDARMoeForCausalLM does not support pipeline parallel (pp_size={self.pp_group.world_size}). " - "Please set pp_size=1." - ) - - self.config = config - self.quant_config = quant_config - alt_stream = torch.cuda.Stream() if _is_cuda else None - - self.model = SDARModel( - config, - quant_config=quant_config, - prefix=add_prefix("model", ""), - alt_stream=alt_stream, - ) - - if self.pp_group.is_last_rank: - tp_size = get_tensor_model_parallel_world_size() - if ( - self.pp_group.world_size == 1 - and config.tie_word_embeddings - and tp_size == 1 - ): - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - use_attn_tp_group=get_global_server_args().enable_dp_lm_head, - prefix=add_prefix("lm_head", prefix), - ) - else: - self.lm_head = PPMissingLayer() - - self.logits_processor = LogitsProcessor(config, return_full_logits=True) - - @property - def start_layer(self): - return self.model.start_layer - - @property - def end_layer(self): - return self.model.end_layer - - @torch.no_grad() - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - input_embeds: Optional[torch.Tensor] = None, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - ) -> torch.Tensor: - hidden_states = self.model( - input_ids=input_ids, - positions=positions, - forward_batch=forward_batch, - input_embeds=input_embeds, - pp_proxy_tensors=pp_proxy_tensors, - ) - if self.pp_group.is_last_rank: - return self.logits_processor( - input_ids, hidden_states, self.lm_head, forward_batch - ) - else: - return hidden_states - - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - params_dict = dict(self.named_parameters()) - for name, loaded_weight in weights: - if not name.startswith("model.") and ( - name.startswith("layers.") - or name.startswith("embed_tokens.") - or name.startswith("norm.") - ): - name = add_prefix(name, "model") - - if name == "model.embed_tokens.weight": - if self.pp_group.is_last_rank and self.config.tie_word_embeddings: - param = params_dict["lm_head.weight"] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - - layer_id = get_layer_id(name) - if ( - layer_id is not None - and hasattr(self.model, "start_layer") - and ( - layer_id < self.model.start_layer - or layer_id >= self.model.end_layer - ) - ): - continue - - if "rotary_emb.inv_freq" in name or "projector" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - if name.startswith("model.vision_tower") and name not in params_dict: - continue - if "scale" in name: - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if name in params_dict.keys(): - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - else: - logger.warning(f"Parameter {name} not found in params_dict") - - -EntryClass = SDARForCausalLM diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py deleted file mode 100644 index a1b20f67c..000000000 --- a/python/sglang/srt/models/sdar_moe.py +++ /dev/null @@ -1,746 +0,0 @@ -# coding=utf-8 -""" -SGLang SDARMoeModelLM (block diffusion / dLLM-style forward) with MoE MLP. -""" - -import logging -from typing import Iterable, Optional, Tuple, Union - -import torch -from torch import nn -from transformers import PretrainedConfig - -from sglang.srt.distributed import ( - get_moe_expert_parallel_world_size, - get_pp_group, - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, -) -from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder -from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation -from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo -from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes -from sglang.srt.layers.dp_attention import ( - get_attention_tp_rank, - get_attention_tp_size, - is_dp_attention_enabled, -) -from sglang.srt.layers.layernorm import RMSNorm -from sglang.srt.layers.linear import ( - QKVParallelLinear, - ReplicatedLinear, - RowParallelLinear, -) -from sglang.srt.layers.logits_processor import LogitsProcessor -from sglang.srt.layers.moe import ( - get_moe_a2a_backend, - should_use_flashinfer_cutlass_moe_fp4_allgather, -) -from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class -from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE -from sglang.srt.layers.moe.topk import TopK -from sglang.srt.layers.moe.utils import ( - RoutingMethodType, - filter_moe_weight_param_global_expert, -) -from sglang.srt.layers.quantization.base_config import QuantizationConfig -from sglang.srt.layers.radix_attention import AttentionType, RadixAttention -from sglang.srt.layers.rotary_embedding import get_rope -from sglang.srt.layers.utils import PPMissingLayer, get_layer_id -from sglang.srt.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -from sglang.srt.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) -from sglang.srt.models.utils import ( - apply_qk_norm, - create_fused_set_kv_buffer_arg, - enable_fused_set_kv_buffer, -) -from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers - -logger = logging.getLogger(__name__) -_is_cuda = is_cuda() - - -class SDARMoeSparseMoeBlock(nn.Module): - """ - Qwen3MoE-style sparse MoE block: - - gate: ReplicatedLinear(hidden, num_experts) - - topk routing: TopK - - experts: get_moe_impl_class(quant_config)(...) - """ - - def __init__( - self, - layer_id: int, - config: PretrainedConfig, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - ): - super().__init__() - self.layer_id = layer_id - self.tp_size = get_tensor_model_parallel_world_size() - - if self.tp_size > config.num_experts: - raise ValueError( - f"Tensor parallel size {self.tp_size} > num_experts {config.num_experts}." - ) - - self.topk = TopK( - top_k=config.num_experts_per_tok, - renormalize=config.norm_topk_prob, - use_grouped_topk=False, - layer_id=layer_id, - ) - - self.experts = get_moe_impl_class(quant_config)( - num_experts=config.num_experts - + get_global_server_args().ep_num_redundant_experts, - top_k=config.num_experts_per_tok, - layer_id=layer_id, - hidden_size=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - quant_config=quant_config, - prefix=add_prefix("experts", prefix), - routing_method_type=RoutingMethodType.Renormalize, - ) - - self.gate = ReplicatedLinear( - config.hidden_size, - config.num_experts, - bias=False, - quant_config=None, - prefix=add_prefix("gate", prefix), - ) - - # Deepep / FuseEP support - if get_moe_a2a_backend().is_deepep(): - self.ep_size = get_moe_expert_parallel_world_size() - self.num_experts = ( - config.num_experts + get_global_server_args().ep_num_redundant_experts - ) - self.top_k = config.num_experts_per_tok - - def forward( - self, - hidden_states: torch.Tensor, - forward_batch: Optional[ForwardBatch] = None, - should_allreduce_fusion: bool = False, - use_reduce_scatter: bool = False, - ) -> torch.Tensor: - if ( - not get_moe_a2a_backend().is_deepep() - and not get_moe_a2a_backend().is_ascend_fuseep() - ): - return self.forward_normal( - hidden_states, - should_allreduce_fusion=should_allreduce_fusion, - use_reduce_scatter=use_reduce_scatter, - ) - else: - assert forward_batch is not None, "deepep/fuseep MoE needs forward_batch" - return self.forward_deepep(hidden_states, forward_batch) - - def forward_normal( - self, - hidden_states: torch.Tensor, - should_allreduce_fusion: bool = False, - use_reduce_scatter: bool = False, - ) -> torch.Tensor: - num_tokens, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) - - router_logits, _ = self.gate(hidden_states) # (T, E) - topk_output = self.topk(hidden_states, router_logits) - out = self.experts(hidden_states, topk_output) # (T, H) - - # TP all-reduce (unless fused / reduce_scatter / fp4 allgather path) - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - ): - out = tensor_model_parallel_all_reduce(out) - - return out.view(num_tokens, hidden_dim) - - def forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): - if hidden_states.shape[0] > 0: - router_logits, _ = self.gate(hidden_states) - topk_output = self.topk( - hidden_states, - router_logits, - num_token_non_padded=forward_batch.num_token_non_padded, - expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( - layer_id=self.layer_id - ), - ) - else: - topk_output = self.topk.empty_topk_output(hidden_states.device) - - out = self.experts(hidden_states=hidden_states, topk_output=topk_output) - return out - - def get_moe_weights(self): - return [ - p.data - for name, p in self.experts.named_parameters() - if name not in ["correction_bias"] - and filter_moe_weight_param_global_expert( - name, p, self.experts.num_local_experts - ) - ] - - -class SDARMoeAttention(nn.Module): - def __init__( - self, - config: PretrainedConfig, - layer_id: int, - quant_config: Optional[QuantizationConfig] = None, - reduce_results: bool = True, - prefix: str = "", - alt_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__() - self.layer_id = layer_id - self.hidden_size = config.hidden_size - self.total_num_heads = config.num_attention_heads - - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() - - assert self.total_num_heads % attn_tp_size == 0 - self.num_heads = self.total_num_heads // attn_tp_size - - self.total_num_kv_heads = config.num_key_value_heads - if self.total_num_kv_heads >= attn_tp_size: - assert self.total_num_kv_heads % attn_tp_size == 0 - else: - assert attn_tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) - - self.head_dim = getattr( - config, "head_dim", self.hidden_size // self.total_num_heads - ) - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scale = self.head_dim**-0.5 - - self.qkv_proj = QKVParallelLinear( - self.hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=getattr(config, "attention_bias", False), - quant_config=quant_config, - prefix=add_prefix("qkv_proj", prefix), - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - self.hidden_size, - bias=getattr(config, "attention_bias", False), - quant_config=quant_config, - reduce_results=reduce_results, - tp_rank=attn_tp_rank, - tp_size=attn_tp_size, - prefix=add_prefix("o_proj", prefix), - ) - - self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - - rope_theta = getattr(config, "rope_theta", 10000.0) - rope_scaling = getattr(config, "rope_scaling", None) - max_pos = getattr(config, "max_position_embeddings", 32768) - self.rotary_dim = self.head_dim - self.rotary_emb = get_rope( - self.head_dim, - rotary_dim=self.rotary_dim, - max_position=max_pos, - base=rope_theta, - rope_scaling=rope_scaling, - ) - - self.attn = RadixAttention( - self.num_heads, - self.head_dim, - self.scale, - num_kv_heads=self.num_kv_heads, - layer_id=layer_id, - attn_type=AttentionType.ENCODER_ONLY, - prefix=add_prefix("attn", prefix), - ) - self.alt_stream = alt_stream - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - ) -> torch.Tensor: - if get_global_server_args().rl_on_policy_target is not None: - hidden_states = hidden_states.bfloat16() - - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = apply_qk_norm( - q=q, - k=k, - q_norm=self.q_norm, - k_norm=self.k_norm, - head_dim=self.head_dim, - alt_stream=self.alt_stream, - ) - q, k = self.rotary_emb( - positions, - q, - k, - fused_set_kv_buffer_arg=( - create_fused_set_kv_buffer_arg( - value=v, - layer=self.attn, - forward_batch=forward_batch, - ) - if enable_fused_set_kv_buffer(forward_batch) - else None - ), - ) - - if get_global_server_args().rl_on_policy_target is not None: - q = q.to(torch.bfloat16) - k = k.to(torch.bfloat16) - - context = self.attn( - q, - k, - v, - forward_batch, - save_kv_cache=not enable_fused_set_kv_buffer(forward_batch), - ) - out, _ = self.o_proj(context) - return out - - -class SDARMoeBlock(nn.Module): - def __init__( - self, - config: PretrainedConfig, - layer_id: int, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - alt_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.layer_id = layer_id - - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_global_server_args().rl_on_policy_target is not None - else {} - ) - self.input_layernorm = RMSNorm( - self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs - ) - self.post_attention_layernorm = RMSNorm( - self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs - ) - - self.self_attn = SDARMoeAttention( - config=config, - layer_id=layer_id, - quant_config=quant_config, - reduce_results=False, - prefix=add_prefix("self_attn", prefix), - alt_stream=alt_stream, - ) - - self.mlp = SDARMoeSparseMoeBlock( - layer_id=layer_id, - config=config, - quant_config=quant_config, - prefix=add_prefix("mlp", prefix), - ) - - self.layer_scatter_modes = LayerScatterModes.init_new( - layer_id=layer_id, - num_layers=config.num_hidden_layers, - is_layer_sparse=True, - is_previous_layer_sparse=True, - is_next_layer_sparse=True, - ) - - self.layer_communicator = LayerCommunicator( - layer_scatter_modes=self.layer_scatter_modes, - input_layernorm=self.input_layernorm, - post_attention_layernorm=self.post_attention_layernorm, - allow_reduce_scatter=True, - is_last_layer=(layer_id == config.num_hidden_layers - 1), - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - residual: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor]: - - hidden_states, residual = self.layer_communicator.prepare_attn( - hidden_states, residual, forward_batch - ) - - if hidden_states.shape[0] != 0: - hidden_states = self.self_attn( - positions=positions, - hidden_states=hidden_states, - forward_batch=forward_batch, - ) - - hidden_states, residual = self.layer_communicator.prepare_mlp( - hidden_states, residual, forward_batch - ) - - should_allreduce_fusion = ( - self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( - forward_batch - ) - ) - use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( - forward_batch - ) - - hidden_states = self.mlp( - hidden_states, - forward_batch=forward_batch, - should_allreduce_fusion=should_allreduce_fusion, - use_reduce_scatter=use_reduce_scatter, - ) - - if should_allreduce_fusion: - hidden_states._sglang_needs_allreduce_fusion = True - else: - hidden_states, residual = self.layer_communicator.postprocess_layer( - hidden_states, residual, forward_batch - ) - - return hidden_states, residual - - -class SDARMoeModel(nn.Module): - def __init__( - self, - config: PretrainedConfig, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - alt_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__() - self.config = config - self.vocab_size = config.vocab_size - self.embed_dim = config.hidden_size - self.pp_group = get_pp_group() - - if self.pp_group.is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - self.embed_dim, - quant_config=quant_config, - use_attn_tp_group=is_dp_attention_enabled(), - prefix=add_prefix("embed_tokens", prefix), - ) - else: - self.embed_tokens = PPMissingLayer() - - self.layers, self.start_layer, self.end_layer = make_layers( - config.num_hidden_layers, - lambda idx, prefix: SDARMoeBlock( - config=config, - layer_id=idx, - quant_config=quant_config, - prefix=prefix, - alt_stream=alt_stream, - ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, - prefix=add_prefix("layers", prefix), - ) - - if self.pp_group.is_last_rank: - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_global_server_args().rl_on_policy_target is not None - else {} - ) - self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps, **norm_kwargs) - else: - self.norm = PPMissingLayer(return_tuple=True) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - input_embeds: torch.Tensor = None, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - ) -> Union[torch.Tensor, PPProxyTensors]: - if self.pp_group.is_first_rank: - hidden_states = ( - self.embed_tokens(input_ids) if input_embeds is None else input_embeds - ) - residual = None - else: - assert pp_proxy_tensors is not None - hidden_states = pp_proxy_tensors["hidden_states"] - residual = pp_proxy_tensors.get("residual", None) - - for i in range(self.start_layer, self.end_layer): - layer = self.layers[i] - with get_global_expert_distribution_recorder().with_current_layer(i): - hidden_states, residual = layer( - positions, hidden_states, forward_batch, residual - ) - - if not self.pp_group.is_last_rank: - return PPProxyTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - - if not forward_batch.forward_mode.is_idle(): - hidden_states, residual = self.norm(hidden_states, residual) - return hidden_states - - -class SDARMoeForCausalLM(nn.Module): - fall_back_to_pt_during_load = False - - def __init__( - self, - config: PretrainedConfig, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - ): - super().__init__() - self.pp_group = get_pp_group() - assert self.pp_group.world_size == 1, ( - f"SDARMoeForCausalLM does not support pipeline parallel (pp_size={self.pp_group.world_size}). " - "Please set pp_size=1." - ) - - self.pp_group = get_pp_group() - self.config = config - self.quant_config = quant_config - alt_stream = torch.cuda.Stream() if _is_cuda else None - - self.model = SDARMoeModel( - config, - quant_config=quant_config, - prefix=add_prefix("model", ""), - alt_stream=alt_stream, - ) - - if self.pp_group.is_last_rank: - tp_size = get_tensor_model_parallel_world_size() - if ( - self.pp_group.world_size == 1 - and getattr(config, "tie_word_embeddings", False) - and tp_size == 1 - ): - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - use_attn_tp_group=get_global_server_args().enable_dp_lm_head, - prefix=add_prefix("lm_head", prefix), - ) - else: - self.lm_head = PPMissingLayer() - - self.logits_processor = LogitsProcessor(config, return_full_logits=True) - - @property - def start_layer(self): - return self.model.start_layer - - @property - def end_layer(self): - return self.model.end_layer - - @torch.no_grad() - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - input_embeds: Optional[torch.Tensor] = None, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - ) -> torch.Tensor: - hidden_states = self.model( - input_ids=input_ids, - positions=positions, - forward_batch=forward_batch, - input_embeds=input_embeds, - pp_proxy_tensors=pp_proxy_tensors, - ) - if self.pp_group.is_last_rank: - return self.logits_processor( - input_ids, hidden_states, self.lm_head, forward_batch - ) - return hidden_states - - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = [ - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - - expert_params_mapping = FusedMoE.make_expert_params_mapping( - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, - ) - - if not hasattr(self, "_cached_params_dict"): - self._cached_params_dict = dict(self.named_parameters()) - params_dict = self._cached_params_dict - - for name, loaded_weight in weights: - if not name.startswith("model.") and ( - name.startswith("layers.") - or name.startswith("embed_tokens.") - or name.startswith("norm.") - ): - name = add_prefix(name, "model") - - if name == "model.embed_tokens.weight": - if self.pp_group.is_last_rank and getattr( - self.config, "tie_word_embeddings", False - ): - if "lm_head.weight" in params_dict: - param = params_dict["lm_head.weight"] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - - layer_id = get_layer_id(name) - if ( - layer_id is not None - and hasattr(self.model, "start_layer") - and ( - layer_id < self.model.start_layer - or layer_id >= self.model.end_layer - ) - ): - continue - - if "rotary_emb.inv_freq" in name or "projector" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - continue - - if "scale" in name: - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - if "mlp.experts" in name: - continue - - name2 = name.replace(weight_name, param_name) - if name2.endswith(".bias") and name2 not in params_dict: - continue - if name2 not in params_dict: - continue - - param = params_dict[name2] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight, shard_id) - break - else: - is_expert_weight = False - for mapping in expert_params_mapping: - param_name, weight_name, expert_id, shard_id = mapping - if weight_name not in name: - continue - is_expert_weight = True - - name2 = name.replace(weight_name, param_name) - if name2 not in params_dict: - continue - - param = params_dict[name2] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader( - param, - loaded_weight, - name2, - shard_id=shard_id, - expert_id=expert_id, - ) - break - else: - if is_expert_weight: - continue - - # 3) regular params - if name.endswith(".bias") and name not in params_dict: - continue - if name not in params_dict: - continue - - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - - if not hasattr(self, "routed_experts_weights_of_layer"): - self.routed_experts_weights_of_layer = LazyValue( - lambda: { - lid: self.model.layers[lid].mlp.get_moe_weights() - for lid in range(self.start_layer, self.end_layer) - if isinstance(self.model.layers[lid].mlp, SDARMoeSparseMoeBlock) - } - ) - - @classmethod - def get_model_config_for_expert_location(cls, config): - return ModelConfigForExpertLocation( - num_layers=config.num_hidden_layers, - num_logical_experts=config.num_experts, - num_groups=None, - ) - - -EntryClass = SDARMoeForCausalLM diff --git a/test/registered/dllm/test_sdar.py b/test/registered/dllm/test_sdar.py deleted file mode 100644 index 5a179b07b..000000000 --- a/test/registered/dllm/test_sdar.py +++ /dev/null @@ -1,91 +0,0 @@ -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci - -register_cuda_ci(est_time=181, suite="stage-b-test-large-1-gpu") -register_amd_ci(est_time=330, suite="stage-b-test-small-1-gpu-amd") - -import unittest -from types import SimpleNamespace - -from sglang.srt.utils import kill_process_tree -from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k -from sglang.test.send_one import BenchArgs, send_one_prompt -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - is_in_amd_ci, - is_in_ci, - popen_launch_server, - write_github_step_summary, -) - - -class TestSDARMini(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = "JetLM/SDAR-8B-Chat" - cls.base_url = DEFAULT_URL_FOR_TEST - - other_args = [ - "--trust-remote-code", - "--mem-fraction-static", - "0.8", - "--max-running-requests", - "64", - "--attention-backend", - "flashinfer", - "--dllm-algorithm", - "LowConfidence", - "--tp", - "1", - ] - - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=other_args, - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - num_shots=5, - data_path=None, - num_questions=200, - max_new_tokens=1024, - parallel=128, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - metrics = run_eval_few_shot_gsm8k(args) - print(f"{metrics=}") - - self.assertGreater(metrics["accuracy"], 0.88) - if is_in_amd_ci(): - self.assertGreater(metrics["output_throughput"], 80) - else: - self.assertGreater(metrics["output_throughput"], 250) - - def test_bs_1_speed(self): - args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048) - acc_length, speed = send_one_prompt(args) - - print(f"{speed=:.2f}") - - if is_in_ci(): - write_github_step_summary( - f"### test_bs_1_speed (SDAR-8B-Chat) with tp1\n" - f"{speed=:.2f} token/s\n" - ) - if is_in_amd_ci(): - self.assertGreater(speed, 10) - else: - self.assertGreater(speed, 70) - - -if __name__ == "__main__": - unittest.main()