Repository: pytorch
Issue: RuntimeError: NVML_SUCCESS == DriverAPI::get()->nvmlInit_v2_() INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":799
Labels: module: cuda,triaged,module: CUDACachingAllocator

### 🐛 Describe the bug

```python 
import torch
import torch.nn as nn
import torch.optim as optim
import os
import argparse
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import ImageFile
from tqdm import tqdm
import numpy as np
from dataset import PairedImagesDataset, sample_paired_images
from model import CroDINO, get_patch_embeddings

def train(model, dino, train_loader, val_loader, device, criterion, optimizer, epochs=1, save_path='untitled', debug=False):
    model.to(device)
    
    model_path = os.path.join('models', save_path)
    os.makedirs(model_path, exist_ok=True)

    best_val_loss = np.inf
    patience_counter = 0
    best_model_path = None

    for epoch in range(epochs):
        model.train()
        running_loss = 0.0

        with tqdm(total=len(train_loader), desc=f'Epoch {epoch+1}/{epochs}', unit='batch') as pbar:
            for ground_image, aerial_image in train_loader:
                ground_image, aerial_image = ground_image.to(device), aerial_image.to(device)

                # Forward pass
                _, attentions = model(ground_image, aerial_image, return_attention=True)
                
                ground_tokens = dino(ground_image)
                aerial_tokens = dino(aerial_image)
                ground_tokens = get_patch_embeddings(dino, ground_image)
                aerial_tokens = get_patch_embeddings(dino, aerial_image)

                # Visualize the final single-head attention layer
                attention = attentions[-1].mean(dim=1)  # average across heads only

                # Calculate the number of patches for ground and aerial images
                patch_size = model.patch_embed.proj.kernel_size[0]
                num_patches_ground = (ground_image.shape[-1] // patch_size) * (ground_image.shape[-2] // patch_size)
                num_patches_aerial = (aerial_image.shape[-1] // patch_size) * (aerial_image.shape[-2] // patch_size)

                # Assuming batch size of 1 for simplicity
                cross_attention_A2G = attention[:, :num_patches_ground, num_patches_ground:].detach()
                cross_attention_G2A = attention[:, num_patches_ground:, :num_patches_ground].detach()

                # Reconstruct the images from the tokens
                reconstructed_aerial = torch.matmul(cross_attention_G2A, ground_tokens)
                reconstructed_ground = torch.matmul(cross_attention_A2G, aerial_tokens)

                if debug:
                    print("attention shape: ", attention.shape)
                    print("attention dtype: ", attention.dtype)
                    print("num_patches_ground: ", num_patches_ground)
                    print("num_patches_aerial: ", num_patches_aerial)
                    print("cross_attention_G2A shape: ", cross_attention_G2A.shape)
                    print("cross_attention_A2G shape: ", cross_attention_A2G.shape)
                    print("reconstructed_ground shape: ", reconstructed_ground.shape)
                    print("reconstructed_aerial shape: ", reconstructed_aerial.shape)

                # Compute MSE loss
                criterion = nn.MSELoss()
                loss_ground = criterion(reconstructed_ground, ground_tokens)
                loss_aerial = criterion(reconstructed_aerial, aerial_tokens)
                
                loss = loss_ground + loss_aerial
                
                running_loss += loss.item()
                
                # Backward pass and optimization
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                # Update progress bar
                pbar.set_postfix({'Loss': running_loss / (pbar.n + 1)})
                pbar.update()

        train_loss = running_loss / len(train_loader)
        
        # Validation
        val_loss = validate(model, dino, val_loader, criterion, device)
        
        # Save the best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            best_model_path = os.path.join(model_path, f'best_model_epoch_{epoch+1}.pth')
            torch.save(model.state_dict(), best_model_path)
        else:
            patience_counter += 1
        
        # Save the last model
        torch.save(model.state_dict(), os.path.join(model_path, f'last_model_epoch_{epoch+1}.pth'))
        
        print(f'Epoch {epoch+1}/{epochs} - Train Loss: {train_loss:.4f} - Val Loss: {val_loss:.4f}')

    print('Training Complete!\nBest Validation Loss:', best_val_loss)

def validate(model, dino, val_loader, criterion, device):
    model.eval()
    val_loss = 0

    with torch.no_grad():
        for ground_image, aerial_image in val_loader:
            ground_image, aerial_image = ground_image.to(device), aerial_image.to(device)
            
            # Forward pass
            output, attentions = model(ground_image, aerial_image, return_attention=True)
            
            # Compute the output of the original DINOv2 model
            dino = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14').to(device)
            ground_tokens = dino(ground_image)
            aerial_tokens = dino(aerial_image)
            
            # Reconstruct the images from the tokens
            patch_size = model.patch_embed.proj.kernel_size[0]
            num_patches_ground = (ground_image.shape[-1] // patch_size) * (ground_image.shape[-2] // patch_size)
            num_patches_aerial = (aerial_image.shape[-1] // patch_size) * (aerial_image.shape[-2] // patch_size)
            
            cross_attention_A2G = attentions[-1][:, :num_patches_ground, num_patches_ground:].detach()
            cross_attention_G2A = attentions[-1][:, num_patches_ground:, :num_patches_ground].detach()
            
            reconstructed_aerial = torch.matmul(cross_attention_G2A, ground_tokens.unsqueeze(0).repeat(cross_attention_G2A.size(0), 1, 1))
            reconstructed_ground = torch.matmul(cross_attention_A2G, aerial_tokens.unsqueeze(0).repeat(cross_attention_A2G.size(0), 1, 1))
            
            # Compute MSE loss
            loss_ground = criterion(reconstructed_ground, ground_tokens.unsqueeze(0).repeat(reconstructed_ground.size(0), 1, 1))
            loss_aerial = criterion(reconstructed_aerial, aerial_tokens.unsqueeze(0).repeat(reconstructed_aerial.size(0), 1, 1))
            loss = loss_ground + loss_aerial
            
            val_loss += loss.item()

    val_loss /= len(val_loader)
    return val_loss

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train a model.')
    parser.add_argument('--save_path', '-p', type=str, default='untitled', help='Path to save the model and results')
    parser.add_argument('--epochs', '-e', type=int, default=1, help='Number of epochs to train')
    args = parser.parse_args()

    # Constants
    image_size = 224
    aerial_scaling = 2
    batch_size = 32
    shuffle = True

    # Select device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Initialize the model
    repo_name="facebookresearch/dinov2"
    model_name="dinov2_vitb14"
    model = CroDINO(repo_name, model_name, pretrained=True).to(device)
    dino = torch.hub.load(repo_name, model_name).to(device)


    # Optimizer
    learning_rate = 1e-5
    weight_decay = 1e-5
    optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

    # Loss function
    criterion = nn.MSELoss()

    # Transformations
    transform_ground = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.CenterCrop((image_size, image_size)),
        transforms.ToTensor()
    ])

    transform_aerial = transforms.Compose([
        transforms.CenterCrop((image_size // aerial_scaling, image_size // aerial_scaling)),
        transforms.ToTensor()
    ])

    # Enable loading truncated images
    ImageFile.LOAD_TRUNCATED_IMAGES = True

    # Sample paired images
    dataset_path = '/home/lrusso/cvusa'
    train_filenames, val_filenames = sample_paired_images(dataset_path, sample_percentage=0.2, split_ratio=0.8, groundtype='cutouts')

    # Instantiate the dataset and dataloader
    train_dataset = PairedImagesDataset(train_filenames, transform_aerial=transform_aerial, transform_ground=transform_ground)
    val_dataset = PairedImagesDataset(val_filenames, transform_aerial=transform_aerial, transform_ground=transform_ground)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=8)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=8)

    # Train the model
    train(model, dino, train_loader, val_loader, device, criterion, optimizer, epochs=args.epochs, save_path=args.save_path)

```


