Repository: pytorch
Issue: RuntimeError: Modules with uninitialized parameters can't be used with `DistributedDataParallel` when using subclass of tensor.
Labels: oncall: distributed,tensor subclass

### 🐛 Describe the bug

DDP init call is failing when using subclass of torch.Tensor, same code works with torch.Tensor.
Command to run the code
python test.py --max-gpus 2 --batch-size 512 --epoch 10

High Level overview of code example.
1) Define subclass of torch.tensor called MyTensor
2) Define subclass of torch.nn.Module called MyLinearLayer which uses tensor of type MyTensor for storing weight.
3) Instantiate a model call Net using MyLinearLayer.
4) Distribute the model Net across given number of GPU     
```ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True,static_graph=True )```
5) During initialization following error is reported.

```
free_gpus=[1, 2]
<MyTensor([1., 2., 3.])>
Running DDP on rank 1 batch_size=512.
Running DDP on rank 0 batch_size=512.
W0909 10:36:55.065000 23054275613888 torch/multiprocessing/spawn.py:146] Terminating process 3930793 via signal SIGTERM
Traceback (most recent call last):
  File "/prj/qct/lasgpu/users/ankushj/builds/osetml/3.py", line 205, in <module>
    main()
  File "/prj/qct/lasgpu/users/ankushj/builds/osetml/3.py", line 199, in main
    mp.spawn(train, args=(world_size, args.batch_size,args.epoch,  args.verbose), nprocs=world_size, join=True)
  File "/prj/qct/lasgpu/users/ankushj/tools/miniconda3/envs/osetml/lib/python3.11/site-packages/torch/multiprocessing/spawn.py", line 282, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/prj/qct/lasgpu/users/ankushj/tools/miniconda3/envs/osetml/lib/python3.11/site-packages/torch/multiprocessing/spawn.py", line 238, in start_processes
    while not context.join():
              ^^^^^^^^^^^^^^
  File "/prj/qct/lasgpu/users/ankushj/tools/miniconda3/envs/osetml/lib/python3.11/site-packages/torch/multiprocessing/spawn.py", line 189, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/prj/qct/lasgpu/users/ankushj/tools/miniconda3/envs/osetml/lib/python3.11/site-packages/torch/multiprocessing/spawn.py", line 76, in _wrap
    fn(i, *args)
  File "/prj/qct/lasgpu/users/ankushj/builds/osetml/3.py", line 126, in train
    ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True,static_graph=False )
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/prj/qct/lasgpu/users/ankushj/tools/miniconda3/envs/osetml/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 784, in __init__
    self._log_and_throw(
  File "/prj/qct/lasgpu/users/ankushj/tools/miniconda3/envs/osetml/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 1127, in _log_and_throw
    raise err_type(err_msg)
RuntimeError: Modules with uninitialized parameters can't be used with `DistributedDataParallel`. Run a dummy forward pass to correctly initialize the modules


``` 

6) Same code works if we replace MyTensor with native tensor. i.e
```
            #weight = MyTensor.ones((in_features, out_features),**factory_kwargs)
            weight = torch.ones((in_features, out_features),**factory_kwargs)
```

```
free_gpus=[1, 2]
<MyTensor([1., 2., 3.])>
Running DDP on rank 0 batch_size=512.
Running DDP on rank 1 batch_size=512.
[rank1]:[W909 10:38:32.530689521 reducer.cpp:1400] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration,  which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())
[rank0]:[W909 10:38:32.583329909 reducer.cpp:1400] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration,  which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())
Training on GPU 0 took 2.22 seconds for 10 epochs for batch_size=512.
Training on GPU 1 took 2.27 seconds for 10 epochs for batch_size=512.
```

7) Underlying issue seems to be is related with call to 
```Parameter(weight,requires_grad=True)```
When the weight is of type torch.tensor it removes the attribute torch.nn.parameter.UninitializedParameter however for subclass MyTensor said attribute persist even after explicit init.

This can be further demonstrated by following unit test case.

```
import torch
from torch.nn import Parameter


def check(tensor):
    print("*"*20)
    print(f'{type(tensor)=}\n{tensor=}')
    param = Parameter(tensor)
    assert not isinstance(param, torch.nn.parameter.UninitializedParameter)
    assert type(param) is torch.nn.parameter.Parameter

# Define a custom tensor subclass
class MyTensor(torch.Tensor):
    def __init__(self, *args, **kwargs):
        super().__init__()

    def custom_method(self):
        return self * 2

size=(2,2)
tensor = torch.randn(size)
# Convert the tensor to a Parameter
check(tensor)

float_tensor = torch.FloatTensor(size).uniform_(-1, 1)
# Convert the FloatTensor to a Parameter
check(float_tensor)

# Create an instance of the custom tensor subclass
my_tensor = MyTensor(torch.randn(size))
check(my_tensor)
```

