🐛 Describe the bug

Hello. I found out that buffers in FSDP are not casted to requested dtype, and code breaks. User is forced to cast buffers each time in forward.

Piece of error:

  File "/home/user/regbuf_compile_debug.py", line 44, in forward
    return nn.functional.conv2d(x, self.kernel, groups=self.kernel.shape[0], stride=2, padding=self.padding)
RuntimeError: expected scalar type Half but found Float

Repro script:

import argparse
import os
from contextlib import nullcontext
from typing import Tuple

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
from torch.distributed.fsdp.wrap import ModuleWrapPolicy
from tqdm.auto import tqdm

torch._dynamo.config.inline_inbuilt_nn_modules = False
torch._dynamo.config.optimize_ddp = False


def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


class NonLearnableConv(nn.Module):
    def __init__(self, kernel: Tuple[int], in_channels: int):
        super().__init__()

        self.padding = (len(kernel) - 1) // 2

        kernel = torch.tensor(kernel, dtype=torch.float32)
        kernel = kernel / kernel.sum()
        kernel = kernel.outer(kernel)[None, None].repeat(in_channels, 1, 1, 1)

        self.register_buffer("kernel", kernel)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        print(x.dtype, self.kernel.dtype)
        return nn.functional.conv2d(x, self.kernel, groups=self.kernel.shape[0], stride=2, padding=self.padding)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--grad_accum_steps", type=int, default=1)
    parser.add_argument("--num_iterations", type=int, default=200)
    parser.add_argument("--use_fsdp", action="store_true")
    parser.add_argument("--use_compile", action="store_true")
    args = parser.parse_args()
    return args


def main(rank, world_size, args):
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    device = torch.device(f"cuda:{rank}")
    dtype = torch.float16

    model = nn.Sequential(
        nn.Sequential(nn.Conv2d(3, 64, 3, padding=1)),
        nn.Sequential(NonLearnableConv((1, 2, 2, 1), 64)),
        nn.Sequential(nn.Conv2d(64, 3, 3, padding=1)),
        nn.Sequential(NonLearnableConv((1, 2, 2, 1), 3)),
    ).to(device)

    if args.use_fsdp:
        model = FSDP(
            module=model,
            device_id=rank,
            use_orig_params=args.use_compile,
            sharding_strategy=ShardingStrategy.HYBRID_SHARD,
            forward_prefetch=True,
            limit_all_gathers=True,
            auto_wrap_policy=ModuleWrapPolicy({nn.Sequential}),
            mixed_precision=MixedPrecision(
                param_dtype=dtype,
                buffer_dtype=dtype,
                reduce_dtype=dtype,
            ),
        )
        loss_amp_context = torch.amp.autocast("cuda", dtype=dtype, enabled=True)
        model_amp_context = nullcontext()
        scaler = ShardedGradScaler(enabled=dtype is torch.float16)
    else:
        loss_amp_context = torch.amp.autocast("cuda", dtype=dtype, enabled=True)
        model_amp_context = loss_amp_context
        scaler = torch.amp.GradScaler("cuda", enabled=dtype is torch.float16)

    if args.use_compile:
        print("Trying compile.")
        model.compile(mode="default", dynamic=False)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-5, betas=(0.9, 0.98))

    iterator = range(args.num_iterations)
    if rank == 0:
        iterator = tqdm(iterator, total=args.num_iterations, miniters=10)

    for _ in iterator:
        for _ in range(args.grad_accum_steps):

            x = torch.randn(args.batch_size, 3, 128, 128, device=device)
            with model_amp_context:
                out = model(x)
            with loss_amp_context:
                loss = out.mean() / args.grad_accum_steps

            loss_test = loss.clone()  # Ensure local loss is not changed by allreduce
            torch.distributed.all_reduce(loss_test)  # Check if any gpu has NaN loss
            if torch.isnan(loss_test):
                raise ValueError("NaN loss.")

            scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
    cleanup()


