class_docstr: Inference model implementation for ragged batching for Llama-2 models.
class_docstr_tok: ["Inference","model","implementation","for","ragged","batching","for","Llama-2","models","."]
docstr: Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
  optimization to fuse the layer norm of the next layer into the current layer.
  
  Arguments:
      layer_idx (int): The index of the layer to execute.
      residual (torch.Tensor): The residual tensor from the previous layer.
      hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
          hidden states after pre normalization.
      ragged_batch_info (RaggedBatchWrapper): The batch metadata.
docstr_tok: ["Executes","one","(","slightly","offset",")","layer","of","the","transformer",".","This","implementation","does","a","peak-ahead","optimization","to","fuse","the","layer","norm","of","the","next","layer","into","the","current","layer",".","Arguments",":","layer_idx","(","int",")",":","The","index","of","the","layer","to","execute",".","residual","(","torch.Tensor",")",":","The","residual","tensor","from","the","previous","layer",".","hidden_states","(","torch.Tensor",")",":","The","hidden","states","from","the","previous","layer",".","This","is","the","hidden","states","after","pre","normalization",".","ragged_batch_info","(","RaggedBatchWrapper",")",":","The","batch","metadata","."]
code: def _forward_transformer_layer(
      self,
      layer_idx: int,
      residual: torch.Tensor,
      hidden_states: torch.Tensor,
      ragged_batch_info: RaggedBatchWrapper,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
      optimization to fuse the layer norm of the next layer into the current layer.
  
      Arguments:
          layer_idx (int): The index of the layer to execute.
          residual (torch.Tensor): The residual tensor from the previous layer.
          hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
              hidden states after pre normalization.
          ragged_batch_info (RaggedBatchWrapper): The batch metadata.
      """
      assert (
          self.config.parallel_attn
      ), "Only parallel attention implementation is supported"
  
      cur_params = self._transformer[layer_idx]
      kv_cache = self.state_manager.get_cache(layer_idx)
  
      attn_ln_out = hidden_states
      attn_hidden_state = self.qkv(
          attn_ln_out, cur_params.qkv_w, b=None
      )
      attn_hidden_state = self.attn(
          attn_hidden_state, kv_cache, ragged_batch_info
      )
      attention_output = self.attn_out(
          attn_hidden_state, cur_params.attn_out_w, b=None
      )
  
      if self.config.new_decoder_architecture:
          residual, mlp_ln_out = self.norm(
              residual,
              None,
              gamma=cur_params.ln_mlp_gamma,
              beta=cur_params.ln_mlp_beta,
          )
      else:
          mlp_ln_out = hidden_states
  
      mlp_hidden_state = self.mlp_1(
          mlp_ln_out, cur_params.mlp_1_w, b=None
      )
      mlp_output = self.mlp_2(
          mlp_hidden_state, cur_params.mlp_2_w, b=None
      )
  
      mlp_output.add_(attention_output)
  
      if self.tp_size > 1:
          dist.all_reduce(mlp_output, group=self._base_mp_group)
  
      if layer_idx != self.num_layers - 1:
          next_params = self._transformer[layer_idx + 1]
          residual, mlp_output = self.norm(
              residual,
              mlp_output,
              next_params.ln_attn_gamma,
              beta=next_params.ln_attn_beta,
          )
      else:
          # On last layer, we just need to perform the residual add. Adding into the residual
          # here is safe.
          residual.add_(mlp_output)
  
      return residual, mlp_output
  
imports: from typing import Iterable, Optional, Tuple
  import torch
  import deepspeed.comm
  from ...allocator import empty_from
  from ...inference_utils import ActivationType, DtypeEnum
  from ..None import *
  from ...modules.configs import *
  from ...modules.interfaces import *
  from ...ragged import RaggedBatchWrapper
  from .container import FalconNonTransformerContainer, FalconTransformerContainer
label_desc: Use image node_id 15 for calling the FalconInferenceModel obj's underlying member method code with example usage: obj._forward_transformer_layer(layer_idx, residual, hidden_states, ragged_batch_info) and returns: residual,  mlp_output