fabric precision: `32` bit using far less memory than `bf16`

Question

fabric precision: `32` bit using far less memory than `bf16`

Closed this issue 5 months ago · 12 comments

Hi, was playing with fabric precision and noticed that bf16 uses quite a bit more memory than 32 (23.83Gi vs 13.37Gi). It also trains a little bit slower (1.77 Hz vs 1.91 Hz)
Wondering if it is expected behavior, or something's wrong somewhere?

Config:

defaults:
  - dreamer_v3
  - override /algo: dreamer_v3_XL
  - override /env: atari
  - _self_

# Experiment
seed: 5

# Environment
env:
  num_envs: 6
  max_episode_steps: 27000
  id: MsPacmanNoFrameskip-v4
  input_buffer: 12
  use_input_buf_as_obs: True

# Checkpoint
checkpoint:
  every: 2000

# Buffer
buffer:
  size: 100000
  checkpoint: True

# Algorithm
algo:
  learning_starts: 1024
  total_steps: 1200000
  train_every: 12
  per_rank_batch_size: 8
  per_rank_sequence_length: 98
  mlp_keys:
    encoder: [actions_0, actions_1, actions_2, actions_3, actions_4, actions_5, actions_6, actions_7, actions_8, actions_9, actions_10, actions_11]

fabric:
  precision: bf16

Memory using precision: 32, Training at 1.91 Hz.

Memory using precision: bf16, Training at 1.77 Hz.

Answer 1 · 2024-03-25T07:52:45.000Z

Hi @geranim0, we are looking into the precision because we sometimes escape from it. There's already an open issue (#236) and a branch where I've already made some changes. If you want you can contribute by looking at the branch 🤟 in the next few days I'll comment with all the activities I've done

Answer 2 · 2024-03-27T16:51:53.000Z

Hi @geranim0, can you try out this branch and tell me if anything changes?

Answer 3 · 2024-03-27T20:40:46.000Z

Hi @belerico ,

Thanks for that branch.

Here are the results of the tests (bf16 still uses more memory than 32)

python sheeprl.py fabric.accelerator=gpu fabric.precision=bf16 exp=dreamer_v3_100k_ms_pacman yields 3.90/24 Gi

python sheeprl.py fabric.accelerator=gpu fabric.precision=32 exp=dreamer_v3_100k_ms_pacman yields 3.61/24 Gi

Testing the same experiment but overriding to use the XL model yields 13.53/24 Gi for 32, and 18.79/24 Gi for bf16

Also, this command using 16-true makes it crash
python sheeprl.py fabric.accelerator=gpu fabric.precision=16-true exp=dreamer_v3_100k_ms_pacman

