Closed as duplicate of#154241
Closed as duplicate of#154241
Description
🐛 Describe the bug
Hi,
I am experimenting with parameter sharing across different networks. The idea is to have different networks share the same parameters using the .state_dict() + .load_dict() idiom:
net1 = create_network()
net2 = create_network()
net2.load_dict(net1.state_dict(), assign=True)
# net1 and net2 are different objects sharing the memory of their parameters
# updates are shared between the two networksThis is useful in my reinforcement learning application using torchrl and tensordict. When different TensorDictModules share the same underlying module but have different keys: we initialise two identically structured TensorDictModule with the different keys I need, and then share parameters using state dicts.
This works well for modules containing nested fully connected layers.
However, problems start when a nn.LSTM is involved in the mix. After debugging, I discovered that the culprit is the RNNBase._flat_weight logic, which is used to guarantee that the rnn's cudnn kernel is called with flattened parameters (ie, all the weights and biases of the module are contiguous in memory). The problem here is that load_state_dict changes the cached references to the flat weights, and this triggers a flatten_parameters() call, which re-initialises the parameters, breaking the parameter sharing I tried to obtain with load_state_dict.
A reproducing script:
from collections.abc import Sequence
import logging
from typing import Any
import weakref
import torch
from torch import nn
from torch.nn import RNNBase
log = logging.Logger(__file__, level=0)
def all_parameters_shared(policy, policy_, /, *, return_different=False) -> bool | list[str]:
# Validate identical structure and parameter sharing across policies
# between the cross-play and self-play. Guarantees that the `cross_play_training` policy is on the correct device
different_params_keys = [(k, k_)
for (k,p), (k_,p_) in zip(
policy.named_parameters(),
policy_.named_parameters(),
strict=True,
) if p.data_ptr() != p_.data_ptr()]
if not return_different:
if len(different_params_keys) > 0:
return False
return True
else:
return different_params_keys
def all_parameters_equal(policy, policy_, /) -> bool:
# Validate identical structure and parameter sharing across policies
# between the cross-play and self-play. Guarantees that the `cross_play_training` policy is on the correct device
different_params_keys = [(k, k_)
for (k,p), (k_,p_) in zip(
policy.named_parameters(),
policy_.named_parameters(),
strict=True,
) if torch.any(p != p_)]
if len(different_params_keys) > 0:
return False
return True
BATCH_SIZE = 20
SEQ_LEN = 4
INPUT_SIZE = 10
HIDDEN_SIZE = 7
def build_network() -> nn.LSTM:
return nn.LSTM(
input_size=INPUT_SIZE,
hidden_size=HIDDEN_SIZE,
batch_first=True
)
if __name__ == "__main__":
device = "cuda"
x = torch.rand((BATCH_SIZE, SEQ_LEN, INPUT_SIZE)).to(device)
recurrent_state_h = torch.rand((1, BATCH_SIZE, HIDDEN_SIZE)).to(device)
recurrent_state_c = torch.rand((1, BATCH_SIZE, HIDDEN_SIZE)).to(device)
hidden = (recurrent_state_c, recurrent_state_h)
net1 = build_network().to(device)
out, hidden = net1(x, hidden)
assert out.shape == (BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE)
assert hidden[0].shape == (1, BATCH_SIZE, HIDDEN_SIZE)
assert hidden[1].shape == (1, BATCH_SIZE, HIDDEN_SIZE)
net2 = build_network().to(device)
out, hidden = net2(x, hidden)
# clearly, at the start, they do not share parameters
assert not all_parameters_shared(net1, net2)
assert not all_parameters_equal(net1, net2)
# try to share parameters
net2.load_state_dict(net1.state_dict(), assign=True)
# everything seems to work...
assert all_parameters_shared(net1, net2)
assert all_parameters_equal(net1, net2)
## THE PROBLEM
# the lstm actually thinks its weights are not contiguous
assert net2._weights_have_changed()
# a reallocation is triggered with the next `.forward()`
out, hidden = net2(x, hidden)
if device.startswith("cuda"):
# Compacting only happens when using a cuda device
# Because there we have the use of the cuDNN kernel
assert not all_parameters_shared(net1, net2)
else:
assert all_parameters_shared(net1, net2)
assert all_parameters_equal(net1, net2)
## A STRANGE DIFFERENCE
# The same DOES NOT HAPPEN when using a `.to(same device)` which calls
# something similar to`_apply(lambda param: param.to())` recursively on modules
net2_ = net2.to(device) # `.to()` to the same device is a no-op
# However RNNBase._apply() always invokes a `flatten_parameters()`
# this calls the `torch._cudnn_rnn_flatten_weight` which for an unknown reason
# does not reallocate the buffer containing the params.
# I checked the C++ code without success to find any logic justifying this behavior.
assert all_parameters_shared(net2, net2_)
# This does not happen after the `load_state_dict` however, ie, parameters are changed
#
# QUESTION?: Is there is some buffer-tracking logic in the C++ code that makes
# sure that already flattened parameters are not re-instantiated?cc @vmoens
Versions
PyTorch version: 2.7.0+cu126
Is debug build: False
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A
OS: Ubuntu 22.04.5 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.24.1
Libc version: glibc-2.35
Python version: 3.11.12 (main, Apr 9 2025, 04:04:00) [Clang 20.1.0 ] (64-bit runtime)
Python platform: Linux-5.10.236-228.935.amzn2.x86_64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla V100-SXM2-16GB
Nvidia driver version: 550.163.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.2
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: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 32
On-line CPU(s) list: 0-31
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) CPU E5-2686 v4 @ 2.30GHz
CPU family: 6
Model: 79
Thread(s) per core: 2
Core(s) per socket: 16
Socket(s): 1
Stepping: 1
CPU max MHz: 3000.0000
CPU min MHz: 1200.0000
BogoMIPS: 4600.02
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 pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single pti fsgsbase bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx xsaveopt
Hypervisor vendor: Xen
Virtualization type: full
L1d cache: 512 KiB (16 instances)
L1i cache: 512 KiB (16 instances)
L2 cache: 4 MiB (16 instances)
L3 cache: 45 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-31
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: KVM: Mitigation: VMX unsupported
Vulnerability L1tf: Mitigation; PTE Inversion
Vulnerability Mds: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Meltdown: Mitigation; PTI
Vulnerability Mmio stale data: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Vulnerable
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: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect
Metadata
Metadata
Assignees
Labels
No labels