Method
__init__
(self,
moe_config: MoeConfig,
hidden_size: int,
ffn_hidden_size: int,
hidden_act: str,
mapping: Mapping = Mapping(),
bias: bool = True,
dtype=None,
tp_group: List[int] = None,
tp_size: int = 1,
quant_mode=QuantMode(0),
use_shared_gate: bool = False,
use_side_stream: bool = False)
Source from the content-addressed store, hash-verified
| 1466 | class SharedMoE(MOE): |
| 1467 | |
| 1468 | def __init__(self, |
| 1469 | moe_config: MoeConfig, |
| 1470 | hidden_size: int, |
| 1471 | ffn_hidden_size: int, |
| 1472 | hidden_act: str, |
| 1473 | mapping: Mapping = Mapping(), |
| 1474 | bias: bool = True, |
| 1475 | dtype=None, |
| 1476 | tp_group: List[int] = None, |
| 1477 | tp_size: int = 1, |
| 1478 | quant_mode=QuantMode(0), |
| 1479 | use_shared_gate: bool = False, |
| 1480 | use_side_stream: bool = False): |
| 1481 | super().__init__( |
| 1482 | moe_config=moe_config, |
| 1483 | hidden_size=hidden_size, |
| 1484 | ffn_hidden_size=ffn_hidden_size, |
| 1485 | hidden_act=hidden_act, |
| 1486 | mapping=mapping, |
| 1487 | bias=bias, |
| 1488 | dtype=dtype, |
| 1489 | tp_group=tp_group, |
| 1490 | tp_size=tp_size, |
| 1491 | quant_mode=quant_mode, |
| 1492 | use_all_reduce=False, |
| 1493 | ) |
| 1494 | self.shared_expert = MLP( |
| 1495 | hidden_size=hidden_size, |
| 1496 | ffn_hidden_size=moe_config.shared_expert_intermediate_size, |
| 1497 | hidden_act=hidden_act, |
| 1498 | bias=False, |
| 1499 | dtype=self.dtype, |
| 1500 | tp_group=tp_group, |
| 1501 | tp_size=tp_size, |
| 1502 | quant_mode=self.quant_mode, |
| 1503 | is_expert=True, |
| 1504 | ) |
| 1505 | self.use_shared_gate = use_shared_gate |
| 1506 | if use_shared_gate: |
| 1507 | self.shared_expert_gate = RowLinear( |
| 1508 | hidden_size, |
| 1509 | 1, |
| 1510 | bias=False, |
| 1511 | dtype=dtype, |
| 1512 | tp_group=None, |
| 1513 | tp_size=1, |
| 1514 | ) |
| 1515 | else: |
| 1516 | self.shared_expert_gate = None |
| 1517 | self.use_side_stream = use_side_stream |
| 1518 | |
| 1519 | def forward(self, hidden_states, lora_layer_params=None): |
| 1520 | side_stream_id = SideStreamIDType.moe if self.use_side_stream else SideStreamIDType.disable |
Tested by
no test coverage detected