pytorch - 💡(How to fix) Fix weakrefs don't when tensor swap is used

"""Minimal reproducer for why group offloading cannot use torch.utils.swap_tensors()
on regular parameters.

A live torch.compile'd module keeps weakrefs on its parameters (Dynamo guards).
swap_tensors() refuses to operate on tensors that have weakrefs. pytorch#180566 only
clears these when the compile context is torn down -- group offloading runs forwards
through a still-live compiled artifact, so the weakrefs are present at onload time.

Usage:  python repro_swap_weakref.py {swap|data}
Needs:  CUDA + a torch nightly that already contains pytorch#180566.
"""
import sys
import torch

from diffusers import FluxTransformer2DModel
import diffusers.hooks.group_offloading as go

mode = sys.argv[1] if len(sys.argv) > 1 else "swap"

def onload_swap_tensors(target, source):
    # what we'd like to use everywhere (it's required for TorchAO subclass tensors)
    if isinstance(target, torch.nn.Parameter) and not isinstance(source, torch.nn.Parameter):
        source = torch.nn.Parameter(source, requires_grad=target.requires_grad)
    torch.utils.swap_tensors(target, source)

def onload_data_assign(target, source):
    target.data = source  # what diffusers actually uses for regular tensors

go._swap_tensor_data = onload_swap_tensors if mode == "swap" else onload_data_assign

init = dict(patch_size=1, in_channels=4, num_layers=1, num_single_layers=1,
            attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32,
            pooled_projection_dim=32, axes_dims_rope=[4, 4, 8])
inp = dict(
    hidden_states=torch.randn(1, 16, 4, device="cuda"),
    encoder_hidden_states=torch.randn(1, 48, 32, device="cuda"),
    pooled_projections=torch.randn(1, 32, device="cuda"),
    img_ids=torch.randn(16, 3, device="cuda"),
    txt_ids=torch.randn(48, 3, device="cuda"),
    timestep=torch.tensor([1.0], device="cuda"),
)

model = FluxTransformer2DModel(**init).eval()
model.enable_group_offload(onload_device=torch.device("cuda"), offload_device="cpu",
                           offload_type="block_level", num_blocks_per_group=1,
                           use_stream=True, non_blocking=True)
model.compile()

print(f"torch={torch.__version__}  onload={mode}")
with torch.no_grad():
    model(**inp)   # forward 1: compiles blocks -> Dynamo leaves weakrefs on their params
    model(**inp)   # forward 2: onload swaps the now-weakref'd params
print("RESULT: success")

---

"""Minimal reproducer for why group offloading cannot use torch.utils.swap_tensors()
on regular parameters.

THE ISSUE, distilled: torch.utils.swap_tensors(p, ...) raises the instant `p` has *any*
weakref -- it's an unconditional check at the top of the function. That's the whole bug;
it needs neither torch.compile, nor hooks, nor diffusers to reproduce.

Why this is faithful and not a cheat: the weakref that bites in production is created by
Dynamo when it traces a module and wraps each parameter
(torch/_dynamo/variables/builder.py:wrap_tensor), and it is exactly a
torch.utils.weak.TensorWeakRef -- the same type used below.

Why torch.compile isn't in this script: torch.compile attaches that weakref *during*
tracing and cleans it up once the forward returns, so there is no post-forward window in
which to catch it (compile-alone leaves zero persistent weakrefs). The only way the
weakref and a swap coexist is to swap *mid-forward* -- which is precisely what a group
offloading onload hook does. So hooks/diffusers are needed only to recreate that timing
window, not to cause the failure. This script reproduces the failure directly instead.

Run:  python pt_minimal.py   (no CUDA, no diffusers, deterministic)
"""
import weakref

import torch
from torch.utils.weak import TensorWeakRef

p = torch.nn.Parameter(torch.zeros(4, 4))

_ = TensorWeakRef(p)  # exactly what Dynamo's wrap_tensor leaves on a traced parameter

print("weakrefs on param:", len(weakref.getweakrefs(p)))
try:
    torch.utils.swap_tensors(p, torch.nn.Parameter(torch.ones(4, 4)))
    print("swap_tensors: OK")
except RuntimeError as e:
    print("swap_tensors:", str(e).splitlines()[0])

