POC of LLMDecoder on both Pt trace model and Onnx model (LMAdapter part)

[KexinFeng]

Interface stepGenerator will be renamed to LMAdapter

Description

This is a POC which shows that the GPT2 can be traced and exported into DJL for a step inference. Moreover, this step inference allows the following:

1. batch sequence input
2. cached past_key_values input, which is a Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], ...] in python.

In this POC, both GPT2_init.pt and GPT2.pt are used since they have different inputs, the former not having past_key_values input.

Run TextGeneration.java to test it.

Design

The new classes are the following:

public class JavaDecoder {
    private StepGenerator generator;
    
    private DecodeParam decodingParams;
    
    public Text generateText() {};
}

JavaDecoder is made engine agnostic. The inference loop can be implemented inside generateText().

public interface StepGenerator {
    private String modelUrl;  // used to load the model
    
    public Token stepGen(inputIds, positionIds, attentionMask, pastKeyValues);  
    // This function is the same as P(w_n|w_{n-1}, w_{n-2}, ...)
}

This interface is a wrapper over the model files from different sources, e.g. gpt2.pt, gpt2.onnx, etc.
It can be seen as a Java abstraction of a causal language model, i.e. the conditional probability p_\theta(v_t | x_1, ..., x_{t-1}), i.e. given the past tokens up to a certain time x_{< t}, the probability that the next token is v, taken from a vocabulary set V. \theta is the model's weight.

It will be implemented individually to adapt for different models, which applies to different scenarios.

1. Traced GPT2.pt + JNI_for_pytorch. Its POC is done.
2. GPT2.onnx + OnnxEngine. This will be the use case corresponding to the graph above.
3. GPT2 + nueronx
4. FasterTransformer + tensorRTEngine

Model tracing

The onnx model gpt2.onnx is loaded from https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#exporting-a-model-using-past-keysvalues-in-the-decoder.
See also https://github.com/huggingface/optimum/releases.

The gpt2.pt is traced with the following scripts: https://gist.github.com/KexinFeng/4876c6bfb27f40abffe4d5a92c02acff

[zachgk]

Per our discussion, see if you can remake the LMAdapter to be a type of Block. Then, you can load it using Model.load() rather than requiring the special handling in Engine.

[KexinFeng]

Type change to Block is done. But the special handling in Engine is not avoided. GPT2PtLMBlock depends on module pytorch-engines.main and can not be used in engine agnostic frontend. The reason GPT2PtLMBlock is engine specific is that it adapts to different engines and that engine-specific types are used, like IValue for pytorch.

[zachgk]

I'm noticing that these examples are almost identical. Is it possible to refactor out a common helper. So, each one would just initialize the LMAdapter and then pass it into the helper. Maybe the helper would look something like public static void generateInternal(LMAdapter generator);

[KexinFeng]

The testing data used here, like inputIds, positionIds and attentionMask, are more of case by case. It's like those test examples in NDIndexTest.java. They are not likely to be reused. The rest common parts have been factored out.

[KexinFeng]

See the latest PR for the updates.

[KexinFeng]

Merged in a different PR.