Large CPU/GPU numerical discrepancy in tf.linalg.matmul for float32 linear layer equivalent #127346

import numpy as np
import tensorflow as tf

for gpu in tf.config.list_physical_devices("GPU"):
    tf.config.experimental.set_memory_growth(gpu, True)

rng = np.random.default_rng(133006)

x_np = (rng.standard_normal((64, 197, 1536)).astype("float32") * 1000.0)
w_np = (rng.standard_normal((384, 1536)).astype("float32") * 50.0)
b_np = (rng.standard_normal((384,)).astype("float32") * 50.0)

with tf.device("/CPU:0"):
    x = tf.constant(x_np)
    w = tf.constant(w_np)
    b = tf.constant(b_np)
    dropped_cpu = tf.keras.layers.Dropout(rate=0.0)(x, training=False)
    y_cpu = tf.linalg.matmul(dropped_cpu, w, transpose_b=True) + b
    y_cpu = y_cpu.numpy()

with tf.device("/GPU:0"):
    x = tf.constant(x_np)
    w = tf.constant(w_np)
    b = tf.constant(b_np)
    dropped_gpu = tf.keras.layers.Dropout(rate=0.0)(x, training=False)
    y_gpu = tf.linalg.matmul(dropped_gpu, w, transpose_b=True) + b
    y_gpu = y_gpu.numpy()

dropout_diff = np.max(np.abs(dropped_cpu.numpy() - dropped_gpu.numpy()))
linear_diff = np.max(np.abs(y_cpu - y_gpu))
mean_diff = np.mean(np.abs(y_cpu - y_gpu))
max_output = np.max(np.abs(y_cpu))

print("TensorFlow:", tf.__version__)
print("GPUs:", tf.config.list_physical_devices("GPU"))
print("output shape:", y_cpu.shape)
print("dropout CPU/GPU max abs diff:", dropout_diff)
print("linear CPU/GPU max abs diff:", linear_diff)
print("linear CPU/GPU mean abs diff:", mean_diff)
print("max abs output:", max_output)
print("relative max diff:", linear_diff / max_output)

TensorFlow: 2.20.0
GPUs: [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')]
output shape: (64, 197, 384)
dropout CPU/GPU max abs diff: 0.0
linear CPU/GPU max abs diff: 3066.0
linear CPU/GPU mean abs diff: 459.5823
max abs output: 10401001.0
relative max diff: 0.00029477