if __name__ == "__main__":
    args = parse_args()
    world_size = torch.cuda.device_count()
    torch.multiprocessing.freeze_support()
    if world_size == 1:
        main(0, world_size, args)
    else:
        torch.multiprocessing.spawn(fn=main, args=(world_size, args), nprocs=world_size, join=True)

Successful launch without FSDP:
python regbuf_compile_debug.py

To break:
python regbuf_compile_debug.py --use_fsdp

This error propagates in real setup on UNet model (buffers are casted to input time in forward) with attention blocks when using FSDP, compile and gradient accumulation, dtype errors appear + sometimes _dynamo does not even find registered buffer and fails with AttributeError.

Versions

PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
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: 14.0.0-1ubuntu1.1
CMake version: version 3.30.0
Libc version: glibc-2.35

Python version: 3.10.13 | packaged by conda-forge | (main, Dec 23 2023, 15:36:39) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.4.210-39.1-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A800-SXM4-80GB
GPU 1: NVIDIA A800-SXM4-80GB
GPU 2: NVIDIA A800-SXM4-80GB
GPU 3: NVIDIA A800-SXM4-80GB
GPU 4: NVIDIA A800-SXM4-80GB
GPU 5: NVIDIA A800-SXM4-80GB
GPU 6: NVIDIA A800-SXM4-80GB
GPU 7: NVIDIA A800-SXM4-80GB

Nvidia driver version: 550.54.14
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): 120
On-line CPU(s) list: 0-119
Vendor ID: AuthenticAMD
Model name: AMD EPYC 7662 64-Core Processor
CPU family: 23
Model: 49
Thread(s) per core: 1
Core(s) per socket: 1
Socket(s): 120
Stepping: 0
BogoMIPS: 3992.37
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr wbnoinvd arat npt nrip_save umip rdpid arch_capabilities
Virtualization: AMD-V
L1d cache: 7.5 MiB (120 instances)
L1i cache: 7.5 MiB (120 instances)
L2 cache: 60 MiB (120 instances)
L3 cache: 1.9 GiB (120 instances)
NUMA node(s): 4
NUMA node0 CPU(s): 0-29
NUMA node1 CPU(s): 30-59
NUMA node2 CPU(s): 60-89
NUMA node3 CPU(s): 90-119
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 disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected

Versions of relevant libraries:
[pip3] flake8==5.0.4
[pip3] lovely-numpy==0.2.13
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] open-clip-torch==2.24.0
[pip3] pytorch-warmup==0.1.1
[pip3] torch==2.5.1
[pip3] torch-fidelity==0.3.0
[pip3] torch-model-archiver==0.11.1
[pip3] torch-tb-profiler==0.4.3
[pip3] torch-workflow-archiver==0.2.14
[pip3] torchaudio==2.5.1
[pip3] torchdata==0.7.1
[pip3] torchmetrics==1.4.0.post0
[pip3] torchsde==0.2.6
[pip3] torchserve==0.11.1
[pip3] torchvision==0.20.1
[pip3] triton==3.1.0
[conda] lovely-numpy 0.2.13 pypi_0 pypi
[conda] numpy 1.26.4 pypi_0 pypi
[conda] open-clip-torch 2.24.0 pypi_0 pypi
[conda] torch 2.5.1 pypi_0 pypi
[conda] torch-fidelity 0.3.0 pypi_0 pypi
[conda] torch-model-archiver 0.11.1 pypi_0 pypi
[conda] torch-tb-profiler 0.4.3 pypi_0 pypi
[conda] torch-workflow-archiver 0.2.14 pypi_0 pypi
[conda] torchaudio 2.5.1 pypi_0 pypi
[conda] torchdata 0.7.1 pypi_0 pypi
[conda] torchmetrics 1.4.0.post0 pypi_0 pypi
[conda] torchsde 0.2.6 pypi_0 pypi
[conda] torchserve 0.11.1 pypi_0 pypi
[conda] torchvision 0.20.1 pypi_0 pypi
[conda] triton 3.1.0 pypi_0 pypi

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @zhaojuanmao @mrshenli @rohan-varma @chauhang