Output:

```
********************
type(tensor)=<class 'torch.Tensor'>
tensor=tensor([[-1.0664, -1.0035],
        [ 0.5663,  0.8214]])
********************
type(tensor)=<class 'torch.Tensor'>
tensor=tensor([-0.6655, -0.6205])
********************
type(tensor)=<class '__main__.MyTensor'>
tensor=MyTensor([[-1.6115, -0.8384],
          [ 0.0495,  0.1868]])
Traceback (most recent call last):
  File "/prj/qct/lasgpu/users/ankushj/builds/osetml/a1.py", line 35, in <module>
    check(my_tensor)
  File "/prj/qct/lasgpu/users/ankushj/builds/osetml/a1.py", line 9, in check
    assert not isinstance(param, torch.nn.parameter.UninitializedParameter)
AssertionError
```




 
Example code to reproduce the DDP init issue.

```
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.nn.parameter import Parameter, UninitializedParameter
from torch.utils.data import Dataset, DataLoader, DistributedSampler
from torchvision import datasets, transforms
import argparse
import os
import pynvml
import time
import torch
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim


class MyTensor(torch.Tensor):
    def __init__(self, *args,  **kwargs):
        super().__init__()

    @staticmethod
    def __new__(cls, *args,  **kwargs):
        return  super().__new__(cls, *args, **kwargs)

    @classmethod
    def zeros(cls, *args, **kwargs):
        return cls(torch.zeros(*args, **kwargs))

    @classmethod
    def ones(cls, *args, **kwargs):
        return cls(torch.ones(*args, **kwargs))

    @classmethod
    def empty(cls, *args, **kwargs):
        return cls(torch.empty(*args, **kwargs))

    def __repr__(self, *, tensor_contents=None):
            return f"<{super().__repr__()}>"


class MyLinearLayer(torch.nn.Module):
    '''My linear nn.Module'''

    def __init__(self, in_features, out_features,bias: bool = True, device=None, dtype=None):
        '''Creates nn.Parameters for model and configures their initial value'''
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(MyLinearLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        weight = MyTensor.ones((in_features, out_features),**factory_kwargs)
        #weight = torch.ones((in_features, out_features),**factory_kwargs) #Native tensor works
        weight = Parameter(weight,requires_grad=True)
        if not isinstance(weight, torch.nn.parameter.UninitializedParameter):
            print(f'Error : {isinstance(weight, torch.nn.parameter.UninitializedParameter)=} !\n{type(weight)=}\n{weight=}')
        self.weight = weight
        self.reset_parameters()

    def forward(self, input):
        '''Forward pass for simple linear model'''
        y=torch.matmul(input,self.weight)
        return y

    def extra_repr(self):
        #return f'MyLinearLayer : w={self.w} output_fmt=[{self.output_fmt}]'
        buff=''
        for n, p in self.named_parameters():
            buff+=f'MyLinearLayer : {n=}:{p=} '
        return buff

    def reset_parameters(self) -> None:
        self.weight.data.fill_(0)

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = MyLinearLayer(2, 1 )

    def forward(self, x):
        x = self.fc(x)
        return x


def get_free_gpus(threshold=0.20, max_gpus=None):
    free_gpus = []
    device_count = pynvml.nvmlDeviceGetCount()
    for i in range(device_count):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        free_capacity = mem_info.free / mem_info.total
        if free_capacity > threshold:
            free_gpus.append(i)
        if max_gpus and len(free_gpus) >= max_gpus:
            break
    print(f"{free_gpus=}")
    return free_gpus

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12345'
    torch.cuda.set_device(rank)
    torch.distributed.init_process_group(backend='nccl', rank=rank, world_size=world_size)

def cleanup():
    torch.distributed.destroy_process_group()

class SumDataset(Dataset):
    def __init__(self, data_size):
        self.data = [([v[0].item(), v[1].item()], v[0].item() + v[1].item()) for v in zip(torch.rand((data_size,)), torch.rand((data_size,)))]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        inputs, target = self.data[idx]
        return torch.tensor(inputs, dtype=torch.float32), torch.tensor(target, dtype=torch.float32)

def train(rank, world_size, batch_size, epoch, verbose=False):
    print(f"Running DDP on rank {rank} {batch_size=}.")
    setup(rank, world_size)

    # Create model and move it to the corresponding GPU
    model = Net().to(rank)
    if verbose:
        print(f'********** Model Description for {rank=} after dummy pass ********** ')
        for name, param in model.named_parameters():
            print(name, param.data)
            pass

    ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True,static_graph=False )

    # Use a distributed sampler
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

    data_size=2**15
    train_dataset = SumDataset(data_size)
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
    train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, sampler=train_sampler)

    # Loss and optimizer
    criterion = nn.MSELoss().to(rank)
    #optimizer = optim.SGD(ddp_model.parameters(), lr=0.12)
    optimizer = optim.Adam(model.parameters(), lr=0.01)


    # Training loop
    start_time = time.time()
    ddp_model.train()
    count = 0

    while count < epoch:
        count_loss = 0.0

        for data, target in (train_loader):
            data, target = data.to(rank), target.to(rank).view(-1, 1)

            optimizer.zero_grad()
            output = ddp_model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            count_loss += loss.item()
            loss=count_loss / len(train_loader)
        if verbose:
            print(f"Rank {rank}, epoch {count}/{epoch}, Loss: {loss:.4f} ")

        count += 1

    end_time = time.time()
    cleanup()

    print(f"Training on GPU {rank} took {end_time - start_time:.2f} seconds for {epoch} epochs for {batch_size=}.")


def main():
    parser = argparse.ArgumentParser(description="Distributed Training Script")
    parser.add_argument('--batch-size', type=int, default=32, help='Batch size for training')
    parser.add_argument('--epoch', type=int, default=32, help='Epoch size for training')
    parser.add_argument('--max-gpus', type=int, default=1, help='Maximum number of GPUs to use')
    parser.add_argument('--verbose', action='store_true', help='Enable verbose output')

    args = parser.parse_args()
    args.batch_size=max(32,args.batch_size)
    args.epoch=max(2,args.epoch)
    args.max_gpus=max(1,args.max_gpus)

    free_gpus = get_free_gpus(threshold=0.5, max_gpus=args.max_gpus)

    if not free_gpus:
        print("No GPUs with more than 50% free capacity found.")
        return

    os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(map(str, free_gpus))

    world_size = len(free_gpus)

# Example usage
    tensor = MyTensor([1, 2, 3])
    print(tensor)  # tensor([1, 2, 3])

    mp.spawn(train, args=(world_size, args.batch_size,args.epoch,  args.verbose), nprocs=world_size, join=True)

torch._dynamo.config.traceable_tensor_subclasses=set([MyTensor])
if __name__ == "__main__":
# Initialize NVML
    pynvml.nvmlInit()
    main()
# Shutdown NVML
    pynvml.nvmlShutdown()

```

