[Python] Replicating GPU environment across architectures

I am training a Mamba model on two different GPU architectures: RTX 4090 and RTX A6000. Despite setting all random seeds and using deterministic algorithms, I am observing significant non-deterministic behavior in the training process on the RTX 4090, while the RTX A6000 exhibits almost deterministic results. This issue is due to the use of atomic adds in the backward pass of the Mamba model, as discussed in this GitHub issue.

I would like to maintain the current hyperparameters, as they provide the best performance. Therefore, I am looking for any tricks or settings that can make the RTX 4090 behave more like the RTX A6000 environment to achieve consistent results. Here are the details of my setup:

RTX4090

| NVIDIA-SMI 545.23.08 Driver Version: 545.23.08 CUDA Version: 12.3 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA GeForce RTX 4090 On | 00000000:26:00.0 Off | Off |
| 30% 22C P8 12W / 450W | 17MiB / 24564MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
| 0 N/A N/A 426218 G /usr/lib/xorg/Xorg 4MiB |
+---------------------------------------------------------------------------------------+


RTXA6000

| NVIDIA-SMI 545.23.08 Driver Version: 545.23.08 CUDA Version: 12.3 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA RTX A6000 On | 00000000:24:00.0 Off | Off |
| 30% 22C P8 25W / 300W | 12MiB / 49140MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
| 0 N/A N/A 7388 G /usr/lib/xorg/Xorg 4MiB |
+---------------------------------------------------------------------------------------+


So far I've tried setting:

import os
import torch

os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
torch.backends.cuda.matmul.allow_tf32 = False
torch.set_default_dtype(torch.float32)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False


Any more help is appreciated.

Continue reading...