tested and working

    parser = _add_biencoder_args(parser)

    parser = _add_vit_args(parser)

    parser = _add_logging_args(parser)

    parser = _add_inference_args(parser)

    # Custom arguments.

    if extra_args_provider is not None:

    assert getattr(args, arg) is not None, '{} argument is None'.format(arg)

def _add_inference_args(parser):

    group = parser.add_argument_group(title='inference')

    group.add_argument('--inference-batch-times-seqlen-threshold',

                       type=int, default=512,

                       help='During inference, if batch-size times '

                       'sequence-length is smaller than this threshold '

                       'then we will not use pipelining, otherwise we will.')

    return parser

def _add_network_size_args(parser):

    group = parser.add_argument_group(title='network size')

    detokenize_generations)

def generate_and_post_process(model,

                              prompts=None,

                              tokens_to_generate=0,

                              return_output_log_probs=False,

                              return_all_log_probs=False,

                              greedy_sampling=False,

                              top_k_sampling=0,

                              top_p_sampling=0.0,

                              temperature=1.0,

                              add_BOS=False):

    """TO DO ..."""

                              add_BOS=False,

                              use_eod_token_for_early_termination=True):

    """Run inferecne and post-process outputs, i.e., detokenize,

    move to cpu and convert to list."""

    # Main inference.

    tokens, lengths, output_log_probs, all_log_probs = generate(

        tokens_to_generate=tokens_to_generate,

        return_output_log_probs=return_output_log_probs,

        return_all_log_probs=return_all_log_probs,

        greedy_sampling=greedy_sampling,

        top_k_sampling=top_k_sampling,

        top_p_sampling=top_p_sampling,

        temperature=temperature,

        add_BOS=add_BOS)

        add_BOS=add_BOS,

        use_eod_token_for_early_termination=use_eod_token_for_early_termination)

    # Only post-process on first stage.

    if mpu.is_pipeline_first_stage():

    return None

def generate(model,

             prompts=None,

             tokens_to_generate=0,

             return_output_log_probs=False,

             return_all_log_probs=False,

             greedy_sampling=False,

             top_k_sampling=0,

             top_p_sampling=0.0,

             temperature=1.0,

             add_BOS=False):

    """TO DO ..."""

             add_BOS=False,

             use_eod_token_for_early_termination=True):

    """Given prompts and input parameters, run inference and return:

       tokens: prompts plus the generated tokens.

       lengths: length of the prompt + generations. Note that we can

           discard tokens in the tokens tensor that are after the

           corresponding length.

       output_log_probs: log probs of the tokens.

       all_log_probs: full log probs for all of tokens.

    """

    # Make sure input params are avaialble to all ranks.

    values = [tokens_to_generate, return_output_log_probs,

              return_all_log_probs, temperature, add_BOS]

    values_float_tensor = broadcast_float_list(5, float_list=values)

    values = [tokens_to_generate,

              return_output_log_probs, return_all_log_probs,

              greedy_sampling, top_k_sampling, top_p_sampling,

              temperature, add_BOS, use_eod_token_for_early_termination]

    values_float_tensor = broadcast_float_list(9, float_list=values)

    tokens_to_generate = int(values_float_tensor[0].item())

    return_output_log_probs = bool(values_float_tensor[1].item())

    return_all_log_probs = bool(values_float_tensor[2].item())

    temperature = values_float_tensor[3].item()

    add_BOS = bool(values_float_tensor[4].item())

    greedy_sampling = bool(values_float_tensor[3].item())

    top_k_sampling = int(values_float_tensor[4].item())

    top_p_sampling = values_float_tensor[5].item()

    temperature = values_float_tensor[6].item()

    add_BOS = bool(values_float_tensor[7].item())

    use_eod_token_for_early_termination = bool(values_float_tensor[8].item())

    # Tokenize prompts and get the batch.

    # Note that these tensors are broadcaseted to all ranks.

        model, context_tokens_tensor, context_length_tensor,

        return_output_log_probs=return_output_log_probs,

        return_all_log_probs=return_all_log_probs,

        temperature=temperature)

        greedy=greedy_sampling, top_k=top_k_sampling, top_p=top_p_sampling,

        temperature=temperature,

        use_eod_token_for_early_termination=use_eod_token_for_early_termination)

