RoBERTa model conversion does not pass the huggingface test

[arinaruck]

System Info

CPU architecture: x86_64
GPU type: NVIDIA A100-SXM4-40GB
CUDA Version: 12.7
Driver Version: 565.57.01

Who can help?

Hey, @byshiue!
I saw you responding to other Encoder Model related issue, hope you might be the right person for this question.

Thank you!

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Hello!
I am trying to use the TensorRT-LLM/examples/bert to convert a RoBERTa model (FacebookAI/roberta-base) to a TRT-LLM engine.
I am using v.0.16.0 tag and am following the instructions from TensorRT-LLM/examples/bert/README.md

• First, as a sanity check I verify that converting a BERT model (google-bert/bert-base-uncased) passes the run.py huggingface comparisson test (with intermediate checks). With the following commands

CUDA_VISIBLE_DEVICES=0 python /code/tensorrt_llm/examples/bert/convert_checkpoint.py --model=BertModel  --model_dir=google-bert/bert-base-uncased   --output_dir=trt_checkpoints/bert-base

CUDA_VISIBLE_DEVICES=0 trtllm-build --checkpoint_dir trt_checkpoints/bert-base/ --output_dir engines/bert-base --remove_input_padding=disable --max_batch_size=128 --max_seq_len=512  --bert_attention_plugin=disable --context_fmha=disable --enable_debug_output

CUDA_VISIBLE_DEVICES=0 python /code/tensorrt_llm/examples/bert/run.py --engine_dir engines/bert-base/  --hf_model_dir=google-bert/bert-base-uncased --run_hf_test  --debug

This results with final hidden outputs as well as intermediate layer outputs passing the torch.all_close checks with the huggingface model outputs (as implemented in run.py)

• Next, I try to perform the same comparison with a RobertaModel

CUDA_VISIBLE_DEVICES=0 python /code/tensorrt_llm/examples/bert/convert_checkpoint.py --model=RobertaModel  --model_dir=FacebookAI/roberta-base --output_dir=trt_checkpoints/roberta-base 

CUDA_VISIBLE_DEVICES=0 trtllm-build --checkpoint_dir trt_checkpoints/roberta-base/ --output_dir engines/roberta-base --remove_input_padding=disable --max_batch_size=128 --max_seq_len=512  --bert_attention_plugin=disable --context_fmha=disable --enable_debug_output

CUDA_VISIBLE_DEVICES=0 python /code/tensorrt_llm/examples/bert/run.py --engine_dir engines/roberta-base/  --hf_model_dir=FacebookAI/roberta-base --run_hf_test  --debug

Even though I get a pass on the final check results:RobertaModel result is all close to HF reference!, I observe that the intermediate layer outputs do not match, starting from the 4th layer with default tolerance = 1e-2 (or from the 0th encoder layer with a lower tolerance of 1e-3). Here is the output of the default script:

 Embedding are all close                                                                                                       
 BertEncoderLayer_0_output is close: True                                                                                      
 BertEncoderLayer_1_output is close: True                                                                                      
 BertEncoderLayer_2_output is close: True                                                                                      
 BertEncoderLayer_3_output is close: True                                                                                      
 BertEncoderLayer_4_output is close: False                                                                                     
 BertEncoderLayer_5_output is close: True                                                                                      
 BertEncoderLayer_6_output is close: False                                                                                     
 BertEncoderLayer_7_output is close: False                                                                                     
 BertEncoderLayer_8_output is close: False                                                                                     
 BertEncoderLayer_9_output is close: False                                                                                     
 BertEncoderLayer_10_output is close: False     

• When I try to use a fine-tuned roberta-base checkpoint I encounter both the intermediate and the final checks failing

Expected behavior

Getting the result of
CUDA_VISIBLE_DEVICES=0 python /code/tensorrt_llm/examples/bert/run.py --engine_dir engines/roberta-base/ --hf_model_dir=FacebookAI/roberta-base --run_hf_test --debug as:

 Embedding are all close                                                                                                       
 BertEncoderLayer_0_output is close: True                                                                                      
 BertEncoderLayer_1_output is close: True                                                                                      
 BertEncoderLayer_2_output is close: True                                                                                      
 BertEncoderLayer_3_output is close: True                                                                                      
 BertEncoderLayer_4_output is close: True                                                                                   
 BertEncoderLayer_5_output is close: True                                                                                      
 BertEncoderLayer_6_output is close: True                                                                                     
 BertEncoderLayer_7_output is close: True                                                                                     
 BertEncoderLayer_8_output is close: True                                                                                   
 BertEncoderLayer_9_output is close: True                                                                                     
 BertEncoderLayer_10_output is close: True     