Error executing job with overrides: ['fabric.accelerator=gpu', 'fabric.precision=16-true', 'exp=dreamer_v3_100k_ms_pacman']
Traceback (most recent call last):
  File "/home/sam/dev/sheeprl/sheeprl/cli.py", line 347, in run
    run_algorithm(cfg)
  File "/home/sam/dev/sheeprl/sheeprl/cli.py", line 186, in run_algorithm
    fabric.launch(reproducible(command), cfg, **kwargs)
  File "/home/sam/dev/sheeprl/.venv/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 839, in launch
    return self._wrap_and_launch(function, self, *args, **kwargs)
  File "/home/sam/dev/sheeprl/.venv/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 925, in _wrap_and_launch
    return to_run(*args, **kwargs)
  File "/home/sam/dev/sheeprl/.venv/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 930, in _wrap_with_setup
    return to_run(*args, **kwargs)
  File "/home/sam/dev/sheeprl/sheeprl/cli.py", line 182, in wrapper
    return func(fabric, cfg, *args, **kwargs)
  File "/home/sam/dev/sheeprl/sheeprl/algos/dreamer_v3/dreamer_v3.py", line 712, in main
    train(
  File "/home/sam/dev/sheeprl/sheeprl/algos/dreamer_v3/dreamer_v3.py", line 243, in train
    imagined_prior, recurrent_state = world_model.rssm.imagination(imagined_prior, recurrent_state, actions)
  File "/home/sam/dev/sheeprl/sheeprl/algos/dreamer_v3/agent.py", line 461, in imagination
    _, imagined_prior = self._transition(recurrent_state)
  File "/home/sam/dev/sheeprl/sheeprl/algos/dreamer_v3/agent.py", line 442, in _transition
    return logits, compute_stochastic_state(
  File "/home/sam/dev/sheeprl/sheeprl/algos/dreamer_v2/utils.py", line 63, in compute_stochastic_state
    stochastic_state = dist.rsample() if sample else dist.mode
  File "/home/sam/dev/sheeprl/.venv/lib/python3.10/site-packages/torch/distributions/independent.py", line 104, in rsample
    return self.base_dist.rsample(sample_shape)
  File "/home/sam/dev/sheeprl/sheeprl/utils/distribution.py", line 397, in rsample
    samples = self.sample(sample_shape)
  File "/home/sam/dev/sheeprl/sheeprl/utils/distribution.py", line 365, in sample
    indices = self._categorical.sample(sample_shape)
  File "/home/sam/dev/sheeprl/.venv/lib/python3.10/site-packages/torch/distributions/categorical.py", line 132, in sample
    samples_2d = torch.multinomial(probs_2d, sample_shape.numel(), True).T
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

Answer 4 · 2024-03-30T12:07:32.000Z

Hi @geranim0, the training with 16-true precision is ok to fail, in fact even PyTorch guys suggest to switch to bf16 or to the mixed precision. Regarding the increase in memory we are still investigating it.
Thank you for your patience

Answer 5 · 2024-04-10T07:03:38.000Z

Hi @geranim0, in this branch I've added the possibility to print the CUDA memory allocation since the beginning of the training with torch.cuda.max_memory_allocated().

This is what I've obtained running the following command

python sheeprl.py \
fabric.accelerator=gpu \
exp=dreamer_v3_100k_ms_pacman \
algo.total_steps=1152

fabric.precision=bf16-mixed: 2.57GB
fabric.precision=bf16-true: 2.07GB
fabric.precision=16-mixed: 2.58GB
fabric.precision=16-true: nan in the categorical probs
fabric.precision=32: 5.24GB

This is my env collected with python -m torch.utils.collect_env:

Collecting environment information...
PyTorch version: 2.3.0.dev20240314+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Pop!_OS 22.04 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.11.8 (main, Feb 26 2024, 21:39:34) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.8.0-76060800daily20240311-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4060 Laptop GPU
Nvidia driver version: 550.67
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      39 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             16
On-line CPU(s) list:                0-15
Vendor ID:                          GenuineIntel
Model name:                         13th Gen Intel(R) Core(TM) i5-13500H
CPU family:                         6
Model:                              186
Thread(s) per core:                 2
Core(s) per socket:                 12
Socket(s):                          1
Stepping:                           2
CPU max MHz:                        4700.0000
CPU min MHz:                        400.0000
BogoMIPS:                           6374.40
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb intel_pt sha_ni xsaveopt xsavec xgetbv1 xsaves split_lock_detect user_shstk avx_vnni dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req hfi vnmi umip pku ospke waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm md_clear serialize arch_lbr ibt flush_l1d arch_capabilities
Virtualization:                     VT-x
L1d cache:                          448 KiB (12 instances)
L1i cache:                          640 KiB (12 instances)
L2 cache:                           9 MiB (6 instances)
L3 cache:                           18 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced / Automatic IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] mypy==1.2.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] pytorch-lightning==2.2.1
[pip3] pytorch-triton==3.0.0+989adb9a29
[pip3] torch==2.3.0.dev20240314+cu121
[pip3] torchmetrics==1.3.2
[pip3] torchvision==0.18.0.dev20240314+cu121
[pip3] triton==2.2.0
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] pytorch-lightning         2.2.1                    pypi_0    pypi
[conda] pytorch-triton            3.0.0+989adb9a29          pypi_0    pypi
[conda] torch                     2.3.0.dev20240314+cu121          pypi_0    pypi
[conda] torchmetrics              1.3.2                    pypi_0    pypi
[conda] torchvision               0.18.0.dev20240314+cu121          pypi_0    pypi
[conda] triton                    2.2.0                    pypi_0    pypi

Can you also try out this branch and tell us your number?
Thank you

Answer 6 · 2024-04-10T12:00:40.000Z

Current machine doesn't support all of them, but for 16-mixed I get 5.67 GB

With this config

