sdpa backward kernel is required to reduce memory usage
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 9h
- Merged PRs (30d)
- 112
Description
### 🚀 The feature, motivation and pitch
I'm working on intel/AutoRound and would like to support XPU device. However, sdpa op is taking much more memory than expected and causing OOM. I heard that we already have a plan to support sdpa backward kernel for XPU, this issue is used to track the status and for you to reproduce my issue.
To reproduce:
```python
import torch
import time
import gc
from typing import Tuple, Dict, Any
def get_device():
"""????????????????"""
if torch.cuda.is_available():
device = torch.device("cuda")
device_type = "cuda"
elif hasattr(torch, 'xpu') and torch.xpu.is_available():
device = torch.device("xpu")
device_type = "xpu"
else:
device = torch.device("cpu")
device_type = "cpu"
print(f"Using device: {device_type.upper()}")
return device, device_type
def get_memory_info(device_type: str) -> Dict[str, float]:
"""????????????????"""
memory_info = {}
if device_type == "cuda":
memory_info["allocated"] = torch.cuda.memory_allocated() / 1024**3 # GB
memory_info["reserved"] = torch.cuda.memory_reserved() / 1024**3 # GB
memory_info["max_allocated"] = torch.cuda.max_memory_allocated() / 1024**3 # GB
memory_info["max_reserved"] = torch.cuda.max_memory_reserved() / 1024**3 # GB
device_props = torch.cuda.get_device_properties(0)
memory_info["total"] = device_props.total_memory / 1024**3 # GB
elif device_type == "xpu":
memory_info["allocated"] = torch.xpu.memory_allocated() / 1024**3 # GB
memory_info["reserved"] = torch.xpu.memory_reserved() / 1024**3 # GB
memory_info["max_allocated"] = torch.xpu.max_memory_allocated() / 1024**3 # GB
memory_info["max_reserved"] = torch.xpu.max_memory_reserved() / 1024**3 # GB
return memory_info
def print_memory_info(device_type: str, prefix: str = ""):
"""????????????"""
if device_type == "cpu":
return
memory_info = get_memory_info(device_type)
print(f"\n{prefix}Memory Info ({device_type.upper()}):")
print("-" * 50)
for key, value in memory_info.items():
print(f"{key:15}: {value:.3f} GB")
def test_scaled_dot_product_attention(
batch_size: int = 1,
seq_len: int = 4096,
num_heads: int = 32,
head_dim: int = 128,
dtype: torch.dtype = torch.float16,
enable_grad: bool = True,
dropout: float = 0.1,
is_causal: bool = True,
num_iterations: int = 5
) -> Tuple[torch.Tensor, Dict[str, Any]]:
"""
???? scaled_dot_product_attention ???? (Large Model Configuration)
Args:
batch_size: ????????
seq_len: ????????
num_heads: ??????????
head_dim: ????????????
dtype: ????????
enable_grad: ????????????????????????????
dropout: dropout ????
is_causal: ????????????????
num_iterations: ????????
Returns:
??????????????????????
"""
# ????????????
device, device_type = get_device()
print(f"\nConfig: batch_size={batch_size}, seq_len={seq_len}, num_heads={num_heads}, head_dim={head_dim}")
print(f"dtype={dtype}, enable_grad={enable_grad}, dropout={dropout}, is_causal={is_causal}")
# ????????
gc.collect()
if device_type == "cuda":
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
elif device_type == "xpu":
torch.xpu.empty_cache()
torch.xpu.reset_peak_memory_stats()
print_memory_info(device_type, "Initial ")
# ????????????
print("\nCreating input tensors...")
query = torch.randn(batch_size, num_heads, seq_len, head_dim,
device=device, dtype=dtype, requires_grad=enable_grad)
key = torch.randn(batch_size, num_heads, seq_len, head_dim,
device=device, dtype=dtype, requires_grad=enable_grad)
value = torch.randn(batch_size, num_heads, seq_len, head_dim,
device=device, dtype=dtype, requires_grad=enable_grad)
scaling = 1.0 / (head_dim ** 0.5)
print_memory_info(device_type, "After tensor creation ")
# ????????
stats = {
"forward_times": [],
"backward_times": [],
"memory_usage": []
}
print(f"\nRunning {num_iterations} iterations in TRAIN mode...")
for i in range(num_iterations):
print(f"\nIteration {i+1}/{num_iterations}")
# ????????
start_time = time.time()
attn_output = torch.nn.functional.scaled_dot_product_attention(
query,
key,
value,
attn_mask=None,
dropout_p=dropout,
scale=scaling,
is_causal=is_causal,
)
if device_type in ["cuda", "xpu"]:
getattr(torch, device_type).synchronize()
forward_time = time.time() - start_time
stats["forward_times"].append(forward_time)
# ????????
loss = attn_output.sum()
start_time = time.time()
loss.backward()
if device_type in ["cuda", "xpu"]:
getattr(torch, device_type).synchronize()
backward_time = time.time() - start_time
stats["backward_times"].append(backward_time)
# ????????
query.grad = None
key.grad = None
value.grad = None
# ????????????????
memory_info = get_memory_info(device_type)
stats["memory_usage"].append(memory_info)
print(f" Forward time: {forward_time:.4f}s, Backward time: {backward_time:.4f}s")
if device_type in ["cuda", "xpu"] and "allocated" in memory_info:
print(f" Memory allocated: {memory_info['allocated']:.3f} GB")
print_memory_info(device_type, "Final ")
# ????????????
avg_forward_time = sum(stats["forward_times"]) / len(stats["forward_times"])
avg_backward_time = sum(stats["backward_times"]) / len(stats["backward_times"])
print(f"\n{'='*60}")
print("Performance Summary:")
print(f"{'='*60}")
print(f"Average forward time: {avg_forward_time:.4f}s")
print(f"Average backward time: {avg_backward_time:.4f}s")
print(f"Average total time: {avg_forward_time + avg_backward_time:.4f}s")
# ????????????
print(f"\n{'='*60}")
print("Output Tensor Info:")
print(f"{'='*60}")
print(f"Shape: {attn_output.shape}")
print(f"Dtype: {attn_output.dtype}")
print(f"Device: {attn_output.device}")
print(f"Requires_grad: {attn_output.requires_grad}")
return attn_output, stats
def main():
"""?????? - ??????????????"""
print("=" * 60)
print("Scaled Dot Product Attention Memory Test")
print("Large Model Configuration")
print("=" * 60)
try:
output, stats = test_scaled_dot_product_attention()
print(f"\n?? Test completed successfully!")
except Exception as e:
print(f"\n?? Test failed: {str(e)}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()
```
### Alternatives
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.