from megatron import mpu

def recv_from_prev_pipeline_rank_(recv_buffer=None):

    """Receive from previous pipeline stage and update the

    input buffer inplace."""

    if not mpu.is_pipeline_first_stage():

        assert recv_buffer is not None

        recv_prev_op = torch.distributed.P2POp(

            torch.distributed.irecv, recv_buffer,

            mpu.get_pipeline_model_parallel_prev_rank())

        reqs = torch.distributed.batch_isend_irecv([recv_prev_op])

        for req in reqs:

            req.wait()

        # To protect against race condition when using batch_isend_irecv().

        torch.cuda.synchronize()

def send_to_next_pipeline_rank(tensor=None):

    """Send output to the next pipeline stage."""

    if not mpu.is_pipeline_last_stage():

        assert tensor is not None

        send_next_op = torch.distributed.P2POp(

            torch.distributed.isend, tensor,

            mpu.get_pipeline_model_parallel_next_rank())

        reqs = torch.distributed.batch_isend_irecv([send_next_op])

        for req in reqs:

            req.wait()

        # To protect against race condition when using batch_isend_irecv().

        torch.cuda.synchronize()

def broadcast_from_last_pipeline_stage(size, dtype, tensor=None):

    """Broadcast a tensor from last pipeline stage to all ranks."""

            tensor[...] = tensor_