Collecting environment information...
PyTorch version: 2.2.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1080
Nvidia driver version: 537.34
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      39 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             12
On-line CPU(s) list:                0-11
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Core(TM) i7-8700K CPU @ 3.70GHz
CPU family:                         6
Model:                              158
Thread(s) per core:                 2
Core(s) per socket:                 6
Socket(s):                          1
Stepping:                           10
BogoMIPS:                           7391.99
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology cpuid pni pclmulqdq vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti ssbd ibrs ibpb stibp tpr_shadow vnmi ept vpid ept_ad fsgsbase bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap clflushopt xsaveopt xsavec xgetbv1 xsaves md_clear flush_l1d arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          192 KiB (6 instances)
L1i cache:                          192 KiB (6 instances)
L2 cache:                           1.5 MiB (6 instances)
L3 cache:                           12 MiB (1 instance)
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled
Vulnerability L1tf:                 Mitigation; PTE Inversion; VMX conditional cache flushes, SMT vulnerable
Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:             Mitigation; PTI
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Retbleed:             Mitigation; IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; IBRS, IBPB conditional, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Unknown: Dependent on hypervisor status
Vulnerability Tsx async abort:      Mitigation; Clear CPU buffers; SMT Host state unknown

Versions of relevant libraries:
[pip3] mypy==1.2.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] numpy-ringbuffer==0.2.2
[pip3] pytorch-lightning==2.2.0.post0
[pip3] torch==2.2.1
[pip3] torchmetrics==1.3.1
[pip3] triton==2.2.0
[conda] Could not collect

Answer 7 · 2024-04-10T13:36:05.000Z

I can try on windows+wsl too! With fp32 how much memory consumes?

Answer 8 · 2024-04-10T14:05:16.000Z

Sure, for fp32, getting 5.24 GB.

Answer 9 · 2024-04-10T19:10:35.000Z

Hi @geranim0, I've tried also on windows+wsl and I've obtained the following:

fabric.precision=bf16-mixed: 2.57GB
fabric.precision=bf16-true: 2.08GB
fabric.precision=16-mixed: 2.57GB
fabric.precision=32: 5.24GB

with the following env:

/home/belerico/miniconda3/envs/sheeprl-p10/lib/python3.10/runpy.py:126: RuntimeWarning: 'torch.utils.collect_env' found in sys.modules after import of package 'torch.utils', but prior to execution of 'torch.utils.collect_env'; this may result in unpredictable behaviour
  warn(RuntimeWarning(msg))
Collecting environment information...
PyTorch version: 2.2.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4060 Laptop GPU
Nvidia driver version: 551.61
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      39 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             16
On-line CPU(s) list:                0-15
Vendor ID:                          GenuineIntel
Model name:                         13th Gen Intel(R) Core(TM) i5-13500H
CPU family:                         6
Model:                              186
Thread(s) per core:                 2
Core(s) per socket:                 8
Socket(s):                          1
Stepping:                           2
BogoMIPS:                           6374.40
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology tsc_reliable nonstop_tsc cpuid pni pclmulqdq vmx ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves avx_vnni umip waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm serialize flush_l1d arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          384 KiB (8 instances)
L1i cache:                          256 KiB (8 instances)
L2 cache:                           10 MiB (8 instances)
L3 cache:                           18 MiB (1 instance)
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] mypy==1.2.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.3
[pip3] pytorch-lightning==2.2.0.post0
[pip3] torch==2.2.1
[pip3] torchmetrics==1.2.1
[pip3] triton==2.2.0
[conda] numpy                     1.26.3                   pypi_0    pypi
[conda] pytorch-lightning         2.2.0.post0              pypi_0    pypi
[conda] torch                     2.2.1                    pypi_0    pypi
[conda] torchmetrics              1.2.1                    pypi_0    pypi
[conda] triton                    2.2.0                    pypi_0    pypi

It could be the different NVIDIA driver version? Have you tried to update it?

Answer 10 · 2024-04-10T20:17:28.000Z

Updated to the new driver version, unfortunately same results.

Not so sure what's going on there. Will try on the 4090 tomorrow with updated drivers to see if it requires both the new driver and a new card.

Answer 11 · 2024-04-10T20:46:39.000Z

Can you post here the output of your pip freeze please? I can try with the exact same environment as yours tomorrow.

Edit: maybe it's better to use pipreqs?

Answer 12 · 2024-04-10T21:07:04.000Z

