执行代码
python
mport torch
import torch.cuda
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 参数设置
B = 64 # batch size
L = 32 # sequence length
C = 512 # embedding dimension
H = 8 # number of heads
D = C // H # head dimension
# 创建随机张量
q = torch.randn(B, H, L, D).to(device)
k = torch.randn(B, H, D, L).to(device)
v = torch.randn(B, H, L, D).to(device)
x = torch.randn(B, L, C).to(device)
# 记录当前显存使用
def fa( x):
print(f"Initial Memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
prev_memory = torch.cuda.memory_allocated()
# 执行矩阵乘法
attn = (q @ k) * 0.125 # 假设 self.scale = 0.125
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
prev_memory = current_memory # 更新 prev_memory
if True:
# 执行最终的矩阵乘法和重新整形
x = (attn @ v).transpose(1, 2).reshape(B, L, C)
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After final matmul and reshape: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
fa(x)
current_memory = torch.cuda.memory_allocated()
print(f"final : {current_memory / 1024**2:.2f} MB")
结果为
bash
Initial Memory: 16.00 MB
After matmul: 18.00 MB, Change: 2.00 MB
After final matmul and reshape: 22.00 MB, Change: 4.00 MB
final : 16.00 MB
但是执行代码
python
import torch
import torch.cuda
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 参数设置
B = 64 # batch size
L = 32 # sequence length
C = 512 # embedding dimension
H = 8 # number of heads
D = C // H # head dimension
# 创建随机张量
q = torch.randn(B, H, L, D).to(device)
k = torch.randn(B, H, D, L).to(device)
v = torch.randn(B, H, L, D).to(device)
x = torch.randn(B, L, C).to(device)
print(f"Initial Memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
prev_memory = torch.cuda.memory_allocated()
# 执行矩阵乘法
attn = (q @ k) * 0.125 # 假设 self.scale = 0.125
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
prev_memory = current_memory # 更新 prev_memory
if True:
# 执行最终的矩阵乘法和重新整形
x = (attn @ v).transpose(1, 2).reshape(B, L, C)
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After final matmul and reshape: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
结果为
bash
Initial Memory: 16.00 MB
After matmul: 18.00 MB, Change: 2.00 MB
After final matmul and reshape: 18.00 MB, Change: 0.00 MB
主要涉及 PyTorch 的显存管理机制 和 Python 的作用域规则。
1. 作用域与变量生命周期:
在 Python 中,变量的生命周期受到作用域的影响:
- 不在函数中 时:
- 当执行
x = (attn @ v).transpose(1, 2).reshape(B, L, C)
时,旧的x
会被立即覆盖,旧的x
的引用计数变为 0,显存会被立即释放或加入缓存。
- 当执行
- 在函数中 时:
- 局部变量(如
attn
,v
等)在函数执行完之前不会被释放,即使x
被重新赋值,中间张量(如attn @ v
)依然在显存中占用空间。 - 这些局部变量直到函数返回后才会被 Python 的垃圾回收机制回收,导致显存未及时释放,从而造成额外的显存占用。
- 局部变量(如
2. PyTorch 的显存缓存机制:
PyTorch 使用 CUDA 显存缓存机制来优化显存分配,但有以下行为特征:
- 不在函数中时 :
- PyTorch 可能更积极地释放未引用的张量。
- 在函数中时 :
- 中间结果的显存使用可能会被缓存,直到函数执行完毕,导致显存使用看起来增加了。
3. 临时张量未释放:
在函数中,PyTorch 可能会生成一些 临时张量,这些临时张量通常在 PyTorch 后台管理,直到函数执行完毕后才会释放。因此:
- 不在函数中时,这些张量更早地被回收。
- 在函数中时,临时张量的显存延迟释放,导致显存增加。
5 函数中内存管理
python
def fa( x):
print(f"Initial Memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
prev_memory = torch.cuda.memory_allocated()
x + 1
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
prev_memory = current_memory
# 执行矩阵乘法
attn = (q @ k) * 0.125 # 假设 self.scale = 0.125
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
prev_memory = current_memory # 更新 prev_memory
if True:
# 执行最终的矩阵乘法和重新整形
x = (attn @ v).transpose(1, 2).reshape(B, L, C)
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After final matmul and reshape: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
fa(x)
current_memory = torch.cuda.memory_allocated()
print(f"final : {current_memory / 1024**2:.2f} MB")
结果
bash
Initial Memory: 16.00 MB
After matmul: 16.00 MB, Change: 0.00 MB
After matmul: 18.00 MB, Change: 2.00 MB
After final matmul and reshape: 22.00 MB, Change: 4.00 MB
final : 16.00 MB
python
def fa( x):
print(f"Initial Memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
prev_memory = torch.cuda.memory_allocated()
x = x + 1
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
prev_memory = current_memory
# 执行矩阵乘法
attn = (q @ k) * 0.125 # 假设 self.scale = 0.125
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
prev_memory = current_memory # 更新 prev_memory
if True:
# 执行最终的矩阵乘法和重新整形
x = (attn @ v).transpose(1, 2).reshape(B, L, C)
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After final matmul and reshape: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
fa(x)
current_memory = torch.cuda.memory_allocated()
print(f"final : {current_memory / 1024**2:.2f} MB")
bash
Initial Memory: 16.00 MB
After matmul: 20.00 MB, Change: 4.00 MB
After matmul: 22.00 MB, Change: 2.00 MB
After final matmul and reshape: 22.00 MB, Change: 0.00 MB
final : 16.00 MB
1. 操作与赋值的差异:
首先,理解 x = x + 1
和 x + 1
在 PyTorch 中的区别很重要。
x + 1
:不会增加显存
x + 1
是一个普通的张量运算,它会 创建一个新的张量 作为结果,但不会修改原张量x
。- 然而,PyTorch 在执行这种操作时,通常会 复用已有的显存。它会将中间计算结果存放在新的内存位置,并且使用现有的显存池优化内存分配。
- 在这种情况下,如果没有显式的赋值给
x
,就不会创建额外的内存开销。
-x = x + 1:会增加显存
x = x + 1
:会增加显存 x = x + 1
这一操作实际上会执行 "计算+赋值" 。这一过程中:- 中间结果 (即
x + 1
计算出的新张量)会被存储到 新内存位置。 - 原来存储
x
的显存会被新值 覆盖 ,但是由于这是一个"计算+赋值"操作,PyTorch 为了避免覆盖数据,通常会先分配新的内存空间来存储计算结果。 - 这意味着,新张量 和原始张量会在一段时间内 共享显存 ,直到 Python 的垃圾回收机制清理旧张量的内存。
在函数调用期间,新的张量(x + 1
)会占用新的显存,而原来的x
张量的内存要等到函数结束或者垃圾回收时才能释放。因此,暂时增加了显存占用。
- 中间结果 (即