Repository: pytorch
Issue: Pipeline parallel doesn't seem to work with mixed precision
Labels: oncall: distributed

### 🐛 Describe the bug

When running a model with pipeline parallelism, automatic mixed precision doesn't seem to have any effect. I've pasted a minimal example below - when I set `USE_PIPE` to `False` then the print statement inside the model prints `linear.dtype = torch.bfloat16` (as expected) whereas when `USE_PIPE` is set to `True` then the print statement prints `linear.dtype = torch.float32`.

```python
import os
import tempfile
from typing import Final

import numpy as np
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from numpy.typing import NDArray
from torch import nn
from torch.distributed import rpc
from torch.distributed.pipeline.sync import Pipe
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler

USE_PIPE: Final[bool] = True
PIPE_NUM_GPUS: Final[int] = 4
PIPE_CHUNKS: Final[int] = 4

DIM: Final[int] = 1024
NUM_LAYERS: Final[int] = 8
BATCH_SIZE: Final[int] = 8


class MyDataset(Dataset):
    def __len__(self) -> int:
        return 999

    def __getitem__(self, index: int) -> NDArray:
        return np.float32(np.random.rand(DIM))


def collate_fn(batch: list[NDArray]) -> torch.Tensor:
    return torch.stack(batch)


class Layer(nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.linear = nn.Linear(DIM, DIM)
        self.relu = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        linear = self.linear(x)

        print(f"{linear.dtype = }")

        return self.relu(x)


def build_model(rank: int) -> nn.Sequential:
    layers: list[nn.Module] = []

    num_gpus = PIPE_NUM_GPUS if USE_PIPE else 1
    partition_len = ((NUM_LAYERS - 1) // num_gpus) + 1

    for i in range(NUM_LAYERS):
        device = torch.device((num_gpus * rank) + (i // partition_len))

        layers.append(Layer().to(device))

    return nn.Sequential(*layers)


def run_worker(rank: int, world_size: int) -> None:
    if USE_PIPE:
        tmpfile = tempfile.NamedTemporaryFile()

        rpc.init_rpc(
            name="worker",
            # We need to initialize the RPC framework with only a single worker since we’re using a
            # single process to drive multiple GPUs.
            rank=0,
            world_size=1,
            rpc_backend_options=rpc.TensorPipeRpcBackendOptions(
                init_method="file://{}".format(tmpfile.name),
                # # Specifying _transports and _channels is a workaround and we no longer
                # # will have to specify _transports and _channels for PyTorch
                # # versions >= 1.8.1
                # _transports=["ibv", "uv"],
                # _channels=["cuda_ipc", "cuda_basic"],
            )
        )

    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)

    model = build_model(rank)

    if USE_PIPE:
        model = Pipe(model, chunks=PIPE_CHUNKS, checkpoint="never")

    model = DDP(model)
    optimizer = torch.optim.Adam(model.parameters())

    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(
        dataset=dataset,
        batch_size=BATCH_SIZE,
        sampler=sampler,
        pin_memory=True,
    )

    model.train()

    for batch in dataloader:
        optimizer.zero_grad()

        if USE_PIPE:
            batch = batch.to(rank * PIPE_NUM_GPUS)
        else:
            batch = batch.to(rank)

        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
            model(batch)


if __name__ == "__main__":
    if USE_PIPE:
        world_size = int(8 / PIPE_NUM_GPUS)
    else:
        world_size = 8

    mp.spawn(run_worker, args=(world_size,), nprocs=world_size)
```

### Versions

PyTorch version: 2.1.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: (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.6 (main, Dec  6 2023, 13:50:43) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.2.0-1018-azure-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.3.103
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 545.23.08
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):                             96
On-line CPU(s) list:                0-95
Vendor ID:                          AuthenticAMD
Model name:                         AMD EPYC 7V12 64-Core Processor
CPU family:                         23
Model:                              49
Thread(s) per core:                 1
Core(s) per socket:                 48
Socket(s):                          2
Stepping:                           0
BogoMIPS:                           4890.89
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 tsc_reliable nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core ssbd vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru arat umip rdpid
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          3 MiB (96 instances)
L1i cache:                          3 MiB (96 instances)
L2 cache:                           48 MiB (96 instances)
L3 cache:                           384 MiB (24 instances)
NUMA node(s):                       4
NUMA node0 CPU(s):                  0-23
NUMA node1 CPU(s):                  24-47
NUMA node2 CPU(s):                  48-71
NUMA node3 CPU(s):                  72-95
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; untrained return thunk; SMT disabled
Vulnerability Spec rstack overflow: Mitigation; safe RET, no microcode
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; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.2
[pip3] torch==2.1.1
[pip3] triton==2.1.0
[conda] Could not collect

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225