(.venv) sam@sam:~/dev/ml/sheeprl$ pip freeze
absl-py==2.1.0
aiohttp==3.9.3
aiosignal==1.3.1
ale-py==0.8.1
antlr4-python3-runtime==4.9.3
asttokens==2.4.1
async-timeout==4.0.3
attrs==23.2.0
autoflake==2.2.1
AutoROM==0.4.2
AutoROM.accept-rom-license==0.6.1
backcall==0.2.0
beautifulsoup4==4.12.3
black==23.12.1
bleach==6.1.0
box2d-py==2.3.5
certifi==2024.2.2
cfgv==3.4.0
charset-normalizer==3.3.2
click==8.1.7
cloudpickle==3.0.0
contourpy==1.2.0
coverage==7.4.3
crafter==1.8.3
cycler==0.12.1
Cython==0.29.37
decorator==4.4.2
defusedxml==0.7.1
distlib==0.3.8
docopt==0.6.2
etils==1.7.0
exceptiongroup==1.2.0
executing==2.0.1
Farama-Notifications==0.0.4
fastjsonschema==2.19.1
filelock==3.13.1
fonttools==4.49.0
frozenlist==1.4.1
fsspec==2024.2.0
glfw==2.7.0
grpcio==1.62.0
gymnasium==0.29.1
hydra-core==1.3.0
identify==2.5.35
idna==3.6
imageio==2.34.0
imageio-ffmpeg==0.4.9
importlib_resources==6.1.2
iniconfig==2.0.0
ipython==8.12.3
isort==5.13.2
jedi==0.19.1
Jinja2==3.1.3
jsonschema==4.21.1
jsonschema-specifications==2023.12.1
jupyter_client==8.6.1
jupyter_core==5.7.2
jupyterlab_pygments==0.3.0
kiwisolver==1.4.5
lightning==2.1.4
lightning-utilities==0.9.0
lz4==4.3.3
Markdown==3.5.2
markdown-it-py==3.0.0
MarkupSafe==2.1.5
matplotlib==3.8.3
matplotlib-inline==0.1.6
mdurl==0.1.2
melee==0.40.1
melee-env @ git+https://github.com/geranim0/melee-env.git@69013bad870b5c1044a30daabc45cd02da2262ed
mistune==3.0.2
moviepy==1.0.3
mpmath==1.3.0
mujoco==3.1.2
multidict==6.0.5
mypy==1.2.0
mypy-extensions==1.0.0
nbclient==0.10.0
nbconvert==7.16.3
nbformat==5.10.4
networkx==3.2.1
nodeenv==1.8.0
numpy==1.26.4
numpy-ringbuffer==0.2.2
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==8.9.2.26
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.19.3
nvidia-nvjitlink-cu12==12.3.101
nvidia-nvtx-cu12==12.1.105
omegaconf==2.3.0
opencv-python==4.8.0.76
opensimplex==0.4.5
packaging==23.2
pandocfilters==1.5.1
parso==0.8.4
pathspec==0.12.1
pexpect==4.9.0
pickleshare==0.7.5
pillow==10.2.0
pipreqs==0.5.0
platformdirs==4.2.0
pluggy==1.4.0
pre-commit==3.5.0
proglog==0.1.10
prompt-toolkit==3.0.43
protobuf==4.25.3
ptyprocess==0.7.0
pure-eval==0.2.2
py-ubjson==0.16.1
pyenet==1.3.17
pyflakes==3.2.0
pygame==2.5.2
Pygments==2.17.2
PyOpenGL==3.1.7
pyparsing==3.1.1
pyserial==3.5
pytest==7.3.1
pytest-cov==4.1.0
pytest-cover==3.0.0
pytest-coverage==0.0
pytest-timeout==2.1.0
python-dateutil==2.9.0.post0
python-dotenv==1.0.1
pytorch-lightning==2.2.0.post0
PyYAML==6.0.1
pyzmq==25.1.2
referencing==0.34.0
requests==2.31.0
rich==13.5.3
rpds-py==0.18.0
ruamel.yaml==0.18.6
ruamel.yaml.clib==0.2.8
ruff==0.1.11
sheeprl @ file:///home/sam/dev/ml/sheeprl
Shimmy==0.2.1
six==1.16.0
soupsieve==2.5
stack-data==0.6.3
swig==4.2.1
sympy==1.12
tensorboard==2.16.2
tensorboard-data-server==0.7.2
tinycss2==1.2.1
tomli==2.0.1
torch==2.2.2
torchmetrics==1.3.1
tornado==6.4
tqdm==4.66.2
traitlets==5.14.2
triton==2.2.0
typing_extensions==4.10.0
urllib3==2.2.1
virtualenv==20.25.1
wcwidth==0.2.13
webencodings==0.5.1
Werkzeug==3.0.1
yarg==0.1.9
yarl==1.9.4
zipp==3.17.0

Also tried pipreqs but it fails somehow with an invalid byte

Edit:
Just tried on the 4090 with the 550 (updated from 535) drivers, and it works! Guess it takes both the new drivers and a new card.
Thanks for the help 💯

Edit again:
Actually, while getting the same results as you for the cuda allocated memory, the GPU is stilled getting filled up more with bf16-mixed than for 32-true. bf16-true works though, and works faster and with less memory than 32-true.