as well as the outputs of the final layer:

RobertaModel result is all close to HF reference!

actual behavior

 Embedding are all close                                                                                                       
 BertEncoderLayer_0_output is close: True                                                                                      
 BertEncoderLayer_1_output is close: True                                                                                      
 BertEncoderLayer_2_output is close: True                                                                                      
 BertEncoderLayer_3_output is close: True                                                                                      
 BertEncoderLayer_4_output is close: False                                                                                     
 BertEncoderLayer_5_output is close: True                                                                                      
 BertEncoderLayer_6_output is close: False                                                                                     
 BertEncoderLayer_7_output is close: False                                                                                     
 BertEncoderLayer_8_output is close: False                                                                                     
 BertEncoderLayer_9_output is close: False                                                                                     
 BertEncoderLayer_10_output is close: False     

additional notes

Based on the EmbedderLayer outputs matching in all cases, I would expect the difference be in the EncoderLayer, however based on the huggingface implementations those should be exactly the same for BERT and RoBERTa modelling_bert.py, modelling_roberta.py

I would appreciate any help or pointers on how to debug this and make the RobertaModel work with TensorRT-LLM.

[arinaruck]

Hey @symphonylyh, @juney-nvidia, I noticed you were both involved in the original BERT + RoBERTa integration (#778).
Tagging you as well, in case it might be relevant. Thanks!

[symphonylyh]

Hi @arinaruck , have you tried with older versions and noticed the accuracy issue only after 0.16.0, or 0.16.0 is the first one you try?

[symphonylyh]

two questions: (1) is there a reason why you disable the bert attention plugin and context fmha etc.? (2) how would your accuracy check look like if you try fp32?

[arinaruck]

Hey! @symphonylyh Thanks for your reply!

0.16.0 is the first version I tried

(1) I disabled all the plugins to go with the "simplest" setup, I was observing similar issues with e.g --bert_attention_plugin=float16 or --bert_attention_plugin=float16 --context_fmha=enable, with the intermediate check outputs:

Embedding are all close
BertEncoderLayer_0_output is close: True
BertEncoderLayer_1_output is close: True
BertEncoderLayer_2_output is close: True
BertEncoderLayer_3_output is close: True
BertEncoderLayer_4_output is close: False
BertEncoderLayer_5_output is close: False
BertEncoderLayer_6_output is close: False
BertEncoderLayer_7_output is close: False
BertEncoderLayer_8_output is close: False
BertEncoderLayer_9_output is close: False
BertEncoderLayer_10_output is close: False

With one extra output flipped to False (BertEncoderLayer_5_output).

(2)
If I switch dtype to float32 by passing it in the convert checkpoint

CUDA_VISIBLE_DEVICES=0 python /code/tensorrt_llm/examples/bert/convert_checkpoint.py --model=RobertaModel  --model_dir=FacebookAI/roberta-base --output_dir=trt_checkpoints/roberta-base --dtype=float32

(all the other script runs remain unchanged)
I get the same output of the intermediate checks:

Embedding are all close
BertEncoderLayer_0_output is close: True
BertEncoderLayer_1_output is close: True
BertEncoderLayer_2_output is close: True
BertEncoderLayer_3_output is close: True
BertEncoderLayer_4_output is close: False
BertEncoderLayer_5_output is close: True
BertEncoderLayer_6_output is close: False
BertEncoderLayer_7_output is close: False
BertEncoderLayer_8_output is close: False
BertEncoderLayer_9_output is close: False
BertEncoderLayer_10_output is close: False

(3) Additionally, since the last time I wrote I tried to investigate a little bit myself. I thought the reason might be in the Embeddings, since that's the main thing that distinguishes the BertModel and RobertaModel. I tried registering

self.vocab_embedding(input_ids)
self.position_embedding(position_ids)
self.token_embedding(token_type_ids)

from the BertEmbedding with self.register_network_output and compared those with the embeddings from a huggingface model, but those are exactly the same.
I am not sure what else could it be, looking forward to hearing from you.

Thanks!