intel / intel/torch-xpu-ops

sdpa backward kernel is required to reduce memory usage

Open
#2,232 11 comments 2 reactions 1 assignee Claimed by @LuFinch View on GitHub
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.

Image
Image

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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.