# the assignment diffusers actually uses has no such restriction:
p.data = torch.ones(4, 4)
print("target.data = source: OK ->", bool((p == 1).all()))

---

W0609 15:05:11.891000 68436 torch/utils/_pytree.py:630] <enum 'KernelPreference'> is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.
Unable to import `torchao` Tensor objects. This may affect loading checkpoints serialized with `torchao`
W0609 15:05:12.363000 68436 torch/utils/_pytree.py:630] <enum 'ScaleCalculationMode'> is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.
torch=2.13.0.dev20260603+cu130  onload=swap
Traceback (most recent call last):
  File "/home/sayakpaul/diffusers/repro_swap_weakref.py", line 52, in <module>
    model(**inp)   # forward 2: onload swaps the now-weakref'd params
    ^^^^^^^^^^^^
  File "/home/sayakpaul/.dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sayakpaul/.dev/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1166, in compile_wrapper
    result = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/sayakpaul/.dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sayakpaul/diffusers/src/diffusers/hooks/hooks.py", line 188, in new_forward
    args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sayakpaul/diffusers/src/diffusers/hooks/group_offloading.py", line 401, in pre_forward
    self.group.onload_()
  File "/home/sayakpaul/diffusers/src/diffusers/hooks/group_offloading.py", line 416, in torch_dynamo_resume_in_pre_forward_at_401
    self.next_group.onload_()
  File "/home/sayakpaul/.dev/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1446, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/sayakpaul/diffusers/src/diffusers/hooks/group_offloading.py", line 359, in onload_
    self._onload_from_memory()
  File "/home/sayakpaul/diffusers/src/diffusers/hooks/group_offloading.py", line 305, in _onload_from_memory
    self._process_tensors_from_modules(pinned_memory, default_stream=default_stream)
  File "/home/sayakpaul/diffusers/src/diffusers/hooks/group_offloading.py", line 242, in _process_tensors_from_modules
    self._transfer_tensor_to_device(param, source, default_stream)
  File "/home/sayakpaul/diffusers/src/diffusers/hooks/group_offloading.py", line 231, in _transfer_tensor_to_device
    _swap_tensor_data(tensor, moved)
  File "/home/sayakpaul/diffusers/repro_swap_weakref.py", line 24, in onload_swap_tensors
    torch.utils.swap_tensors(target, source)
  File "/home/sayakpaul/.dev/lib/python3.12/site-packages/torch/utils/__init__.py", line 45, in swap_tensors
    raise RuntimeError("Cannot swap t1 because it has weakref associated with it")
RuntimeError: Cannot swap t1 because it has weakref associated with it

---

Collecting environment information...
PyTorch version: 2.13.0.dev20260603+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (aarch64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
Clang version: Could not collect
CMake version: version 3.28.3
Libc version: glibc-2.39

Python version: 3.12.3 (main, Mar 23 2026, 19:04:32) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.17.0-1021-nvidia-aarch64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 13.0.88
CUDA_MODULE_LOADING set to: 
GPU models and configuration: GPU 0: NVIDIA GB10
Nvidia driver version: 580.159.03
cuDNN version: Could not collect
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  20
On-line CPU(s) list:                     0-19
Vendor ID:                               ARM
Model name:                              Cortex-X925
Model:                                   1
Thread(s) per core:                      1
Core(s) per socket:                      10
Socket(s):                               1
Stepping:                                r0p1
Frequency boost:                         disabled
CPU(s) scaling MHz:                      100%
CPU max MHz:                             3900.0000
CPU min MHz:                             1378.0000
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh bti ecv afp wfxt
Model name:                              Cortex-A725
Model:                                   1
Thread(s) per core:                      1
Core(s) per socket:                      10
Socket(s):                               1
Stepping:                                r0p1
CPU(s) scaling MHz:                      100%
CPU max MHz:                             2808.0000
CPU min MHz:                             338.0000
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh bti ecv afp wfxt
L1d cache:                               1.3 MiB (20 instances)
L1i cache:                               1.3 MiB (20 instances)
L2 cache:                                25 MiB (20 instances)
L3 cache:                                24 MiB (2 instances)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-19
Vulnerability Gather data sampling:      Not affected
Vulnerability Ghostwrite:                Not affected
Vulnerability Indirect target selection: 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 Old microcode:             Not affected
Vulnerability Reg file data sampling:    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; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect