| 667 | class CaptionMLP(Module): |
| 668 | |
| 669 | def __init__( |
| 670 | self, |
| 671 | in_features, |
| 672 | hidden_features=None, |
| 673 | out_features=None, |
| 674 | act_layer="gelu", |
| 675 | bias=True, |
| 676 | dtype=None, |
| 677 | mapping=Mapping(), |
| 678 | quant_mode=QuantMode(0), |
| 679 | inner_layernorm=False, |
| 680 | eps=1e-05, |
| 681 | ): |
| 682 | super().__init__() |
| 683 | hidden_act = act_layer |
| 684 | out_features = out_features or in_features |
| 685 | hidden_features = hidden_features or in_features |
| 686 | bias = to_2tuple(bias) |
| 687 | if hidden_act not in ACT2FN: |
| 688 | raise ValueError( |
| 689 | 'unsupported activation function: {}'.format(hidden_act)) |
| 690 | fc_output_size = 2 * hidden_features if hidden_act in [ |
| 691 | 'swiglu', 'gegelu' |
| 692 | ] else hidden_features |
| 693 | self.inner_layernorm = LayerNorm(hidden_features, dtype=dtype, |
| 694 | eps=eps) if inner_layernorm else None |
| 695 | |
| 696 | self.fc1 = ColumnLinear(in_features, |
| 697 | fc_output_size, |
| 698 | bias=bias[0], |
| 699 | dtype=dtype, |
| 700 | tp_group=mapping.tp_group, |
| 701 | tp_size=mapping.tp_size, |
| 702 | gather_output=False) |
| 703 | self.fc2 = RowLinear(hidden_features, |
| 704 | out_features, |
| 705 | bias=bias[1], |
| 706 | dtype=dtype, |
| 707 | tp_group=mapping.tp_group, |
| 708 | tp_size=mapping.tp_size) |
| 709 | |
| 710 | self.in_features = in_features |
| 711 | self.hidden_features = hidden_features |
| 712 | self.out_features = out_features |
| 713 | self.hidden_act = hidden_act |
| 714 | self.dtype = dtype |
| 715 | self.bias = bias |
| 716 | self.mapping = mapping |
| 717 | self.quant_mode = quant_mode |
| 718 | self.eps = eps |
| 719 | |
| 720 | def forward(self, hidden_states, gegelu_limit=None): |
| 721 | inter = self.fc1(hidden_states) |