(crossview) lrusso@Chronos:~/projects/CroDINO$ python train.py -p test
Using device: cuda
Using cache found in /home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main
/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/layers/swiglu_ffn.py:51: UserWarning: xFormers is not available (SwiGLU)
  warnings.warn("xFormers is not available (SwiGLU)")
/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/layers/attention.py:33: UserWarning: xFormers is not available (Attention)
  warnings.warn("xFormers is not available (Attention)")
/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/layers/block.py:40: UserWarning: xFormers is not available (Block)
  warnings.warn("xFormers is not available (Block)")
/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/cuda/__init__.py:628: UserWarning: Can't initialize NVML
  warnings.warn("Can't initialize NVML")
Using cache found in /home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main
Epoch 1/1:   0%|                                                                            | 0/5148 [00:01<?, ?batch/s]
Traceback (most recent call last):
  File "/home/lrusso/projects/CroDINO/train.py", line 214, in <module>
    train(model, dino, train_loader, val_loader, device, criterion, optimizer, epochs=args.epochs, save_path=args.save_path)
  File "/home/lrusso/projects/CroDINO/train.py", line 35, in train
    ground_tokens = dino(ground_image)
                    ^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/models/vision_transformer.py", line 325, in forward
    ret = self.forward_features(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/models/vision_transformer.py", line 261, in forward_features
    x = blk(x)
        ^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/layers/block.py", line 254, in forward
    return super().forward(x_or_x_list)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/layers/block.py", line 113, in forward
    x = x + ffn_residual_func(x)
            ^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/layers/block.py", line 94, in ffn_residual_func
    return self.ls2(self.mlp(self.norm2(x)))
                    ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/layers/mlp.py", line 36, in forward
    x = self.act(x)
        ^^^^^^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/nn/modules/activation.py", line 682, in forward
    return F.gelu(input, approximate=self.approximate)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: NVML_SUCCESS == DriverAPI::get()->nvmlInit_v2_() INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":799, please report a bug to PyTorch.





## Note

When I run the training with batch size = 1, the issue disappears. With any number greater than 1 I have this error.

### Versions

Collecting environment information...
/home/lrusso/miniconda3/envs/crossview/lib/python3.12/site-packages/torch/cuda/__init__.py:628: UserWarning: Can't initialize NVML
  warnings.warn("Can't initialize NVML")
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.4 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.12.2 | packaged by Anaconda, Inc. | (main, Feb 27 2024, 17:35:02) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.5.0-35-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: Could not collect
Nvidia driver version: Could not collect
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:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             32
On-line CPU(s) list:                0-31
Vendor ID:                          AuthenticAMD
Model name:                         AMD Ryzen 9 7950X3D 16-Core Processor
CPU family:                         25
Model:                              97
Thread(s) per core:                 2
Core(s) per socket:                 16
Socket(s):                          1
Stepping:                           2
CPU max MHz:                        5759,0000
CPU min MHz:                        400,0000
BogoMIPS:                           8399.73
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid overflow_recov succor smca fsrm flush_l1d
Virtualization:                     AMD-V
L1d cache:                          512 KiB (16 instances)
L1i cache:                          512 KiB (16 instances)
L2 cache:                           16 MiB (16 instances)
L3 cache:                           128 MiB (2 instances)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-31
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: Mitigation; Safe RET
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; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.2.1
[pip3] torchvision==0.17.1
[conda] blas                      1.0                         mkl
[conda] mkl                       2023.1.0         h213fc3f_46344
[conda] mkl-service               2.4.0           py312h5eee18b_1
[conda] mkl_fft                   1.3.8           py312h5eee18b_0
[conda] mkl_random                1.2.4           py312hdb19cb5_0
[conda] numpy                     1.26.4          py312hc5e2394_0
[conda] numpy-base                1.26.4          py312h0da6c21_0
[conda] torch                     2.2.1                    pypi_0    pypi
[conda] torchvision               0.17.1                   pypi_0    pypi

cc @ptrblck @msaroufim