def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor: r"""Pass the input through the encoder layers in turn. Args: src: the sequence to the encoder (required). mask: the mask for the src sequence (optional). src_key_padding_mask: the mask for the src keys per batch (optional). Shape: see the docs in Transformer class. """ output = src convert_to_nested = False first_layer = self.layers[0] src_key_padding_mask_for_layers = src_key_padding_mask if isinstance(first_layer, torch.nn.TransformerEncoderLayer): if (not first_layer.norm_first and not first_layer.training and first_layer.self_attn.batch_first and first_layer.self_attn._qkv_same_embed_dim and first_layer.activation_relu_or_gelu and first_layer.norm1.eps == first_layer.norm2.eps and src.dim() == 3 and self.enable_nested_tensor and src_key_padding_mask is not None and torch._nested_tensor_from_mask_left_aligned( src, src_key_padding_mask.logical_not())): if src_key_padding_mask is not None and not output.is_nested and mask is None: tensor_args = ( src, first_layer.self_attn.in_proj_weight, first_layer.self_attn.in_proj_bias, first_layer.self_attn.out_proj.weight, first_layer.self_attn.out_proj.bias, first_layer.norm1.weight, first_layer.norm1.bias, first_layer.norm2.weight, first_layer.norm2.bias, first_layer.linear1.weight, first_layer.linear1.bias, first_layer.linear2.weight, first_layer.linear2.bias, ) if not torch.overrides.has_torch_function(tensor_args): if not torch.is_grad_enabled() or all( [not x.requires_grad for x in tensor_args]): if output.is_cuda or 'cpu' in str(output.device): convert_to_nested = True output = torch._nested_tensor_from_mask( output, src_key_padding_mask.logical_not()) src_key_padding_mask_for_layers = None for mod in self.layers: output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask_for_layers) if convert_to_nested: output = output.to_padded_tensor(0.) if self.norm is not None: output = self.norm(output) return output
def grad_test_func(inpt): nt = torch._nested_tensor_from_mask(inpt, mask) # This implicitly tests to_padded_tensor grads return nt.to_padded_tensor(0)
def forward_scriptable( self, src_tokens, src_lengths, return_all_hiddens: bool = False, ): """ Args: src_tokens (LongTensor): tokens in the source language of shape `(batch, src_len)` src_lengths (torch.LongTensor): lengths of each source sentence of shape `(batch)` return_all_hiddens (bool, optional): also return all of the intermediate hidden states (default: False). Returns: dict: - **encoder_out** (Tensor): the last encoder layer's output of shape `(src_len, batch, embed_dim)` - **encoder_padding_mask** (ByteTensor): the positions of padding elements of shape `(batch, src_len)` - **encoder_embedding** (Tensor): the (scaled) embedding lookup of shape `(batch, src_len, embed_dim)` - **encoder_states** (List[Tensor]): all intermediate hidden states of shape `(src_len, batch, embed_dim)`. Only populated if *return_all_hiddens* is True. """ if self.pre_encoder is not None: x, src_lengths, encoder_padding_mask = self.pre_encoder( src_tokens, src_lengths ) else: x, encoder_padding_mask = ( src_tokens, ~speech_utils.sequence_mask(src_lengths, src_tokens.size(1)), ) has_pads = src_tokens.device.type == "xla" or encoder_padding_mask.any() if self.fc0 is not None: x = self.dropout_module(x) x = self.fc0(x) x = self.embed_scale * x if self.embed_positions is not None: # 0s in `~encoder_padding_mask` are used as pad_idx for positional embeddings x = x + self.embed_positions((~encoder_padding_mask).int()) if self.layernorm_embedding is not None: x = self.layernorm_embedding(x) x = self.dropout_module(x) if self.quant_noise is not None: x = self.quant_noise(x) # account for padding while computing the representation if has_pads: x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x)) # B x T x C -> T x B x C x = x.transpose(0, 1) encoder_states = [] fc_results = [] if return_all_hiddens: encoder_states.append(x) attn_mask = self.get_attn_mask(src_lengths) # nested tensor and BT enable layer = self.layers[0] BT_flag = False NT_flag = False # torch version check, BT>=1.12.0 and NT>=1.13.0.dev20220613 # internal format is '1.13.0a0+fb' # external format is '1.13.0.dev20220613'(cpu&gpu) for nightly or "1.11.0"(cpu) or '1.11.0+cu102'(gpu) for stable BT_version = False NT_version = False if "fb" in torch.__version__: BT_version = True NT_version = True else: if "+" in torch.__version__: torch_version = torch.__version__.split("+")[0] else: torch_version = torch.__version__ torch_version = torch_version.split(".") int_version = ( int(torch_version[0]) * 1000 + int(torch_version[1]) * 10 + int(torch_version[2]) ) if len(torch_version) == 3: if int_version >= 1120: BT_version = True if int_version >= 1131: NT_version = True elif len(torch_version) == 4: if int_version >= 1130: BT_version = True # Consider _nested_tensor_from_mask_left_aligned is landed after "20220613" if int_version >= 1131 or ( int_version == 1130 and torch_version[3][3:] >= "20220613" ): NT_version = True if ( BT_version and x.dim() == 3 and layer.load_to_BT and not layer.return_fc and layer.can_use_fastpath and not layer.training and not layer.ever_training and not layer.cfg_checkpoint_activations ): # Batch first can not be justified but needs user to make sure x = x.transpose(0, 1) # Check mask conditions for nested tensor if NT_version: if ( encoder_padding_mask is not None and torch._nested_tensor_from_mask_left_aligned( x, encoder_padding_mask.logical_not() ) ): if not torch.is_grad_enabled() or not x.requires_grad: x = torch._nested_tensor_from_mask( x, encoder_padding_mask.logical_not() ) NT_flag = True BT_flag = True # encoder layers if NT_flag: processing_mask = None else: processing_mask = encoder_padding_mask encoder_padding_mask_out = processing_mask if has_pads else None for layer in self.layers: lr = layer( x, encoder_padding_mask=encoder_padding_mask_out, attn_mask=attn_mask, ) if isinstance(lr, tuple) and len(lr) == 2: x, fc_result = lr else: x = lr fc_result = None if return_all_hiddens and not torch.jit.is_scripting(): assert encoder_states is not None encoder_states.append(x) fc_results.append(fc_result) # change back to non-nested and Batch second if NT_flag: x = x.to_padded_tensor(0.0) if NT_flag or BT_flag: x = x.transpose(0, 1) if self.layer_norm is not None: x = self.layer_norm(x) # The Pytorch Mobile lite interpreter does not supports returning NamedTuple in # `forward` so we use a dictionary instead. # TorchScript does not support mixed values so the values are all lists. # The empty list is equivalent to None. return { "encoder_out": [x], # T x B x C "encoder_padding_mask": [encoder_padding_mask] if encoder_padding_mask.any() else [], # B x T "encoder_embedding": [], "encoder_states": encoder_states, # List[T x B x C] "fc_results": fc_results, # List[T x B x C] "src_tokens": [], "src_lengths": [src_lengths], # List[B] }
def _create_nested_tensor_from_mask(self, requires_grad=False): data = torch.randn(2, 3, 4, requires_grad=requires_grad) mask = torch.ones_like(data[:, :, 0]).bool() return torch._nested_tensor_from_mask(data, mask)
def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor: r"""Pass the input through the encoder layers in turn. Args: src: the sequence to the encoder (required). mask: the mask for the src sequence (optional). src_key_padding_mask: the mask for the src keys per batch (optional). Shape: see the docs in Transformer class. """ output = src convert_to_nested = False first_layer = self.layers[0] src_key_padding_mask_for_layers = src_key_padding_mask why_not_sparsity_fast_path = '' str_first_layer = "self.layers[0]" if not isinstance(first_layer, torch.nn.TransformerEncoderLayer): why_not_sparsity_fast_path = f"{str_first_layer} was not TransformerEncoderLayer" elif first_layer.norm_first: why_not_sparsity_fast_path = f"{str_first_layer}.norm_first was True" elif first_layer.training: why_not_sparsity_fast_path = f"{str_first_layer} was in training mode" elif not first_layer.self_attn.batch_first: why_not_sparsity_fast_path = f" {str_first_layer}.self_attn.batch_first was not True" elif not first_layer.self_attn._qkv_same_embed_dim: why_not_sparsity_fast_path = f"{str_first_layer}.self_attn._qkv_same_embed_dim was not True" elif not first_layer.activation_relu_or_gelu: why_not_sparsity_fast_path = f" {str_first_layer}.activation_relu_or_gelu was not True" elif not (first_layer.norm1.eps == first_layer.norm2.eps): why_not_sparsity_fast_path = f"{str_first_layer}.norm1.eps was not equal to {str_first_layer}.norm2.eps" elif not src.dim() == 3: why_not_sparsity_fast_path = f"input not batched; expected src.dim() of 3 but got {src.dim()}" elif not self.enable_nested_tensor: why_not_sparsity_fast_path = "enable_nested_tensor was not True" elif src_key_padding_mask is None: why_not_sparsity_fast_path = "src_key_padding_mask was None" elif (((not hasattr(self, "mask_check")) or self.mask_check) and not torch._nested_tensor_from_mask_left_aligned( src, src_key_padding_mask.logical_not())): why_not_sparsity_fast_path = "mask_check enabled, and src and src_key_padding_mask was not left aligned" elif output.is_nested: why_not_sparsity_fast_path = "NestedTensor input is not supported" elif mask is not None: why_not_sparsity_fast_path = "src_key_padding_mask and mask were both supplied" elif first_layer.self_attn.num_heads % 2 == 1: why_not_sparsity_fast_path = "num_head is odd" if not why_not_sparsity_fast_path: tensor_args = ( src, first_layer.self_attn.in_proj_weight, first_layer.self_attn.in_proj_bias, first_layer.self_attn.out_proj.weight, first_layer.self_attn.out_proj.bias, first_layer.norm1.weight, first_layer.norm1.bias, first_layer.norm2.weight, first_layer.norm2.bias, first_layer.linear1.weight, first_layer.linear1.bias, first_layer.linear2.weight, first_layer.linear2.bias, ) if torch.overrides.has_torch_function(tensor_args): why_not_sparsity_fast_path = "some Tensor argument has_torch_function" elif not (src.is_cuda or 'cpu' in str(src.device)): why_not_sparsity_fast_path = "src is neither CUDA nor CPU" elif torch.is_grad_enabled() and any(x.requires_grad for x in tensor_args): why_not_sparsity_fast_path = ( "grad is enabled and at least one of query or the " "input/output projection weights or biases requires_grad") if (not why_not_sparsity_fast_path) and (src_key_padding_mask is not None): convert_to_nested = True # simplify on or after on 8/16/2022 to unconditionally call with mask_check=False # we have established that either (1) the mask is OK with the check above, # or (2) that we don't need a mask check with mask_check=False in the init if not torch.jit.is_scripting(): output = torch._nested_tensor_from_mask( output, src_key_padding_mask.logical_not(), mask_check=False) else: # When scripting, make a simpler call until the FC bar passes on 8/16/2022 output = torch._nested_tensor_from_mask( output, src_key_padding_mask.logical_not()) src_key_padding_mask_for_layers = None for mod in self.layers: output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask_for_layers) if convert_to_nested: output = output.to_padded_tensor(0.) if self.norm is not None: output = self.norm(output) return output