Batch the sequences with ContrastiveSeqBatchScheduler

[KexinFeng]

This PR succeeds PR #2547 #2509 #2557, which contains the benchmark outputs of the searching results.

The traced models can be downloaded here: https://d2l-java-resources.s3.amazonaws.com/tmp_model_unittest/model/nlp/text_generation/ai/djl/pytorch/gpt2/metadata.json and https://d2l-java-resources.s3.amazonaws.com/tmp_model_unittest/gpt2_decoder_model_merged.onnx.gz

Design

Overview

These PRs consists of the following three parts: LMAdapter, LMSearch and SeqBatchScheduler. There are three testing code corresponding to them: TestLMAdapter.java, TestLMSearch.java, TestSeqBatchScheduler.java. The TestSeqBatchScheduler.java utilizes all previous features.

New classes

/* 
This is a scheduler, serving as an API to the consumer of the systme, allowing for three major actions: initForward, addBatch, fastForward, collectResults. An optimal 
control sequence should be solved, after taking into the time consumption of each action, the batch size and 
sequence length of queueing requests.  Such optimal control solver needs additional effort. Primitive policy is 
setting several thresholds.
*/
SeqBatchScheduler {

    LMAdapter lmAdapter;

    SeqBatcher seqBatcher;

    public SeqBatchScheduler(LMAdapter lmAdapter) {
        this.lmAdapter = lmAdapter;
    }

    public abstract SeqBatcher initForward(
            NDArray inputIds, NDArray batchUids, SearchConfig config);

    public void incrementForward(int count) {}

    public void addRequest(NDArray inputIds, NDArray batchUids, SearchConfig config) {}

    public ConcurrentHashMap<Long, NDArray> collectResults() {}
}

// This stores the search state (BatchList), the control variables (eg seqLength, offSets, etc), and batch operations on BatchList.
public class SeqBatcher {

    NDManager manager;

    long batchSize;

    long seqLength;

    NDArray batchUid;

    NDArray offSets;

    BatchList data;

    ConcurrentHashMap<Long, NDArray> finishedSequences;

    List<Pair<Long, Long>> exitIndexEndPosition;

    SeqBatcher(BatchList data, NDArray batchUid, NDArray offSets, NDManager manager) {}

    /** Add new batch. Modify the batch dimension and the left padding. */
    public void addBatch(SeqBatcher seqBatcherNew) {   }

    /** Merge two batchers together. Modify the batch dimension and the left padding. */
    private void merge(SeqBatcher seqBatcher1, SeqBatcher seqBatcher2, long seqDelta) {}

    /** Check which batch needs to exit, according certain criteria like EOS or maxLength. */
    public void exitCriteria(NDArray outputIds, long maxLength, long eosTokenId) {}

    /** Collect the finished sequences and trim the left padding. */
    public void collectAndTrim() { }
}

// BatchList is a struct consisting of NDArrays, whose first dimension is batch, and also contains 
// sequence dimension (whose position in tensor's shape is specified by seqDimOrder)
// It represents a search state, and the NDArrays inside are updated in each iteration of the autoregressive loop.
// The batch operators in SeqBatcher will be applied on them.
public abstract class BatchList {

    // [batch, seq_past]. seq-dim-size == |past_seq| + |inputIds|. Will grow.
    public NDArray pastOutputIds;

    // [batch, seq_past]
    // The cache of past attentionMask. seq-dim-size == |past_seq| + |inputIds|. Will grow.
    public NDArray pastAttentionMask;

    // (k, v) * numLayer,
    // kv: [batch, heads, seq_past, kvfeature]
    // The cache of past sequence. seq-dim-size == |past_seq| + |inputIds|. Will grow.
    public NDList pastKeyValues;

    // Sequence dimension order among all dimensions for each element in the batch list.
    public long[] seqDimOrder;

    BatchList() {}

    BatchList(NDList list, long[] seqDimOrder) {}

    BatchList(
            NDArray pastOutputIds,
            NDArray pastAttentionMask,
            NDList pastKeyValues,
            long[] seqDimOrder) {}

    public abstract BatchList fromList(NDList inputList, long[] seqDimOrder);

    // The pastOutputIds has to be the first in the output list
    public abstract NDList getList();

    public long[] getSeqDimOrder() {
        return seqDimOrder;
    }
   } 
}

This is partly adapted from the searching algorithms in TestLMSearch.java. Currently only the contrastiveSearch has been adapted. The greedySearch and beamSearch will be similarly implemented. On top of that, batching operations are implemented. The system is designed as following, which contains the batching operations like addRequest() and collectAdnTrim()

The changes have been resolved

[KexinFeng]

Features have been merged in #2637