Merge branch 'main' into main_dedup

#!/bin/bash

# Compute embeddings for each entry of a given dataset (e.g. Wikipedia)

RANK=0

WORLD_SIZE=1

# Wikipedia data can be downloaded from the following link:

# https://github.com/facebookresearch/DPR/blob/master/data/download_data.py

EVIDENCE_DATA_DIR=<Specify path of Wikipedia dataset>

EMBEDDING_PATH=<Specify path to store embeddings>

CHECKPOINT_PATH=<Specify path of pretrained ICT model>

python tools/create_doc_index.py \

    --num-layers 12 \

    --hidden-size 768 \

    --num-attention-heads 12 \

    --tensor-model-parallel-size 1 \

    --micro-batch-size 128 \

    --checkpoint-activations \

    --seq-length 512 \

    --retriever-seq-length 256 \

    --max-position-embeddings 512 \

    --load ${CHECKPOINT_PATH} \

    --evidence-data-path ${EVIDENCE_DATA_DIR} \

    --embedding-path ${EMBEDDING_PATH} \

    --indexer-log-interval 1000 \

    --indexer-batch-size 128 \

    --vocab-file bert-vocab.txt \

    --num-workers 2 \

    --fp16

#!/bin/bash

# Evaluate natural question test data given Wikipedia embeddings and pretrained

# ICT model

# Datasets can be downloaded from the following link:

# https://github.com/facebookresearch/DPR/blob/master/data/download_data.py

EVIDENCE_DATA_DIR=<Specify path of Wikipedia dataset>

EMBEDDING_PATH=<Specify path of the embeddings>

CHECKPOINT_PATH=<Specify path of pretrained ICT model>

QA_FILE=<Path of the natural question test dataset>

python tasks/main.py \

    --task ICT-ZEROSHOT-NQ \

    --tokenizer-type BertWordPieceLowerCase \

    --num-layers 12 \

    --hidden-size 768 \

    --num-attention-heads 12 \

    --tensor-model-parallel-size 1 \

    --micro-batch-size 128 \

    --checkpoint-activations \

    --seq-length 512 \

    --max-position-embeddings 512 \

    --load ${CHECKPOINT_PATH} \

    --evidence-data-path ${EVIDENCE_DATA_DIR} \

    --embedding-path ${EMBEDDING_PATH} \

    --retriever-seq-length 256 \

    --vocab-file  bert-vocab.txt\

    --qa-data-test ${QA_FILE} \

    --num-workers 2 \

    --faiss-use-gpu \

    --retriever-report-topk-accuracies 1 5 20 100 \

    --fp16

#! /bin/bash

# Runs the "217M" parameter biencoder model for ICT retriever

RANK=0

WORLD_SIZE=1

PRETRAINED_BERT_PATH=<Specify path of pretrained BERT model>

TEXT_DATA_PATH=<Specify path and file prefix of the text data>

TITLE_DATA_PATH=<Specify path and file prefix od the titles>

CHECKPOINT_PATH=<Specify path>

python pretrain_ict.py \

        --num-layers 12 \

        --hidden-size 768 \

        --num-attention-heads 12 \

        --tensor-model-parallel-size 1 \

        --micro-batch-size 32 \

        --seq-length 256 \

        --max-position-embeddings 512 \

        --train-iters 100000 \

        --vocab-file bert-vocab.txt \

        --tokenizer-type BertWordPieceLowerCase \

        --DDP-impl torch \

        --bert-load ${PRETRAINED_BERT_PATH} \

        --log-interval 100 \

        --eval-interval 1000 \

        --eval-iters 10 \

        --retriever-report-topk-accuracies 1 5 10 20 100 \

        --retriever-score-scaling \

        --load $CHECKPOINT_PATH \

        --save $CHECKPOINT_PATH \

        --data-path ${TEXT_DATA_PATH} \

        --titles-data-path ${TITLE_DATA_PATH} \

        --lr 0.0001 \

        --lr-decay-style linear \

        --weight-decay 1e-2 \

        --clip-grad 1.0 \

        --lr-warmup-fraction 0.01 \

        --save-interval 4000 \

        --exit-interval 8000 \

        --query-in-block-prob 0.1 \

        --fp16

