Refactor inference.py for LLM and RoBERTa support

[andrea-fasoli]

This PR implements a substantial refactoring of inference.py which becomes the single entry point for LLMs and RoBERTa models. Support covers non-quantized, GPTQ W4A16, and INT8 models.

inference.py code has been streamlined. It is now structured into the following sections:

args_parsing            define script arguments across all model configurations
aiu_setup               set up AIU environment variables
model_setup             define model dtype, device, and distributed strategy
quantization_setup      import FMS-MO addons and define linear_config for FMS
direct_quantization     quantize a non-quantized model to INT8 (WIP)
decoders                run token generation task with LLMs
encoders                run QA or MLM task with RoBERTa

Extensive code validation is needed prior merging.

[andrea-fasoli]

cc for review: @ani300 @JRosenkranz

[ani300]

maybe it's no longer needed? if you can't find any reference to it feel free to delete

[ani300]

this is only needed if compile_dynamic is disabled, can we gate it?

[ani300]

I think some of these are already set by default on the e2e_stable image, can we check and remove the ones we don't need anymore?

[ani300]

confirmed these are no longer needed

[andrea-fasoli]

is os.environ.setdefault("DTCOMPILER_KEEP_EXPORT", "true") still needed or not? It's the env var that was set after these ones

[ani300]

this one is only needed for decoder models, for roberta it will probably make it not work

[ani300]

confirmed it will not work for roberta, we need to set it depending on the kind of model

[ani300]

I think VFIO is now PF or VF, depending on the AIU setup, we can set PF by default as most cards are running in PF mode

[ani300]

confirmed it is now PF for all clusters we have access to and will eventually be VF

[andrea-fasoli]

I set it to PF but do we need an argument for this?
One using: choices=["VF", "PF"], default="PF"

[ani300]

this comment can be upgraded to "If set to True, this will unfuse any fused weights in the model" as the way it's done doesn't involve "unfuse_weights" anymore

[ani300]

run seed can also be relevant for randomly initialized models, which we sometimes use

[ani300]

only for cpu/cuda, aiu doesn't have any change

[ani300]

where is this used?

[andrea-fasoli]

mostly in the quantization functions of int8 roberta for now (printing out model parameters if so desired - it's crucial for debug), but I supposed it could be a useful flag for decoders too. I am not using the count functionality at this time, could be just a True/False flag.

[ani300]

isn't this repeated? or at least very similar to the decoder arguments. It might be worth using argument groups (https://docs.python.org/3/library/argparse.html#argument-groups) and making the roberta and decoder-specific arguments mutually exclusive based on what you pass to some other argument

[andrea-fasoli]

there is some overlap indeed.
I combined max_seq_length (enc) with max_prompt_len (dec), as they share the same meaning (although the default values differed).
pad_to_max_length (enc) conceptually overlaps with min_pad_length and fixed_prompt_length (dec), but the first is a boolean while the others are int. I couldn't find a clean way to implement a 3-way exclusivity but I'll add some argument validation at the time of loading encoder vs. decoder.

[ani300]

maybe add a real "is_encoder" argument and that would help with the mutually exclusive argument groups

[ani300]

make this import gated and fail only if args.is_quantized is True in inference.py

[ani300]

this is now deprecated in pytorch, only assume and automatic are needed

[ani300]

please add a TODO to add FP8 inference once that lands too

[ani300]

add the int8 prefix

[ani300]

add the int8 prefix

[ani300]

there is a new --compile_dynamic_sendnn that needs to be added here

[JRosenkranz]

Now that we are cleaning up the args, I wonder if it would make sense to combine compile_dynamic and compile_dynamic_sendnn. What would happen if a user does compile_dynamic_sendnn and not compile_dynamic? We may want to make this something like --compile_dynamic=<static_inputs, symbolic_inputs>

[ani300]

should we add a prefix to mark these as encoder-specific?

[ani300]

These are wrong types

[ani300]

There's plenty of things to change, and I'd like to think more on the general architecture... Given how different the parameters and general flow is for both encoder tasks and decoders, it might be worth splitting inference.py into encoders.py and decoders.py. Most of the code can be reused anyways if the arguments are groups using the API for this and then just picking the relevant groups for each script.

Of course, documentation (README.md) also needs to be updated for this new structure.

[JRosenkranz]

should we use min_pad_length here? If the min_pad_length is not specified, then it will implicitly pad to max length. If min_pad_length is specified and is larger than the largest sequence, this will add extra pads (if we want to simulate a different sequence length)

[andrea-fasoli]

I put some thoughts into this and it is surprisingly tricky because for encoders the tokenization is performed under the hood of a transformers' PretrainedTokenizer, which does not handle truncation and padding the same way as FMS (i.e, with explicit calls to FMS truncate_prompts_to_max_length and pad_input_ids).

PretrainedTokenizer receives a max_length argument for truncation AND padding, and a padding argument which can be True (pad to max sequence), False (do not pad), or string "max_length" (which will pad to the max_length argument).

To make it behave like our decoder tokenization, we would need to adjust max_length and padding based on the tokenized sequence length... which is not known yet at the time of PretrainedTokenizer call.

We could eventually change the whole tokenization and feature preparation process (which right now is mostly based on a pytorch example, would be nice to rework it), but for the time being I would keep pad_to_max_length argument for encoders only.

[JRosenkranz]

Is this handling the case where tokenizer.bos_token_id != tokenizer.eos_token_id?

[andrea-fasoli]

I reproduced this from inference.py. Only difference is the location of where self.add_special_tokens is defined. Now it's updated at the beginning of process_eval_set (before any ids_for_prompt call) as:

        self.add_special_tokens = (
            self.tokenizer.bos_token_id != self.tokenizer.eos_token_id
        )

Should be correct.

[JRosenkranz]

We maybe able to reuse the function ids_for_prompt in utils/__init__.py

[andrea-fasoli]

good catch, it's the same function. It is duplicated also in the current inference.py...

[andrea-fasoli]

at a second thought, this function may belong to the DecoderInfer class, as the encoders use a different approach. I could remove the duplicate from utils/__init__.py

[andrea-fasoli]

never mind, ids_for_prompt in utils/__init__.py is also used to generate ids for testing. inference.py and validation.py duplicate this function though. We should sort this out

[JRosenkranz]

can we use the same arguments for this as decoder rather than introducing pad_to_max_length?

[andrea-fasoli]

see answer above regarding tokenization in encoders