[optimize] add two stream norm for qwen3 (#7740)

Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Yi Zhang
2025-07-03 09:59:17 -07:00
committed by GitHub
co-authored by ispobock
parent 646cef2e2e
commit 264dc6e744
4 changed files with 54 additions and 10 deletions
+3
View File
@@ -291,6 +291,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
layer_id: int,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
alt_stream: Optional[torch.cuda.Stream] = None,
) -> None:
super().__init__()
self.config = config
@@ -393,6 +394,7 @@ class Qwen2MoeModel(nn.Module):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
decoder_layer_type: type[nn.Module] = Qwen2MoeDecoderLayer,
alt_stream: Optional[torch.cuda.Stream] = None,
) -> None:
super().__init__()
self.padding_idx = config.pad_token_id
@@ -418,6 +420,7 @@ class Qwen2MoeModel(nn.Module):
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,