The output of Gemma 3 4B for TensorRT and Transformers is not the same, even when using float32

[Alireza3242]

System Info

A100
tensorrt_llm 0.20.0rc3

Who can help?

@kaiyux

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

I ran the Gemma 3 4B model using both TensorRT and Transformers, but I’m not getting identical outputs.

Serving with Transformers:

from transformers import AutoTokenizer, Gemma3ForCausalLM
ckpt = "/home/dev/ra_workspace/ra_workspace/gemma3_4b/"

model = Gemma3ForCausalLM.from_pretrained(
    ckpt, torch_dtype=torch.float32, device_map="auto", attn_implementation="eager"
)

tokenizer = AutoTokenizer.from_pretrained(ckpt)

messages = [
    {"role": "user", "content": [{"type": "text", "text": "سلام چه طوری؟"}]}
]

inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True,
    return_dict=True, return_tensors="pt",
).to(model.device)

generation = model.generate(**inputs, max_new_tokens=256, do_sample=False)

Serving with TensorRT-LLM:

In this case, we used version 0.20.0rc3. The model was converted and built using the following arguments:

"convert": {
  "model-dir": model_files + "/model",
  "output-model-dir": "/app/data/tllm_checkpoint",
  "dtype": "float",
  "ckpt-type": "hf"
},
"build": {
  "checkpoint_dir": "/app/data/tllm_checkpoint",
  "output_dir": "/app/model_repository/tensorrt_llm/1",
  "gemm_plugin": "disable",
  "max_batch_size": "32"
}

Serve-time arguments:

"args": {
  "triton_max_batch_size": "128",
  "tokenizer_dir": model_files + "/model",
  "decoupled_mode": "true",
  "engine_dir": "/app/model_repository/tensorrt_llm/1",
  "batch_scheduler_policy": "max_utilization",
  "max_tokens_in_paged_kv_cache": "16384",
  "postprocessing_instance_count": "1",
  "preprocessing_instance_count": "1",
  "bls_instance_count": "1",
  "add_special_tokens": "false",
  "triton_backend": "tensorrtllm",
  "max_queue_delay_microseconds": "1000000",
  "max_beam_width": "1",
  "batching_strategy": "inflight_fused_batching",
  "max_attention_window_size": "1024/,1024/,1024/,1024/,1024/,16384",
  "accumulate_tokens": "",
  "tensorrt_llm_model_name": "tensorrt_llm",
  "max_queue_size": "100000",
  "encoder_engine_dir": "",
  "exclude_input_in_output": "true",
  "encoder_input_features_data_type": "TYPE_BF16",
  "logits_datatype": "TYPE_FP32",
  "xgrammar_tokenizer_info_path": ""
}

message='<bos><start_of_turn>user\nسلام چه طوری؟<end_of_turn>\n<start_of_turn>model\n'
data = {
  "text_input": message,
  "temperature": 0.00,
  "top_k": 1,
  "max_tokens": 1024,
  "stream": False
}

url = 'http://0.0.0.0:8000/v2/models/ensemble/generate_stream'

response = requests.post(url, json=data, headers={"Content-Type": "application/json"})

Expected behavior

exact output

actual behavior

they are not same as each other

additional notes

We have modified some parts of the tensorrt_llm library to make it behave more similarly to transformers. The following changes were made:

The value of eps in the layer_norm within the attention layer was not the same.
The eps value should be changed to 1e-6.

In the file /usr/local/lib/python3.12/dist-packages/tensorrt_llm/functional.py,
inside the function:

def rms_norm(input: Tensor,
             normalized_shape: Union[int, Tuple[int]],
             num_groups: int = 1,
             weight: Optional[Tensor] = None,
             eps: float = 1e-06) -> Tensor:

We modified the section related to with precision("float32") as follows:

with precision("float32"):
    input_dtype = input.dtype
    fp32_input = cast(input, "float32")
    varx = pow(fp32_input, 2.0)

    varx = varx.mean(dim=dim, keepdim=True)
    denom = varx + eps
    denom = denom.sqrt()
    fp32_y = fp32_input / denom
    if num_groups > 1:
        fp32_y = fp32_y.view(old_shape)

    if weight is not None:
        fp32_y = fp32_y * weight
    y = cast(fp32_y, input_dtype)

3-
In the file /usr/local/lib/python3.12/dist-packages/tensorrt_llm/models/gemma/convert.py,
in the following section:

elif any(keyword in name for keyword in (
        "pre_attention_norm.scale",
        "pre_ffw_norm.scale",
        "final_norm.scale",
        "pre_attention_norm/vars/0",
        "pre_ffw_norm/vars/0",
        "rms_normalization/vars/0",
        "input_layernorm",
        "post_attention_layernorm",
        "pre_feedforward_layernorm",
        "post_feedforward_layernorm",
        "model.norm.weight",
        "q_norm.weight",
        "k_norm.weight",
)):
    param = param + 1.0  # upcasted to float32 in case of bfloat16
    add_trt_llm_weight(weights, trt_llm_name, param, 'float32')

The value of trt_llm_config.dtype was changed to "float32".

4-
In /usr/local/lib/python3.12/dist-packages/tensorrt_llm/models/gemma/model.py,
we changed:

self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size,
                               eps=config.norm_epsilon,
                               dtype=config.dtype)

to:

self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size,
                               eps=config.norm_epsilon,
                               dtype="float32")

However, despite all these changes, the outputs are still different.

[VALLIS-NERIA]

@brb-nv Can you take a look at this Gemma 3 issue?
@Alireza3242 We just added Gemma 3 support of pytorch backend (at main branch), can you please give a try?

[Alireza3242]

@VALLIS-NERIA
I also loaded Gemma 3 using the code below, but the outputs are still different from those of Transformers.

from tensorrt_llm import SamplingParams
from tensorrt_llm._torch import LLM

def main():
    prompts = [
        # '<bos><start_of_turn>user\nسلام چه طوری؟<end_of_turn>\n<start_of_turn>model\n',
        '<start_of_turn>user\nسلام چه طوری؟<end_of_turn>\n<start_of_turn>model\n',
    ]
    sampling_params = SamplingParams(max_tokens=1000)

    llm = LLM(model='/app/data/gemma-3-4b-it/model', dtype="bfloat16")
    outputs = llm.generate(prompts, sampling_params)

    for i, output in enumerate(outputs):
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"[{i}] Prompt: {prompt!r}, Generated text: {generated_text!r}")

if __name__ == '__main__':
    main()

[brb-nv]

Hi @Alireza3242, we added support mainly for Gemma-3-1b-it and I remember the outputs of that were very close (if not identical) to reference HF implementation with greedy sampling. That said, we're working on adding support for other Gemma3 variants - 4b, 12b, 27b. So, fixing these discrepancies will help those efforts as well.

• Can you please point me to the attention layernorm you're referring to? Are you referring to the q_norm, k_norm from here? If yes, we're using the eps value from config and not hardcoding as such. So, unsure where the discrepancy is coming from.

[Alireza3242]

@brb-nv
This eps:

TensorRT-LLM/tensorrt_llm/layers/attention.py

Line 369 in 039f7e3

eps=1e-05,

[brb-nv]

Hi @Alireza3242, we investigated the discrepancy for Gemma3 1B model in Pytorch backend: #5564

Key findings:

• Gemma3's sliding window isn't left-inclusive in their HF reference implementation. TRTLLM's attention backends (trtllm, flashinfer, vanilla attention(with changes in this MR)) are left-inclusive.
• Subtle difference in position embedding math introduces non-trivial delta in logits. If applying position embeddings is turned off, atol and rtol could be brought down significantly.

Hope this helps.

[liyi-xia]

@Alireza3242 Hi, have you figured out the issue? I am also facing this issue. When I look at intermediate layer result, it is more or less close. But the final logits have diff around 0.02.

[liyi-xia]

I found two methods to reduce diff

1. Configure correct position embedding accordly, Gemma3 use different position embedding for local and non-local layers;
2. Change all norm layer to be float32

Is there any method to further improve accuracy?

[brb-nv]

One other thing: #5773
This is specific to Gemma3 VLMs which use rope scaling.

[liyi-xia]

@brb-nv Thank you for share. However, I am working on TRT backend as I need prompt embedding input. After several fix, the diff is acceptable but still a bit large compared with other TRT-LLM supported models.

Btw, do you know if pytorch backend support prompt embeds (a.k.a prompt table in TRT backend) as input? I saw there is something named as multimodal embedding in llm.py. But is there any complete example or documentation of pytorch backend I could refer to?

[zxcjhg]

Hi @Alireza3242, we investigated the discrepancy for Gemma3 1B model in Pytorch backend: #5564

Key findings:

* Gemma3's sliding window isn't left-inclusive in their HF reference implementation. TRTLLM's attention backends (trtllm, flashinfer, vanilla attention(with changes in this MR)) are left-inclusive.

* Subtle difference in position embedding math introduces non-trivial delta in logits. If applying position embeddings is turned off, atol and rtol could be brought down significantly.

Hope this helps.

Hey @brb-nv, thanks for the breakdown. I am just getting started with TensorRT LLM for my project and was thinking of either monkey patching or putting together a PR to apply the same changes to the TRT workflow.

Just wanted to check first if this can be done in the open source parts or if the logic is in the closed source side of TensorRT?

[geraldstanje]

any infos about that?