I am trying to run a simple MLP on A100 GPU with multi-dimensional batch inputs of shape (b1, b2, input_dim) and output shape (b1, b2, 1). Flax outputs when passing the entire input (b1, b2, input_dim) v/s passing a single input (1, 1, input_dim) iteratively are not matching. When I run the same code example on CPU or run the equivalent PyTorch version, it matches exactly. Please see the minimal code example, colab link and outputs below.
System information
OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04.6 LTS
Flax, jax, jaxlib versions (obtain with pip show flax jax jaxlib:
Output mismatch while running the following minimal example on GPU:
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import numpy as np
# Flax imports
import jax.random as jr
import jax.numpy as jnp
import flax.linen as nn
# PyTorch imports
import torch
# Common constants
b1, b2 = 2, 3
input_dim = 2
hidden_dim = 2
output_dim = 1
batch_shape = (b1, b2)
# Flax code
tiny_model = nn.Sequential([nn.Dense(hidden_dim), nn.Dense(output_dim)])
tiny_params = tiny_model.init(jr.PRNGKey(1234), jnp.ones((*batch_shape, input_dim)))
x = jr.normal(jr.PRNGKey(5678), (*batch_shape, input_dim))
batch_out = tiny_model.apply(tiny_params, x)
individual_out = np.zeros_like(batch_out)
for i in range(b1):
for j in range(b2):
individual_out[i:i+1, j:j+1, :] = tiny_model.apply(tiny_params, x[i:i+1, j:j+1, :])
print(f"Flax output match: {jnp.all(batch_out == individual_out)}")
display(batch_out.squeeze().tolist(), individual_out.squeeze().tolist())
# PyTorch code
torch.manual_seed(1234)
model = torch.nn.Sequential(torch.nn.Linear(input_dim, hidden_dim), torch.nn.Linear(hidden_dim, output_dim))
batch_out = model(torch.tensor(x.tolist()))
individual_out = torch.ones_like(batch_out)
for i in range(b1):
for j in range(b2):
individual_out[i, j, :] = model(torch.tensor(x[i:i+1, j:j+1, :].tolist())).squeeze()
print(f"PyTorch output match: {torch.all(batch_out == individual_out)}")
display(batch_out.squeeze().tolist(), individual_out.squeeze().tolist())
What you expected to happen:
Output should match
Steps to reproduce:
Here is a colab link. Surprisingly it did not brake on Google colab on all random experiments. Luckily I found one combination where it fails. However the differences on A100 GPU system are way off compared to colab GPU. Here are the outputs on colab v/s on my A100 GPU system.
Hey @patel-zeel, Flax doesn't have any platform specific code. Please post this issue on the JAX repo, they might be able to better help you solve the problem.
Hi,
I am trying to run a simple MLP on A100 GPU with multi-dimensional batch inputs of shape (b1, b2, input_dim) and output shape (b1, b2, 1). Flax outputs when passing the entire input (b1, b2, input_dim) v/s passing a single input (1, 1, input_dim) iteratively are not matching. When I run the same code example on CPU or run the equivalent PyTorch version, it matches exactly. Please see the minimal code example, colab link and outputs below.
System information
OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04.6 LTS
Flax, jax, jaxlib versions (obtain with
pip show flax jax jaxlib
:Python version: Python 3.9.12
GPU/TPU model and memory: NVIDIA A100-SXM4-80GB
CUDA version (if applicable):
NVIDIA Driver version: 525.85.12
Problem you have encountered:
Output mismatch while running the following minimal example on GPU:
What you expected to happen:
Output should match
Steps to reproduce:
Here is a colab link. Surprisingly it did not brake on Google colab on all random experiments. Luckily I found one combination where it fails. However the differences on A100 GPU system are way off compared to colab GPU. Here are the outputs on colab v/s on my A100 GPU system.
Outputs on colab GPU
Outputs on colab CPU
Output on A100 GPU system's GPU
Output on A100 GPU system's CPU