import os

import torch

from megatron import fused_kernels

def parse_args(extra_args_provider=None, defaults={},

               ignore_unknown_args=False):

    parser = _add_validation_args(parser)

    parser = _add_data_args(parser)

    parser = _add_autoresume_args(parser)

    parser = _add_realm_args(parser)

    parser = _add_biencoder_args(parser)

    parser = _add_vit_args(parser)

    parser = _add_logging_args(parser)

    model_parallel_size = args.pipeline_model_parallel_size * \

                          args.tensor_model_parallel_size

    assert args.world_size % model_parallel_size == 0, 'world size is not'\

        ' divisible by tensor parallel size ({}) times pipeline paralle ' \

        ' divisible by tensor parallel size ({}) times pipeline parallel ' \

        'size ({})'.format(args.world_size, args.tensor_model_parallel_size,

                           args.pipeline_model_parallel_size)

    args.data_parallel_size = args.world_size // model_parallel_size

            print('setting global batch size to {}'.format(

                args.global_batch_size), flush=True)

    assert args.global_batch_size > 0

    if args.num_layers_per_virtual_pipeline_stage is not None:

        assert args.pipeline_model_parallel_size > 2, \

            'pipeline-model-parallel size should be greater than 2 with ' \

            'interleaved schedule'

        assert args.num_layers % args.num_layers_per_virtual_pipeline_stage == 0, \

            'number of layers is not divisible by number of layers per virtual ' \

            'pipeline stage'

        args.virtual_pipeline_model_parallel_size = \

            (args.num_layers // args.pipeline_model_parallel_size) // \

            args.num_layers_per_virtual_pipeline_stage

    else:

        args.virtual_pipeline_model_parallel_size = None

    # Parameters dtype.

    args.params_dtype = torch.float

    if args.fp16:

        assert not args.bf16

        args.params_dtype = torch.half

    if args.bf16:

        assert not args.fp16

        args.params_dtype = torch.bfloat16

    if args.rank == 0:

        print('using {} for parameters ...'.format(args.params_dtype),

              flush=True)

    # If we do accumulation and all-reduces in fp32, we need to have

    # local DDP and we should set the use-contiguous-buffers-in-ddp. 

    if args.accumulate_allreduce_grads_in_fp32:

        assert args.DDP_impl == 'local'

        args.use_contiguous_buffers_in_ddp = True

    if args.dataloader_type is None:

        args.dataloader_type = 'single'

    if args.fp16_lm_cross_entropy:

        assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.'

    if args.fp32_residual_connection:

        assert args.fp16, \

            'residual connection in fp32 only supported when using fp16.'

        assert args.fp16 or args.bf16, \

            'residual connection in fp32 only supported when using fp16 or bf16.'

    # Activation checkpointing.

    if args.distribute_checkpointed_activations:

        assert args.checkpoint_activations, \

            'for distribute-checkpointed-activations to work you '\

            'need to enable checkpoint-activations'

    # custom kernel constraints check

    seq_len = args.seq_length

    attn_batch_size = \

        (args.num_attention_heads / args.tensor_model_parallel_size) * \

        args.micro_batch_size

    # constraints on sequence length and attn_batch_size to enable warp based

    # optimization and upper triangular optimization (for causal mask)

    custom_kernel_constraint = seq_len > 16 and seq_len <=2048 and \

        seq_len % 4 == 0 and attn_batch_size % 4 == 0

    if args.fp16 and custom_kernel_constraint and args.masked_softmax_fusion:

        print('WARNING: constraints for invoking optimized'

            ' fused softmax kernel are not met. We default back to unfused'

            ' kernel invocations.')

    # Load scaled_masked_softmax_fusion_kernels

    if args.masked_softmax_fusion:

        fused_kernels.load_scaled_upper_triang_masked_softmax_fusion_kernel()

        fused_kernels.load_scaled_masked_softmax_fusion_kernel()

    # Load mixed precision fused layer norm.

    if args.fp32_residual_connection:

        fused_kernels.load_fused_mix_prec_layer_norm_kernel()

    _print_args(args)

    return args

    group.add_argument('--log-params-norm', action='store_true',

                       help='If set, calculate and log parameters norm.')

    group.add_argument('--log-num-zeros-in-grad', action='store_true',

                       help='If set, calculate and log the number of zeros in gradient.')

    group.add_argument('--tensorboard-log-interval', type=int, default=1,

                       help='Report to tensorboard interval.')

    group.add_argument('--tensorboard-queue-size', type=int, default=1000,

    group.add_argument('--fp16', action='store_true',

                       help='Run model in fp16 mode.')

    group.add_argument('--bf16', action='store_true',

                       help='Run model in bfloat16 mode.')

    group.add_argument('--loss-scale', type=float, default=None,

                       help='Static loss scaling, positive power of 2 '

                       'values can improve fp16 convergence. If None, dynamic'

                       help='Run attention masking and softmax in fp32. '

                       'This flag is ignored unless '

                       '--no-query-key-layer-scaling is specified.')

    group.add_argument('--fp32-allreduce', action='store_true',

                       help='All-reduce in fp32')

    group.add_argument('--accumulate-allreduce-grads-in-fp32',

                       action='store_true',

                       help='Gradient accumulation and all-reduce in fp32.')

    group.add_argument('--fp16-lm-cross-entropy', action='store_true',

                       help='Move the cross entropy unreduced loss calculation'

                       'for lm head to fp16.')

    group.add_argument('--model-parallel-size', type=int, default=None,

                       help='Old model parallel argument, do not use. Use '

                       '--tensor-model-parallel-size instead.')

    group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None,

                       help='Number of layers per virtual pipeline stage')

    group.add_argument('--distributed-backend', default='nccl',

                       choices=['nccl', 'gloo'],

                       help='Which backend to use for distributed training.')

                       choices=['local', 'torch'],

                       help='which DistributedDataParallel implementation '

                       'to use.')

    group.add_argument('--use-contiguous-buffers-in-ddp', action='store_true',

                       help='If set, use contiguous buffer in DDP. Note that '

                       'this option only works woth local DDP.' )

    group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false',

                       help='Use scatter/gather to optimize communication of tensors in pipeline',

                       dest='scatter_gather_tensors_in_pipeline')

    group.add_argument('--local_rank', type=int, default=None,

                       help='local rank passed from distributed launcher.')

    group.add_argument('--lazy-mpu-init', type=bool, required=False,

                       'This should be exclusive of --seq-length')

    group.add_argument('--decoder-seq-length', type=int, default=None,

                       help="Maximum decoder sequence length to process.")

    group.add_argument('--retriever-seq-length', type=int, default=256,

                       help='Maximum sequence length for the biencoder model '

                        ' for retriever')

    group.add_argument('--sample-rate', type=float, default=1.0,

                       help='sample rate for training data. Supposed to be 0 '

                            ' < sample_rate < 1')

    group.add_argument('--mask-prob', type=float, default=0.15,

                       help='Probability of replacing a token with mask.')

    group.add_argument('--short-seq-prob', type=float, default=0.1,

    return parser

def _add_realm_args(parser):

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

def _add_biencoder_args(parser):

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

    # network size

    group.add_argument('--ict-head-size', type=int, default=None,

                       help='Size of block embeddings to be used in ICT and '

                        'REALM (paper default: 128)')

    group.add_argument('--biencoder-projection-dim', type=int, default=0,

                       help='Size of projection head used in biencoder (paper'

                        ' default: 128)')

    group.add_argument('--biencoder-shared-query-context-model', action='store_true',

                        help='Whether to share the parameters of the query '

                        'and context models or not')

    # checkpointing

    group.add_argument('--ict-load', type=str, default=None,

                       'ICT dataset')

    group.add_argument('--use-one-sent-docs', action='store_true',

                       help='Whether to use one sentence documents in ICT')

    group.add_argument('--evidence-data-path', type=str, default=None,

                       help='Path to Wikipedia Evidence frm DPR paper')

    # training

    group.add_argument('--report-topk-accuracies', nargs='+', default=[],

                       help="Which top-k accuracies to report (e.g. '1 5 20')")

    group.add_argument('--retriever-report-topk-accuracies', nargs='+', type=int,

                        default=[], help="Which top-k accuracies to report "

                        "(e.g. '1 5 20')")

    group.add_argument('--retriever-score-scaling', action='store_true',

                       help='Whether to scale retriever scores by inverse '

                        'square root of hidden size')

    # faiss index

    group.add_argument('--faiss-use-gpu', action='store_true',

                       help='Whether create the FaissMIPSIndex on GPU')

    group.add_argument('--block-data-path', type=str, default=None,

                       help='Where to save/load BlockData to/from')

    group.add_argument('--embedding-path', type=str, default=None,

                       help='Where to save/load Open-Retrieval Embedding'

                        ' data to/from')

    # indexer

    group.add_argument('--indexer-batch-size', type=int, default=128,

import numpy as np

import torch

from torch.nn.parallel import DistributedDataParallel as torchDDP

from megatron import (get_args,

                      mpu,

                      print_rank_0,

                      update_num_microbatches)

                      update_num_microbatches,

                      utils)

_CHECKPOINT_VERSION = None

    args = get_args()

    # Only rank zero of the data parallel writes to the disk.

    if isinstance(model, torchDDP):

        model = model.module

    model = utils.unwrap_model(model)

    print_rank_0('saving checkpoint at iteration {:7d} to {}'.format(

        iteration, args.save))

        state_dict['args'] = args

        state_dict['checkpoint_version'] = 3.0

        state_dict['iteration'] = iteration

        state_dict['model'] = model.state_dict_for_save_checkpoint()

        if len(model) == 1:

            state_dict['model'] = model[0].state_dict_for_save_checkpoint()

        else:

            for i in range(len(model)):

                mpu.set_virtual_pipeline_model_parallel_rank(i)

                state_dict['model%d' % i] = model[i].state_dict_for_save_checkpoint()

        # Optimizer stuff.

        if not args.no_save_optim:

    return t

def fix_query_key_value_ordering(model, checkpoint_version):

    """Fix up query/key/value matrix ordering if checkpoint

    version is smaller than 2.0

    """

    if checkpoint_version < 2.0:

        for name, param in model.named_parameters():

            if name.endswith(('.query_key_value.weight', '.query_key_value.bias')):

                if checkpoint_version == 0:

                    fixed_param = _transpose_first_dim(param.data, 3, True, model)

                elif checkpoint_version == 1.0:

                    fixed_param = _transpose_first_dim(param.data, 3, False, model)

                else:

                    print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")

                    sys.exit()

                param.data.copy_(fixed_param)

            if name.endswith(('.key_value.weight', '.key_value.bias')):

                if checkpoint_version == 0:

                    fixed_param = _transpose_first_dim(param.data, 2, True, model)

                elif checkpoint_version == 1.0:

                    fixed_param = _transpose_first_dim(param.data, 2, False, model)

                else:

                    print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")

                    sys.exit()

                param.data.copy_(fixed_param)

        print_rank_0(" succesfully fixed query-key-values ordering for"

                    " checkpoint version {}".format(checkpoint_version))

def load_checkpoint(model, optimizer, lr_scheduler, load_arg='load', strict=True):

    """Load a model checkpoint and return the iteration.

    strict (bool): whether to strictly enforce that the keys in

    args = get_args()

    load_dir = getattr(args, load_arg)

    if isinstance(model, torchDDP):

        model = model.module

    model = utils.unwrap_model(model)

    # Read the tracker file and set the iteration.

    tracker_filename = get_checkpoint_tracker_filename(load_dir)

        print_rank_0('could not find arguments in the checkpoint ...')

    # Model.

    model.load_state_dict(state_dict['model'], strict=strict)

    if len(model) == 1:

        model[0].load_state_dict(state_dict['model'], strict=strict)

    else:

        for i in range(len(model)):

            mpu.set_virtual_pipeline_model_parallel_rank(i)

            model[i].load_state_dict(state_dict['model%d' % i], strict=strict)

    # Fix up query/key/value matrix ordering

    if get_checkpoint_version() < 2.0:

    # Fix up query/key/value matrix ordering if needed

    checkpoint_version = get_checkpoint_version()

        for name, param in model.named_parameters():

            if name.endswith(('.query_key_value.weight', '.query_key_value.bias')):

                if checkpoint_version == 0:

                    fixed_param = _transpose_first_dim(param.data, 3, True, model)

                elif checkpoint_version == 1.0:

                    fixed_param = _transpose_first_dim(param.data, 3, False, model)

                else:

                    print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")

                    sys.exit()

                param.data.copy_(fixed_param)

            if name.endswith(('.key_value.weight', '.key_value.bias')):

                if checkpoint_version == 0:

                    fixed_param = _transpose_first_dim(param.data, 2, True, model)

                elif checkpoint_version == 1.0:

                    fixed_param = _transpose_first_dim(param.data, 2, False, model)

                else:

                    print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")

                    sys.exit()

                param.data.copy_(fixed_param)

    print_rank_0(f' checkpoint version {checkpoint_version}')

    fix_query_key_value_ordering(model, checkpoint_version)

    # Optimizer.

    if not release and not args.finetune and not args.no_load_optim:

    return iteration

def load_ict_checkpoint(model, only_query_model=False, only_block_model=False, from_realm_chkpt=False):

    """selectively load ICT models for indexing/retrieving from ICT or REALM checkpoints"""

def load_biencoder_checkpoint(model, only_query_model=False,

        only_context_model=False, custom_load_path=None):

    """

    selectively load retrieval models for indexing/retrieving 

    from saved checkpoints

    """

    args = get_args()

    if isinstance(model, torchDDP):

        model = model.module

    model = utils.unwrap_model(model)

    load_path = args.load if from_realm_chkpt else args.ict_load

    load_path = custom_load_path if custom_load_path is not None else args.load

    tracker_filename = get_checkpoint_tracker_filename(load_path)

    with open(tracker_filename, 'r') as f:

        iteration = int(f.read().strip())

    # assert iteration > 0

    checkpoint_name = get_checkpoint_name(load_path, iteration, False)

    if mpu.get_data_parallel_rank() == 0:

        print('global rank {} is loading checkpoint {}'.format(

            torch.distributed.get_rank(), checkpoint_name))

    state_dict = torch.load(checkpoint_name, map_location='cpu')

    ict_state_dict = state_dict['model']

    if from_realm_chkpt and mpu.get_data_parallel_rank() == 0:

        print(" loading ICT state dict from REALM", flush=True)

        ict_state_dict = ict_state_dict['retriever']['ict_model']

    ret_state_dict = state_dict['model']

    if only_query_model:

        ict_state_dict.pop('context_model')

    if only_block_model:

        ict_state_dict.pop('question_model')

        ret_state_dict.pop('context_model')

    if only_context_model:

        ret_state_dict.pop('query_model')

    model.load_state_dict(ict_state_dict)

    assert len(model) == 1

    model[0].load_state_dict(ret_state_dict)

    torch.distributed.barrier()

    if mpu.get_data_parallel_rank() == 0:

        print(' successfully loaded {}'.format(checkpoint_name))

    return model