def broadcast_tensor(size, dtype, tensor=None, rank=0):

    """ Given size and type of a tensor on all ranks and the tensor value

        only on a specific rank, broadcast from that rank to all other ranks.

    return tensor

def broadcast_list(size, dtype, list_values=None, rank=0):

    """Broadcast a list of values with a given type."""

    return broadcast_tensor(size, dtype, tensor=tensor, rank=rank)

def broadcast_int_list(size, int_list=None, rank=0):

    """Broadcast a list of interger values."""

    return broadcast_list(size, torch.int64, list_values=int_list, rank=rank)

def broadcast_float_list(size, float_list=None, rank=0):

    """Broadcast a list of float values."""

from megatron import (

    get_args,

    mpu)

from .communication import (

    send_to_next_pipeline_rank,

    recv_from_prev_pipeline_rank_)

class InferenceParams:

    """Inference parameters that are passed to the main model in order

    to efficienly calculate and store the context during inference."""

    def __init__(self, max_batch_size, max_sequence_len):

        """Note that offsets are set to zero and we always set the

        flag to allocate memory. After the first call, make sure to

        set this flag to False."""

        self.max_sequence_len = max_sequence_len

        self.max_batch_size = max_batch_size

        self.sequence_len_offset = 0

class ForwardStep:

    """Forward step function with all the communications.

    We use a class here to hide the inference parameters

    from the outside caller."""

    def __init__(self, model, max_batch_size, max_sequence_len):

        """Set values so we don't need to do it multiple times."""

        # Make sure model is in eval mode.

        if isinstance(model, Iterable):

            for this_model in model:

                this_model.eval()

        else:

        assert not isinstance(model, Iterable), \

            'interleaving schedule is not supported for inference'

        model.eval()

        self.model = model

        self.constant = 512

        # Initialize inference parameters.

        self.inference_params = InferenceParams(max_batch_size,

                                                max_sequence_len)

        # Pipelining arguments.

        args = get_args()

        self.pipeline_size_larger_than_one = args.pipeline_model_parallel_size

        # Threshold of pipelining.

        self.pipelining_batch_x_seqlen = \

            args.inference_batch_times_seqlen_threshold

    def __call__(self, tokens, position_ids, attention_mask):

        if tokens.size(0) * tokens.size(1) >= self.constant:

            micro_batch_size = max(1, self.constant // tokens.size(1))

            return _with_pipelining_forward_step(self.model, tokens,

        """Invocation of the forward methods. Note that self.inference_params

        is being modified by the forward step."""

        # Pipelining case.

        if self.pipeline_size_larger_than_one:

            current_batch_x_seqlen = tokens.size(0) * tokens.size(1)

            if current_batch_x_seqlen >= self.pipelining_batch_x_seqlen:

                micro_batch_size = \

                    max(1, self.pipelining_batch_x_seqlen // tokens.size(1))

                return _with_pipelining_forward_step(self.model,

                                                     tokens,

                                                     position_ids,

                                                     attention_mask,

                                                     self.inference_params,

                                                     micro_batch_size)

        else:

            return _no_pipelining_forward_step(self.model, tokens,

        return _no_pipelining_forward_step(self.model,

                                           tokens,

                                           position_ids,

                                           attention_mask,

                                           self.inference_params)

        recv_buffer = _allocate_recv_buffer(batch_size, sequence_length)

    # Receive from previous stage.

    if not mpu.is_pipeline_first_stage():

        torch.distributed.recv(recv_buffer,

                               src=mpu.get_pipeline_model_parallel_prev_rank())

    recv_from_prev_pipeline_rank_(recv_buffer)

    # Forward pass through the model.

    model.set_input_tensor(recv_buffer)

                          inference_params=inference_params)

    # Send output to the next stage.

    if not mpu.is_pipeline_last_stage():

        torch.distributed.send(output_tensor,

                               mpu.get_pipeline_model_parallel_next_rank())

    send_to_next_pipeline_rank(output_tensor)

    # Make sure we do not allocate context memory anymore.

    if inference_params.allocate_key_value_memory:

def _no_pipelining_forward_step(model, tokens, position_ids, attention_mask,

                                inference_params, recv_buffer=None):

    """If recv_buffer is none, we will allocate one on the fly."""

    # Run a simple forward pass.

    output_tensor = _forward_step_helper(model, tokens, position_ids,

                                         attention_mask, inference_params,

    return logits

def _with_pipelining_forward_step(model, tokens, position_ids, attention_mask,

                                  inference_params, micro_batch_size):

    """No interleaving is supported."""

    sequence_length = tokens.size(1)

    batch_size = tokens.size(0)

        model, tokens, lengths,

        return_output_log_probs=False,

        return_all_log_probs=False,

        temperature=1.0):

        greedy=False, top_k=0, top_p=0.0,

        temperature=1.0,

        use_eod_token_for_early_termination=True):

    """Main token generation function.

    Arguments:

        model: XXX

        model: no interleaving is supported.

        tokens: prompt tokens extended to be of size [b, max-sequence-length]

        lengths: original prompt length, size: [b]

        return_output_log_probs: flag to calculate the log probability of

        return_all_log_probs: flag to calculate the log probability of across

            all the tokens (vocab size). Note that the log probability is the

            one after logits are modifed for sampling.

        greedy, top_k, top_p: greedy, top-k, and top-p sampling parameters.

            Note that these three paramters are exclusive meaning that:

                if greedy = true then we should have top-k=top-p=0.

                if top-k > 0 then we expect greedy=false and top-p=0.

                if top-p > 0 then we check for greedy=false and top-k=0.

        temperature: sampling temperature.

        use_eod_token_for_early_termination: if True, do early termination if

            all the sequences have reached this token.

    Note: Outside of model, other parameters only need to be available on

          rank 0.

    Outputs: Note that is size is adjusted to a lower value than

    # Run infernece

    # =============

    with torch.no_grad():

        attention_mask, position_ids = _build_attention_mask_and_position_ids(

            tokens)

    with torch.no_grad():

        prev_context_length = 0

        for context_length in range(min_prompt_length, max_sequence_length):

                last_token_logits = logits[:, -1, :]

                new_sample, updated_last_token_logits = sample(

                    last_token_logits,

                    greedy=args.greedy,

                    top_k=args.top_k,

                    top_p=args.top_p,

                    greedy=greedy,

                    top_k=top_k,

                    top_p=top_p,

                    temperature=temperature,

                    vocab_size=tokenizer.vocab_size)

                # Now that we have the sample and updated logits,

                done = torch.all(is_generation_done)

            done = broadcast_from_last_pipeline_stage(1, torch.uint8,

                                                      tensor=done)

            #if done:

            #    break

            if use_eod_token_for_early_termination and done:

                break

    # ===================================================

    # Update the length of based on max generated length.