### Versions

```
$ python collect_env.py
Collecting environment information...

GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 525.60.13
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
Byte Order:                      Little Endian
Address sizes:                   43 bits physical, 48 bits virtual
CPU(s):                          256
On-line CPU(s) list:             0-255
Thread(s) per core:              2
Core(s) per socket:              64
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           49
Model name:                      AMD EPYC 7742 64-Core Processor
Stepping:                        0
Frequency boost:                 enabled
CPU MHz:                         1499.983
CPU max MHz:                     2250.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        4500.04
Virtualization:                  AMD-V
L1d cache:                       4 MiB
L1i cache:                       4 MiB
L2 cache:                        64 MiB
L3 cache:                        512 MiB
NUMA node0 CPU(s):               0-63,128-191
NUMA node1 CPU(s):               64-127,192-255
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   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; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
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 nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic 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 sme ssbd mba sev ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] onnx==1.15.0
[pip3] onnxruntime==1.18.1
[pip3] onnxsim==0.4.36
[pip3] torch==2.1.1
[pip3] torchaudio==2.4.0
[pip3] torchvision==0.16.1
[pip3] triton==2.1.0
[conda] blas                      1.0                         mkl
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344
[conda] mkl-service               2.4.0           py311h5eee18b_1
[conda] mkl_fft                   1.3.8           py311h5eee18b_0
[conda] mkl_random                1.2.4           py311hdb19cb5_0
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] numpy-base                1.26.4          py311hf175353_0
[conda] pytorch                   2.4.0           py3.11_cuda12.1_cudnn9.1.0_0    pytorch
[conda] pytorch-cuda              12.1                 ha16c6d3_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     2.1.1                    pypi_0    pypi
[conda] torchaudio                2.4.0               py311_cu121    pytorch
[conda] torchtriton               3.0.0                     py311    pytorch
[conda] torchvision               0.16.1                   pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi
```

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @ezyang @albanD