When '--gather_all_token_logits' is enabled, the first token appears to be abnormal."

[StarrickLiu]

Problem Desciption:

When building the engine with the '--gather_all_token_logits' option, there seems to be an issue.

If constructed with '--gather_all_token_logits', there is a high probability of garbled characters appearing in the first token.

However, if built without '--gather_all_token_logits' while keeping other parameters consistent, the first token appears normal.

llama1 7B:

With --gather_all_token_logits'

Build Command:

python3 build.py --model_dir=/path/to/llama-7b-hf/ \
                  --dtype bfloat16 \
                  --use_gpt_attention_plugin bfloat16 \
                  --use_gemm_plugin bfloat16 \
                  --output_dir /path/to/llama-7b-trt/0.6.1-cf-pe1-gatl-mb-bf16-8_gpu-8k-2k-bs4 \
                  --world_size 8 \
                  --tp_size 8 \
                  --max_input_len 8192 \
                  --max_output_len 2048 \
                  --max_batch_size 4 \
                  --remove_input_padding \
                  --enable_context_fmha \
                  --parallel_build \
                  --multi_block_mode \
                  --gather_all_token_logits \
                  --use_parallel_embedding \
                  --embedding_sharding_dim 1

Without --gather_all_token_logits'

Build Command:

python3 build.py --model_dir=/path/to/llama-7b-hf/ \
                  --dtype bfloat16 \
                  --use_gpt_attention_plugin bfloat16 \
                  --use_gemm_plugin bfloat16 \
                  --output_dir /path/to/llama-7b-trt/0.6.1-cf-pe1-mb-bf16-8_gpu-8k-2k-bs4 \
                  --world_size 8 \
                  --tp_size 8 \
                  --max_input_len 8192 \
                  --max_output_len 2048 \
                  --max_batch_size 4 \
                  --remove_input_padding \
                  --enable_context_fmha \
                  --parallel_build \
                  --multi_block_mode \
                  --use_parallel_embedding \
                  --embedding_sharding_dim 1

Worth noting is that this issue has been tested in previous versions as well as in version 0.6.1.

[byshiue]

It should be a bug of 0.6.1, and should be fixed in latest main branch. Please take a try.

[StarrickLiu]

In testing with the new version, everything is fine. Thank you.

[salaki]

@StarrickLiu, wondering if you successfully got the logtis. Is the logits for outputed tokens or all token for vocab?