forked from ccf-ai-infra/GPUCodeForces
fix CrossEntropyLoss #6
This commit is contained in:
parent
e8d83740df
commit
6968605bfb
|
|
@ -0,0 +1,156 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
BATCH_SIZE = 4096
|
||||
N_CLASSES = 1024
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
|
||||
// C++ 接口
|
||||
torch::Tensor cross_entropy_forward_cuda(
|
||||
torch::Tensor logits,
|
||||
torch::Tensor target
|
||||
);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
#include <float.h>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
|
||||
__global__ void cross_entropy_fused_kernel(
|
||||
const float* __restrict__ logits_data, // (N, C)
|
||||
const int64_t* __restrict__ target_data, // (N,)
|
||||
float* __restrict__ loss_per_row_out, // (N,)
|
||||
int N,
|
||||
int C
|
||||
) {
|
||||
// 当前处理的 Batch 索引
|
||||
int row_idx = blockIdx.x;
|
||||
if (row_idx >= N) return;
|
||||
|
||||
// 当前行的指针
|
||||
const float* row_logits = logits_data + row_idx * C;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
// 共享内存：用于 Max 和 Sum 的归约
|
||||
__shared__ float s_data[BLOCK_SIZE];
|
||||
|
||||
float thread_max = -FLT_MAX;
|
||||
|
||||
// Grid-Stride Loop 遍历类别 C
|
||||
for (int c = tid; c < C; c += BLOCK_SIZE) {
|
||||
float val = row_logits[c];
|
||||
if (val > thread_max) {
|
||||
thread_max = val;
|
||||
}
|
||||
}
|
||||
s_data[tid] = thread_max;
|
||||
__syncthreads();
|
||||
|
||||
// 块内归约 (Max)
|
||||
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
|
||||
if (tid < offset) {
|
||||
if (s_data[tid + offset] > s_data[tid]) {
|
||||
s_data[tid] = s_data[tid + offset];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
|
||||
float row_max_val = s_data[0];
|
||||
__syncthreads();
|
||||
|
||||
float thread_sum_exp = 0.0f;
|
||||
|
||||
for (int c = tid; c < C; c += BLOCK_SIZE) {
|
||||
float val = row_logits[c];
|
||||
thread_sum_exp += expf(val - row_max_val);
|
||||
}
|
||||
s_data[tid] = thread_sum_exp;
|
||||
__syncthreads();
|
||||
|
||||
// 块内归约 (Sum)
|
||||
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
|
||||
if (tid < offset) {
|
||||
s_data[tid] += s_data[tid + offset];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
float row_sum_exp = s_data[0];
|
||||
float log_sum_exp = logf(row_sum_exp) + row_max_val;
|
||||
|
||||
int64_t target_class = target_data[row_idx];
|
||||
float target_logit = row_logits[target_class];
|
||||
|
||||
// Cross Entropy Formula
|
||||
loss_per_row_out[row_idx] = -target_logit + log_sum_exp;
|
||||
}
|
||||
}
|
||||
|
||||
// C++ 封装函数
|
||||
torch::Tensor cross_entropy_forward_cuda(
|
||||
torch::Tensor logits,
|
||||
torch::Tensor target
|
||||
) {
|
||||
TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
|
||||
TORCH_CHECK(target.is_cuda(), "target must be a CUDA tensor");
|
||||
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
|
||||
TORCH_CHECK(target.dim() == 1, "target must be 1D");
|
||||
|
||||
// 确保连续
|
||||
logits = logits.contiguous();
|
||||
target = target.contiguous();
|
||||
|
||||
int N = logits.size(0); // Batch Size
|
||||
int C = logits.size(1); // Num Classes
|
||||
|
||||
TORCH_CHECK(target.size(0) == N, "Target size mismatch");
|
||||
|
||||
auto losses = torch::empty({N}, logits.options());
|
||||
|
||||
// 启动配置:
|
||||
// Grid: N (每个 Batch 一个 Block)
|
||||
// Block: 256
|
||||
dim3 grid_dim(N);
|
||||
dim3 block_dim(BLOCK_SIZE);
|
||||
|
||||
cross_entropy_fused_kernel<<<grid_dim, block_dim>>>(
|
||||
logits.data_ptr<float>(),
|
||||
target.data_ptr<int64_t>(),
|
||||
losses.data_ptr<float>(),
|
||||
N, C
|
||||
);
|
||||
|
||||
// 返回 Mean Reduction
|
||||
return losses.mean();
|
||||
}
|
||||
"""
|
||||
|
||||
self.ce_op = load_inline(
|
||||
name="cross_entropy_op_v1",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["cross_entropy_forward_cuda"],
|
||||
extra_cuda_cflags=["-O3"],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
return self.ce_op.cross_entropy_forward_cuda(logits, target)
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
BATCH_SIZE = 4096
|
||||
N_CLASSES = 1024
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.criterion = nn.CrossEntropyLoss(reduction='mean')
|
||||
|
||||
def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
return self.criterion(logits, target)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(BATCH_SIZE, N_CLASSES, dtype=torch.float32)
|
||||
target = torch.randint(0, N_CLASSES, (BATCH_SIZE,), dtype=torch.long)
|
||||
return [logits, target]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
import torch
|
||||
import time
|
||||
from CrossEntropyLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from CrossEntropyLoss_cuda import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用")
|
||||
return
|
||||
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 准备输入数据
|
||||
inputs = [x.cuda(device=device) for x in get_inputs()]
|
||||
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
|
||||
|
||||
# 初始化模型
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
# 预热GPU
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# 正式测试
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
# 精度验证
|
||||
abs_diff = torch.abs(output_torch - output_cuda)
|
||||
max_diff = torch.max(abs_diff).item()
|
||||
mean_diff = torch.mean(abs_diff).item()
|
||||
|
||||
if max_diff < 1e-4 and mean_diff < 1e-5:
|
||||
print(f"✅ 精度对齐：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
|
||||
precision_flag = True
|
||||
else:
|
||||
print(f"❌ 精度不一致：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
|
||||
precision_flag = False
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# 预热GPU
|
||||
for _ in range(10):
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# PyTorch模型计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = torch_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
torch_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
# 自定义CUDA内核计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = cuda_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
cuda_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
print(f"PyTorch内置Swish平均执行时间: {torch_time:.6f}秒")
|
||||
print(f"自定义CUDA Swish平均执行时间: {cuda_time:.6f}秒")
|
||||
speedup = torch_time / cuda_time if cuda_time > 0 else 0
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue