Different outputs when run on CPU vs GPU (CUDA)

Describe the issue

I am attempting to export a model from HuggingFace from PyTorch to Onnx. After exporting the model, I am trying to confirm the outputs are still correct however it appears that when executing the model on the GPU using the CUDAExecutionProvider the outputs of the model are not close enough to the target embeddings produced by the model before exporting. When executing the model on the CPU however, the model does pass the test.

Seems similar to issue #4488 but maybe a new CUDA version or something re-triggered it?

OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 22.04.4 LTS running on WSL2 via Windows 11 64bit
CPU: 13th Gen Intel(R) Core(TM) i9-13900K 3.00 GHz
GPU: Nvidia GeForce RTX 4090 24GB
ONNX Runtime installed from (source or binary): binary (by pip install onnxruntime-gpu)
ONNX Runtime version: 1.19.0
Python version: 3.8.19

To reproduce

`import torch import onnxruntime

import torch.nn.functional as F import numpy as np

from sentence_transformers import SentenceTransformer from transformers import AutoTokenizer, AutoModel

def mean_pooling(last_hidden_state, attention_mask): '''Apply a mean pooling operation to the last hidden state output by the model'''

input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
return torch.sum(last_hidden_state * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

def main():

#Create a large input to make sure we hit the 384 max window before exporting the model
sentences = ['All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy. \
              All work and no play makes Jack a dull boy.']

#Create gold embeddings using the SentenceTransformer module
sent_tran_model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2').to('cuda')
gold_embed = torch.tensor(sent_tran_model.encode(sentences))

#Load the tokenizer and model from HuggingFace
tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-mpnet-base-v2')
model = AutoModel.from_pretrained('sentence-transformers/all-mpnet-base-v2').to('cuda')

#Tokenize the input before passing it into the model
encoded_input = tokenizer(sentences,
                          padding = True,
                          truncation = True,
                          return_tensors = 'pt',
                          max_length = 384).to('cuda')

#Export the model to Onnx
torch.onnx.export(model,
                  (encoded_input['input_ids'], encoded_input['attention_mask']),
                  'all-mpnet-base-v2.onnx',
                  export_params = True,
                  opset_version = 13,
                  do_constant_folding = True,
                  input_names = ['input_ids',
                                 'attention_mask'],
                  output_names = ['last_hidden_state',
                                  'pooler_output'],
                  dynamic_axes = {'input_ids': {0: 'batch_size'},
                                  'attention_mask': {0: 'batch_size'},
                                  'last_hidden_state': {0: 'batch_size'},
                                  'pooler_output': {0: 'batch_size'}})

####################################################################################################
# Run the Onnx model on the CPU and show the embeddings are close to the target embeddings
####################################################################################################

#Run the Onnx model on the CPU
ort_session_cpu = onnxruntime.InferenceSession('all-mpnet-base-v2.onnx', providers = ['CPUExecutionProvider'])
ort_outs_cpu = ort_session_cpu.run(('last_hidden_state', 'pooler_output'), 
                                   {k: v.numpy(force = True) for k, v in encoded_input.items()})
ort_outs_cpu = [torch.tensor(i) for i in ort_outs_cpu]

#Apply the mean pooling operation and normalization operation as described on the HuggingFace page for the model
pool_output_cpu = mean_pooling(ort_outs_cpu[0], encoded_input['attention_mask'].to('cpu'))
sent_embed_cpu = F.normalize(pool_output_cpu, p = 2, dim = 1)

#Assert the CPU embeddings are close to the targets
torch.testing.assert_close(gold_embed, sent_embed_cpu)

####################################################################################################
# Run the Onnx model on the GPU and show the embeddings are NOT close to the target embeddings
####################################################################################################

#Run the Onnx model on the GPU
ort_session_cpu = onnxruntime.InferenceSession('all-mpnet-base-v2.onnx', providers = ['CUDAExecutionProvider'])
ort_outs_cpu = ort_session_cpu.run(('last_hidden_state', 'pooler_output'), 
                                   {k: v.numpy(force = True) for k, v in encoded_input.items()})
ort_outs_cpu = [torch.tensor(i).to('cuda') for i in ort_outs_cpu]

#Apply the mean pooling operation and normalization operation as described on the HuggingFace page for the model
pool_output_cpu = mean_pooling(ort_outs_cpu[0], encoded_input['attention_mask'].to('cuda'))
sent_embed_cpu = F.normalize(pool_output_cpu, p = 2, dim = 1)

#Assert the GPU embeddings are close to the targets
torch.testing.assert_close(gold_embed.to('cuda'), sent_embed_cpu)

if name == "main": main()`

Urgency

Somewhat urgent, attempting to optimize a model to use Onnx so I can use it in Nvidia Triton.

Platform

Linux

OS Version

22.04.4

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.19.0

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Default CPU, CUDA

Execution Provider Library Version

CUDA 12.6

microsoft / onnxruntime