forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish DistillationLoss #108' (#468) from gsd123/GPUCodeForces:gsd108 into main
This commit is contained in:
commit
df374a044d
|
|
@ -0,0 +1,193 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
// Helper for atomicMax on floats using CAS loop
|
||||
__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
|
||||
int* addr_as_i = (int*)addr;
|
||||
int old = *addr_as_i;
|
||||
int assumed;
|
||||
do {
|
||||
assumed = old;
|
||||
if (__int_as_float(assumed) >= value) {
|
||||
return __int_as_float(assumed);
|
||||
}
|
||||
old = atomicCAS(addr_as_i, assumed, __float_as_int(value));
|
||||
} while (assumed != old);
|
||||
return __int_as_float(old);
|
||||
}
|
||||
|
||||
// Warp Reduction Helper for Max
|
||||
template <typename T>
|
||||
__device__ void warpReduceMax(T& val) {
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val = max(val, __shfl_down_sync(0xffffffff, val, offset));
|
||||
}
|
||||
|
||||
// Warp Reduction Helper for Sum
|
||||
template <typename T>
|
||||
__device__ void warpReduceSum(T& val) {
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
}
|
||||
|
||||
__global__ void distillation_loss_kernel(
|
||||
const float* __restrict__ student_logits,
|
||||
const float* __restrict__ teacher_logits,
|
||||
float* __restrict__ output_loss,
|
||||
const int batch_size,
|
||||
const int num_classes,
|
||||
const float temperature) {
|
||||
|
||||
int row = blockIdx.x;
|
||||
if (row >= batch_size) return;
|
||||
|
||||
// Shared memory layout:
|
||||
// [0]: s_max, [1]: t_max, [2]: s_sum, [3]: t_sum, [4]: kl_sum
|
||||
extern __shared__ float shared_mem[];
|
||||
float* s_max_shared = shared_mem;
|
||||
float* t_max_shared = shared_mem + 1;
|
||||
float* s_sum_shared = shared_mem + 2;
|
||||
float* t_sum_shared = shared_mem + 3;
|
||||
float* kl_sum_shared = shared_mem + 4;
|
||||
|
||||
// Initialize shared memory
|
||||
if (threadIdx.x == 0) {
|
||||
*s_max_shared = -3.402823466e+38F;
|
||||
*t_max_shared = -3.402823466e+38F;
|
||||
*s_sum_shared = 0.0f;
|
||||
*t_sum_shared = 0.0f;
|
||||
*kl_sum_shared = 0.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 1. Find Max (for numerical stability)
|
||||
float local_s_max = -3.402823466e+38F;
|
||||
float local_t_max = -3.402823466e+38F;
|
||||
|
||||
for (int col = threadIdx.x; col < num_classes; col += blockDim.x) {
|
||||
float s_val = student_logits[row * num_classes + col] / temperature;
|
||||
float t_val = teacher_logits[row * num_classes + col] / temperature;
|
||||
local_s_max = max(local_s_max, s_val);
|
||||
local_t_max = max(local_t_max, t_val);
|
||||
}
|
||||
|
||||
warpReduceMax(local_s_max);
|
||||
warpReduceMax(local_t_max);
|
||||
|
||||
if (threadIdx.x % 32 == 0) {
|
||||
atomicMaxFloat(s_max_shared, local_s_max);
|
||||
atomicMaxFloat(t_max_shared, local_t_max);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float s_max = *s_max_shared;
|
||||
float t_max = *t_max_shared;
|
||||
|
||||
// 2. Compute Exp Sum
|
||||
float local_s_sum = 0.0f;
|
||||
float local_t_sum = 0.0f;
|
||||
|
||||
for (int col = threadIdx.x; col < num_classes; col += blockDim.x) {
|
||||
float s_val = student_logits[row * num_classes + col] / temperature;
|
||||
float t_val = teacher_logits[row * num_classes + col] / temperature;
|
||||
local_s_sum += expf(s_val - s_max);
|
||||
local_t_sum += expf(t_val - t_max);
|
||||
}
|
||||
|
||||
warpReduceSum(local_s_sum);
|
||||
warpReduceSum(local_t_sum);
|
||||
|
||||
if (threadIdx.x % 32 == 0) {
|
||||
atomicAdd(s_sum_shared, local_s_sum);
|
||||
atomicAdd(t_sum_shared, local_t_sum);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float s_sum = *s_sum_shared;
|
||||
float t_sum = *t_sum_shared;
|
||||
float log_s_sum = logf(s_sum);
|
||||
float log_t_sum = logf(t_sum);
|
||||
|
||||
// 3. Compute KL Divergence
|
||||
// KL = sum( p_teacher * (log(p_teacher) - log(p_student)) )
|
||||
// log(p) = logits - max - log(sum)
|
||||
float local_kl_sum = 0.0f;
|
||||
|
||||
for (int col = threadIdx.x; col < num_classes; col += blockDim.x) {
|
||||
float s_val = student_logits[row * num_classes + col] / temperature;
|
||||
float t_val = teacher_logits[row * num_classes + col] / temperature;
|
||||
|
||||
float p_t = expf(t_val - t_max) / t_sum;
|
||||
float log_p_s = s_val - s_max - log_s_sum;
|
||||
float log_p_t = t_val - t_max - log_t_sum;
|
||||
|
||||
local_kl_sum += p_t * (log_p_t - log_p_s);
|
||||
}
|
||||
|
||||
warpReduceSum(local_kl_sum);
|
||||
|
||||
if (threadIdx.x % 32 == 0) {
|
||||
atomicAdd(kl_sum_shared, local_kl_sum);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
output_loss[row] = *kl_sum_shared;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor distillation_loss_cuda_launcher(torch::Tensor student, torch::Tensor teacher, float T) {
|
||||
int batch_size = student.size(0);
|
||||
int num_classes = student.size(1);
|
||||
auto output_loss = torch::zeros({batch_size}, student.options());
|
||||
|
||||
const int threads = 256;
|
||||
const int blocks = batch_size;
|
||||
const int shared_mem_size = 5 * sizeof(float);
|
||||
|
||||
distillation_loss_kernel<<<blocks, threads, shared_mem_size>>>(
|
||||
student.data_ptr<float>(),
|
||||
teacher.data_ptr<float>(),
|
||||
output_loss.data_ptr<float>(),
|
||||
batch_size,
|
||||
num_classes,
|
||||
T
|
||||
);
|
||||
|
||||
return output_loss;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor distillation_loss_cuda_launcher(torch::Tensor student, torch::Tensor teacher, float T);
|
||||
"""
|
||||
|
||||
# Compile the inline CUDA code
|
||||
distillation_loss = load_inline(
|
||||
name='distillation_loss_cuda',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['distillation_loss_cuda_launcher'],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, temperature):
|
||||
super(ModelNew, self).__init__()
|
||||
self.temperature = temperature
|
||||
self.distillation_loss = distillation_loss
|
||||
|
||||
def forward(self, student_logits, teacher_logits):
|
||||
# Call the custom CUDA kernel
|
||||
batch_losses = self.distillation_loss.distillation_loss_cuda_launcher(
|
||||
student_logits, teacher_logits, self.temperature
|
||||
)
|
||||
# Apply temperature scaling factor and reduction
|
||||
return batch_losses.mean() * (self.temperature ** 2)
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
|
||||
def __init__(self, temperature):
|
||||
super(Model, self).__init__()
|
||||
self.temperature = temperature
|
||||
|
||||
def forward(self, student_logits: torch.Tensor, teacher_logits: torch.Tensor) -> torch.Tensor:
|
||||
soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
|
||||
soft_teacher = F.softmax(teacher_logits / self.temperature, dim=1)
|
||||
kl_div = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
|
||||
return kl_div * (self.temperature ** 2)
|
||||
|
||||
|
||||
batch_size = 128
|
||||
num_classes = 1000
|
||||
temperature = 4.0
|
||||
|
||||
|
||||
def get_inputs():
|
||||
student = torch.randn(batch_size, num_classes, requires_grad=True)
|
||||
teacher = torch.randn(batch_size, num_classes)
|
||||
return [student, teacher]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [temperature]
|
||||
|
|
@ -0,0 +1,60 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given GeGLU architecture to get speedups.
|
||||
|
||||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
|
||||
Custom CUDA kernel extension via torch.utils.cpp_extension.load_inline
|
||||
|
||||
Knowledge distillation loss computation (KL divergence with temperature scaling)
|
||||
|
||||
Warp-level reduction templates for max and sum operations
|
||||
|
||||
Atomic max operation on floats using compare-and-swap (CAS) loop
|
||||
|
||||
Numerically stable softmax with max subtraction
|
||||
|
||||
Shared memory caching for intermediate statistics (max, sum, KL)
|
||||
|
||||
Per-batch parallel processing (one CUDA block per sample)
|
||||
|
||||
Temperature scaling applied to logits before softmax
|
||||
|
||||
KL divergence calculation: Σ p_teacher·(log(p_teacher) - log(p_student))
|
||||
|
||||
Contiguous tensor handling for memory coalescing
|
||||
|
||||
Temperature-squared scaling in final loss reduction
|
||||
|
||||
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
|
||||
def __init__(self, temperature):
|
||||
super(Model, self).__init__()
|
||||
self.temperature = temperature
|
||||
|
||||
def forward(self, student_logits: torch.Tensor, teacher_logits: torch.Tensor) -> torch.Tensor:
|
||||
soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
|
||||
soft_teacher = F.softmax(teacher_logits / self.temperature, dim=1)
|
||||
kl_div = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
|
||||
return kl_div * (self.temperature ** 2)
|
||||
|
||||
|
||||
batch_size = 128
|
||||
num_classes = 1000
|
||||
temperature = 4.0
|
||||
|
||||
|
||||
def get_inputs():
|
||||
student = torch.randn(batch_size, num_classes, requires_grad=True)
|
||||
teacher = torch.randn(batch_size, num_classes)
|
||||
return [student, teacher]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [temperature]
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from DistillationLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from DistillationLoss_cuda import ModelNew
|
||||
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 初始化模型
|
||||
init_inputs = get_init_inputs()
|
||||
init_inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
|
||||
]
|
||||
inputs = get_inputs()
|
||||
inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
|
||||
]
|
||||
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# 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 torch.relu 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag, speedup
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue