def forward(self,
                input_ids: torch.Tensor,
                position_ids: Optional[torch.LongTensor] = None,
                sequence_length: int = None,
                indices: torch.Tensor = None,
                cu_seqlens: int = None,
                max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]:
        """

        Args:
            input_ids: input ids. shape (1, totol_nnz)
            position_ids: position ids. shape (batch_size, seq_length)

        Returns:

        """
        inputs_embeds = self.embed_tokens(input_ids)  # (1, total_nnz) -> (1, total_nnz, hidden_size)

        # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size)
        inputs_embeds = inputs_embeds.transpose(0, 1)
        if self.megatron_config.sequence_parallel:
            inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds)

        hidden_states = inputs_embeds
        for idx, decoder_layer in enumerate(self.layers):
            layer_outputs = decoder_layer(hidden_states,
                                          position_ids=position_ids,
                                          sequence_length=sequence_length,
                                          indices=indices,
                                          cu_seqlens=cu_seqlens,
                                          max_seqlen_in_batch=max_seqlen_in_batch)

            hidden_states = layer_outputs

        hidden_states = self.norm(hidden_states)

        return hidden_states


        Args:
            input_ids: input ids. shape (1, totol_nnz)
            position_ids: position ids. shape (batch_size, seq_length)

        Returns:

        
verl/models/llama/megatron/modeling_llama_megatron.py