| 876 | """ |
| 877 | |
| 878 | def forward(self, inputs, **kwargs): |
| 879 | assert torch.is_tensor(inputs) or isinstance(inputs, tuple) |
| 880 | if torch.is_tensor(inputs) or len(inputs) == 2: |
| 881 | # No attention mask forwarded, search for args.attn_mask |
| 882 | if not hasattr(self, "_args"): |
| 883 | self._args = get_args() |
| 884 | hidden_states, query_hidden_state = inputs |
| 885 | attention_mask = self._args.attn_mask |
| 886 | return super().forward(hidden_states, query_hidden_state, attention_mask, **kwargs) |
| 887 | elif len(inputs) == 3: |
| 888 | # Attention mask is an activation. |
| 889 | hidden_states, query_hidden_state, attention_mask = inputs[0], inputs[1] |
| 890 | return super().forward(*inputs, **kwargs), attention_mask |
| 891 | else: |
| 892 | raise RuntimeError("Received more inputs than understood.") |
| 893 | |
| 894 | |
| 895 | class ParallelTransformer(MegatronModule): |