aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-27 20:25:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #653 from ademeure/cublaslt_refactor

Browse source 
Matmul refactor using only cuBLASLt + GELU Fusion
This commit is contained in:
Andrej 2024-07-01 11:53:02 -07:00 • committed by GitHub
commit a876282eb8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 148 additions and 104 deletions
Download patch file Download diff file Expand all files Collapse all files

40
llmc/attention.cuh
View file

@ -199,7 +199,6 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
// Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer.
// Its contents will be overwritten by this function.
const int block_size = 256;
const float alpha = 1.0f, beta = 0.0f;
// inp is (B, T, 3C) QKV
// preatt, att are (B, NH, T, T)
@ -215,16 +214,8 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
int num_blocks = CEIL_DIV(total_threads, block_size);
permute_kernel<<<num_blocks, block_size, 0, stream>>>(q, k, v, inp, B, T, NH, HS);
floatX* preatt = inp;
cublasCheck(cublasSetStream(cublas_handle, stream));
cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
CUBLAS_OP_T, CUBLAS_OP_N,
T, T, HS, &alpha,
k, CUBLAS_LOWP, HS, T * HS,
q, CUBLAS_LOWP, HS, T * HS,
&beta, preatt, CUBLAS_LOWP, T, T * T,
B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
floatX* preatt = inp; // reuse inp as scratch buffer
matmul_cublaslt(preatt, k, q, nullptr, T, T, HS, stream, true, false, B * NH, T * HS, T * HS, T * T);
// multiply all elements of preatt elementwise by scale
float scale = 1.f / sqrtf(HS);
@ -234,13 +225,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
// new approach: first cuBLAS another batched matmul
floatX* vaccum = inp;
// y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
CUBLAS_OP_N, CUBLAS_OP_N,
HS, T, T, &alpha,
v, CUBLAS_LOWP, HS, T * HS,
att, CUBLAS_LOWP, T, T * T,
&beta, vaccum, CUBLAS_LOWP, HS, T * HS,
B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
matmul_cublaslt(vaccum, v, att, nullptr, HS, T, T, stream, false, false, B * NH, T * HS, T * T, T * HS);
// now unpermute
// y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
@ -258,7 +243,6 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* datt, floatX* scrat
NVTX_RANGE_FN();
const int block_size = 256;
const int HS = C / NH; // head size
const float alpha = 1.0f, beta = 0.0f;
// unpack convenience pointers into q, k, v
const floatX *q, *k, *v;
@ -274,27 +258,17 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* datt, floatX* scrat
int num_blocks = CEIL_DIV(B * T * C, block_size);
unpermute_kernel_backward<<<num_blocks, block_size, 0, stream>>>(scratch, dout, B, T, NH, HS);
// backward into datt
cublasCheck(cublasSetStream(cublas_handle, stream));
cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha,
v, CUBLAS_LOWP, HS, T * HS, scratch, CUBLAS_LOWP, HS, T * HS, &beta,
datt, CUBLAS_LOWP, T, T * T, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
matmul_cublaslt(datt, v, scratch, nullptr, T, T, HS, stream, true, false, B * NH, T * HS, T * HS, T * T);
// backward into dv
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha,
scratch, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, &beta,
dv, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
matmul_cublaslt(dv, scratch, att, nullptr, HS, T, T, stream, false, true, B * NH, T * HS, T * T, T * HS);
const float scale = 1.0f / sqrtf((float)HS);
// backward into preatt. this is an in-place operation; datt turns into dpreatt here
softmax_autoregressive_backward_inplace_kernel<<<dim3(T / 4, B * NH), 256>>>(datt, att, B, T, C, scale);
const floatX* dpreatt = datt;
// backward into q
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha,
k, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta,
dq, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
matmul_cublaslt(dq, k, dpreatt, nullptr, HS, T, T, stream, false, false, B * NH, T * HS, T * T, T * HS);
// backward into k
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha,
q, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta,
dk, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
matmul_cublaslt(dk, q, dpreatt, nullptr, HS, T, T, stream, false, true, B * NH, T * HS, T * T, T * HS);
// backward into inp
num_blocks = CEIL_DIV(B * NH * T * HS, block_size);
permute_kernel_backward<<<num_blocks, block_size, 0, stream>>>(dinp, dq, dk, dv, B, T, NH, HS);

4
llmc/cublas_common.h
View file

@ -14,13 +14,10 @@ cuBLAS related utils
// cuBLAS Precision settings
#if defined(ENABLE_FP32)
typedef float floatX;
#define CUBLAS_LOWP CUDA_R_32F
#elif defined(ENABLE_FP16)
typedef half floatX;
#define CUBLAS_LOWP CUDA_R_16F
#else // default to bfloat16
typedef __nv_bfloat16 floatX;
#define CUBLAS_LOWP CUDA_R_16BF
#endif
@ -32,7 +29,6 @@ const size_t cublaslt_workspace_size = 32 * 1024 * 1024;
void* cublaslt_workspace = NULL;
cublasComputeType_t cublas_compute = CUBLAS_COMPUTE_32F;
cublasLtHandle_t cublaslt_handle;
cublasHandle_t cublas_handle;
// ----------------------------------------------------------------------------
// Error checking

172
llmc/matmul.cuh
View file

@ -7,6 +7,8 @@ Matrix Multiplication, with help from cuBLASLt
#include "cuda_common.h"
#include "cuda_utils.cuh"
#include "cublas_common.h"
// GELU can be either fused (cublasLt) or non-fused (gelu.h)
#include "gelu.cuh"
// ----------------------------------------------------------------------------
// CUDA kernels
@ -102,83 +104,149 @@ __global__ void reduce_add_sum_kernel(floatX* dst, const float* src, size_t n, s
// ----------------------------------------------------------------------------
// kernel launchers
// Wrapper around cublasLtMatmul that is meant to support everything we need in llm.c
// https://docs.nvidia.com/cuda/cublas/#cublasltmatmul
void matmul_forward_cublaslt(floatX* out,
floatX* inp, floatX* weight, floatX* bias,
int B, int T, int C, int OC, cudaStream_t stream) {
void matmul_cublaslt(floatX* d, const floatX* a, const floatX* b, const floatX* bias,
int m, int n, int k, cudaStream_t stream=0, bool transA=true, bool transB=false,
int batch_count=0, size_t strideA=0, size_t strideB=0, size_t strideOut=0,
bool accumulate=false, floatX* pre_gelu=NULL, bool backward=false)
{
NVTX_RANGE_FN();
int has_bias = (bias != NULL);
bool has_bias = (bias != NULL);
bool has_gelu = (pre_gelu != NULL);
// check bias alignment
if(((uintptr_t)bias % 16) != 0) {
printf("Bias pointer is not aligned (cuBLASLt requirement)!\n");
// check alignment (some modes work unaligned but it always best to be aligned for performance)
if(((uintptr_t)a % 16) != 0 || ((uintptr_t)b % 16) != 0 || ((uintptr_t)d % 16) != 0 || ((uintptr_t)bias % 16) != 0) {
printf("All cuBLASLt pointers must be aligned!\n");
exit(EXIT_FAILURE);
}
// these need to be in FP16 if and only if alpha/beta are CUBLAS_COMPUTE_16F
const float alpha = 1.0f, beta = 0.0f;
// create the operation descriptor
cublasLtMatmulDesc_t operationDesc;
cublasCheck(cublasLtMatmulDescCreate(&operationDesc, cublas_compute, CUDA_R_32F));
int returnedResults = 0;
cublasLtMatmulDesc_t operationDesc;
cublasLtMatmulPreference_t preference;
cublasLtMatrixLayout_t weightLayout;
cublasLtMatrixLayout_t inputLayout;
cublasLtMatrixLayout_t outputLayout;
cublasLtMatrixLayout_t biasLayout;
cublasLtMatmulHeuristicResult_t heuristic;
// create the operation descriptor
cublasOperation_t opNoTranspose = CUBLAS_OP_N;
cublasOperation_t opTranspose = CUBLAS_OP_T;
cublasLtEpilogue_t epilogueBias = has_bias ? CUBLASLT_EPILOGUE_BIAS : CUBLASLT_EPILOGUE_DEFAULT;
cublasCheck(cublasLtMatmulDescCreate(&operationDesc, cublas_compute, CUDA_R_32F)); // FP16 if CUBLAS_COMPUTE_16F
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opTranspose, sizeof(opTranspose)));
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opNoTranspose, sizeof(opNoTranspose)));
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogueBias, sizeof(epilogueBias)));
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)));
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, (transA) ? &opTranspose : &opNoTranspose, sizeof(opTranspose)));
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, (transB) ? &opTranspose : &opNoTranspose, sizeof(opNoTranspose)));
// define matrix layouts
cublasCheck(cublasLtMatrixLayoutCreate(&weightLayout, CUBLAS_LOWP, C, OC, C));
cublasCheck(cublasLtMatrixLayoutCreate(&inputLayout, CUBLAS_LOWP, C, B*T, C));
cublasCheck(cublasLtMatrixLayoutCreate(&outputLayout, CUBLAS_LOWP, OC, B*T, OC));
cublasCheck(cublasLtMatrixLayoutCreate(&biasLayout, CUBLAS_LOWP, OC, 1, OC));
cublasLtMatrixLayout_t ALayout;
cublasLtMatrixLayout_t BLayout;
cublasLtMatrixLayout_t DLayout;
cublasLtMatrixLayout_t CLayout;
if (transA) {
cublasCheck(cublasLtMatrixLayoutCreate(&ALayout, CUBLAS_LOWP, k, m, k));
} else {
cublasCheck(cublasLtMatrixLayoutCreate(&ALayout, CUBLAS_LOWP, m, k, m));
}
if (transB) {
cublasCheck(cublasLtMatrixLayoutCreate(&BLayout, CUBLAS_LOWP, n, k, n));
} else {
cublasCheck(cublasLtMatrixLayoutCreate(&BLayout, CUBLAS_LOWP, k, n, k));
}
// cuBLASLt requires C in FP8 mode to be BF16 or FP32... (sigh)
cublasCheck(cublasLtMatrixLayoutCreate(&CLayout, (sizeof(floatX) == 1) ? CUDA_R_16BF : CUBLAS_LOWP, m, n, m));
cublasCheck(cublasLtMatrixLayoutCreate(&DLayout, CUBLAS_LOWP, m, n, m));
// Strided Batched GEMM (used for non-flash attention, equivalent to cublasGemmStridedBatchedEx)
if (batch_count) {
cublasCheck(cublasLtMatrixLayoutSetAttribute(ALayout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count)));
cublasCheck(cublasLtMatrixLayoutSetAttribute(BLayout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count)));
cublasCheck(cublasLtMatrixLayoutSetAttribute(CLayout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count)));
cublasCheck(cublasLtMatrixLayoutSetAttribute(DLayout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count)));
cublasCheck(cublasLtMatrixLayoutSetAttribute(ALayout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &strideA, sizeof(strideA)));
cublasCheck(cublasLtMatrixLayoutSetAttribute(BLayout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &strideB, sizeof(strideB)));
cublasCheck(cublasLtMatrixLayoutSetAttribute(CLayout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &strideOut, sizeof(strideOut)));
cublasCheck(cublasLtMatrixLayoutSetAttribute(DLayout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &strideOut, sizeof(strideOut)));
}
// create a preference handle with specified max workspace
cublasCheck(cublasLtMatmulPreferenceCreate(&preference));
cublasCheck(cublasLtMatmulPreferenceSetAttribute(preference,
CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &cublaslt_workspace_size, sizeof(cublaslt_workspace_size)));
cublasCheck(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&cublaslt_workspace_size, sizeof(cublaslt_workspace_size)));
// find a suitable algorithm
cublasCheck(cublasLtMatmulAlgoGetHeuristic(cublaslt_handle, operationDesc,
weightLayout, inputLayout, outputLayout, outputLayout,
preference, 1, &heuristic, &returnedResults));
// setup epilogue and associated pointers for bias & gelu
cublasLtEpilogue_t epilogue;
if (has_gelu) {
int64_t gelu_ld = m; // todo - is this affected by anything else?
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &gelu_ld, sizeof(gelu_ld)));
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER, &pre_gelu, sizeof(pre_gelu)));
if (backward) {
assert(!has_bias); // we shouldn't have any backward matmuls that use both GELU and bias
epilogue = CUBLASLT_EPILOGUE_DGELU;
} else {
epilogue = has_bias ? CUBLASLT_EPILOGUE_GELU_AUX_BIAS : CUBLASLT_EPILOGUE_GELU_AUX;
}
} else if(has_bias){
epilogue = backward ? CUBLASLT_EPILOGUE_BGRADB : CUBLASLT_EPILOGUE_BIAS;
} else {
epilogue = CUBLASLT_EPILOGUE_DEFAULT;
}
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue)));
if (has_bias) {
// cuBLASLt requires bias in FP8 mode to be BF16... (sigh)
cublasDataType_t bias_data_type = (sizeof(floatX) == 1) ? CUDA_R_16BF : CUBLAS_LOWP; // force BF16 bias for FP8 mode
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE, &bias_data_type, sizeof(bias_data_type)));
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)));
}
// set scale type to FP32 (needs to be FP16 if and only if using CUBLAS_COMPUTE_16F, so it's FP32 even for FP8!)
cublasDataType_t scale_type = CUDA_R_32F;
cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type)));
// find a suitable algorithm (cached internally so shouldn't take much CPU time in practice)
cublasLtMatmulAlgoGetHeuristic(cublaslt_handle, operationDesc, ALayout, BLayout, CLayout, DLayout,
preference, 1, &heuristic, &returnedResults);
if (returnedResults == 0) {
printf("No cuBLASLt algorithm: B: %d, T: %d, C: %d, OC: %d, bias: %d\n", B, T, C, OC, has_bias);
printf("No cuBLASLt algorithm: m: %d, n: %d, k: %d, bias: %d\n", n, m, k, has_bias);
exit(EXIT_FAILURE);
}
// set whether to accumulate (i.e. D += C) or not - note this isn't considered in algorithm selection (?!)
const float alpha = 1.0f, beta = accumulate ? 1.0f : 0.0f;
// call the matmul
cublasCheck(cublasLtMatmul(cublaslt_handle, operationDesc,
&alpha, weight, weightLayout, inp, inputLayout, &beta,
out, outputLayout, out, outputLayout, &heuristic.algo,
cublaslt_workspace, cublaslt_workspace_size, stream));
&alpha, a, ALayout, b, BLayout, &beta, d, CLayout, d, DLayout,
&heuristic.algo, cublaslt_workspace, cublaslt_workspace_size, stream));
// cleanups
cublasCheck(cublasLtMatmulPreferenceDestroy(preference));
cublasCheck(cublasLtMatmulDescDestroy(operationDesc));
cublasCheck(cublasLtMatrixLayoutDestroy(weightLayout));
cublasCheck(cublasLtMatrixLayoutDestroy(inputLayout));
cublasCheck(cublasLtMatrixLayoutDestroy(outputLayout));
cublasCheck(cublasLtMatrixLayoutDestroy(biasLayout));
cublasCheck(cublasLtMatrixLayoutDestroy(ALayout));
cublasCheck(cublasLtMatrixLayoutDestroy(BLayout));
cublasCheck(cublasLtMatrixLayoutDestroy(CLayout));
cublasCheck(cublasLtMatrixLayoutDestroy(DLayout));
cudaCheck(cudaGetLastError());
}
// small wrapper around matmul_cublaslt for the forward pass (keeping historical order of arguments)
void matmul_forward_cublaslt(floatX* out,
floatX* inp, floatX* weight, floatX* bias,
int B, int T, int C, int OC, cudaStream_t stream,
floatX* pre_gelu=NULL, int gelu_fusion=1) {
// By default only fuse GELU for H100+ as cuBLAS seems to be inefficient for fused GELU on Ada/Ampere (?)
if (gelu_fusion < 1 && pre_gelu) {
matmul_cublaslt(pre_gelu, weight, inp, bias, OC, B*T, C, stream, true, false, 0, 0, 0, 0, false, NULL, false);
gelu_forward(out, pre_gelu, B*T*OC, stream);
} else {
matmul_cublaslt(out, weight, inp, bias, OC, B*T, C, stream, true, false, 0, 0, 0, 0, false, pre_gelu, false);
}
}
void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias,
floatX* dout, floatX* inp, floatX* weight,
float* dbias_buffer,
int B, int T, int C, int OC, cudaStream_t stream) {
int B, int T, int C, int OC, cudaStream_t stream,
floatX* pre_gelu=NULL, int gelu_fusion=1) {
NVTX_RANGE_FN();
float one = 1.0f, zero = 0.0f;
// backward to bias, if given, does a +=
if (dbias != NULL) {
@ -204,17 +272,19 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias,
reduce_add_sum_kernel<<<CEIL_DIV(OC, 256 * f128::size), 256, 0, stream>>>(dbias, dbias_buffer, OC, grid_size_y);
cudaCheck(cudaGetLastError());
}
dbias = NULL; // prevent dbias calculation from also being fused in matmul_cublaslt below (if we enabled fusion)
}
// backward to input, uses = in the backward pass (set the gradient)
cublasCheck(cublasSetStream(cublas_handle, stream));
cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B*T, OC, &one,
weight, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &zero,
dinp, CUBLAS_LOWP, C, cublas_compute, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
matmul_cublaslt(dinp, weight, dout, NULL, C, B*T, OC, stream, false, false, 0, 0, 0, 0, false,
gelu_fusion >= 2 ? pre_gelu : NULL, true);
// backward GELU (if it wasn't fused into the matmul above)
if (gelu_fusion < 2 && pre_gelu) {
gelu_backward_inplace(dinp, pre_gelu, B*T*C, stream);
}
// backward to weight, uses += in the backward pass (accumulate the gradient) by setting alpha=one
cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one,
inp, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &one,
dweight, CUBLAS_LOWP, C, cublas_compute, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
cudaCheck(cudaGetLastError());
matmul_cublaslt(dweight, inp, dout, NULL /*dbias*/, C, OC, B*T, stream, false, true, 0, 0, 0, 0,
true /* accumulate */, NULL, true);
}

2
test_gpt2.cu
View file

@ -118,9 +118,11 @@ int main(int argc, char *argv[]) {
for (int i = 1; i < argc; i+=2) {
if (i + 1 >= argc) { exit(EXIT_FAILURE); } // must have arg after flag
if (!(strlen(argv[i]) == 2 || strlen(argv[i]) == 3)) { exit(EXIT_FAILURE); } // must be -x[y] (one dash, one or two letters)
if (argv[i][0] != '-') { exit(EXIT_FAILURE); } // must start with dash
if (argv[i][1] == 'w') { model.use_master_weights = atoi(argv[i+1]); }
else if (argv[i][1] == 'r') { model.recompute = atoi(argv[i+1]); }
else if (argv[i][1] == 'g' && argv[i][2] == 'e') { model.gelu_fusion = atoi(argv[i+1]); }
}
// load additional information that we will use for debugging and error checking

34
train_gpt2.cu
View file

@ -47,8 +47,8 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage.
#include "llmc/encoder.cuh"
// defines: layernorm_forward, residual_forward, fused_residual_forward5, layernorm_backward
#include "llmc/layernorm.cuh"
// defines: gelu_forward, gelu_backward_inplace
#include "llmc/gelu.cuh"
// defines: matmul_cublaslt, matmul_forward, matmul_backward, gelu_forward, gelu_backward_inplace
#include "llmc/matmul.cuh"
#ifdef ENABLE_CUDNN
// defines: create_cudnn, destroy_cudnn, attention_forward_cudnn, attention_backward_cudnn
#include "llmc/cudnn_att.h"
@ -56,8 +56,6 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage.
// defines: attention_forward, attention_backward
#include "llmc/attention.cuh"
#endif
// defines: matmul_forward, matmul_backward
#include "llmc/matmul.cuh"
// defines: fused_classifier
#include "llmc/fused_classifier.cuh"
// defines: adamw_kernel3
@ -366,6 +364,7 @@ typedef struct {
float* cpu_losses; // CPU buffer to copy the losses to, allocated with cudaMallocHost
unsigned long long rng_state; // the RNG state for seeding stochastic rounding etc.
int use_master_weights; // keep master weights copy in float for optim update? 0|1
int gelu_fusion; // fuse gelu via cuBLASLt (0=none, 1=forward, 2=forward+backward)
int recompute; // recompute gelu | layernorm forward during model backward? 0|1|2
// todo - if other functions need cpu scratch buffers in the future, reuse as generic scratch?
int* workload_indices; // encoder_backward, B*T*num_c_groups (int)
@ -400,6 +399,7 @@ void gpt2_init_common(GPT2 *model) {
model->rng_state = 13371337 + multi_gpu_config.process_rank; // used in stochastic rounding
model->use_master_weights = 1; // safe default: do keep master weights in fp32
model->recompute = 1; // good default: recompute gelu but not layernorm
model->gelu_fusion = 0; //deviceProp.major >= 9 ? 2 : 0; // default: off for now (default must match main())
}
void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) {
@ -685,8 +685,7 @@ void gpt2_forward(GPT2 *model, const int* inputs, size_t B, size_t T) {
matmul_forward_cublaslt(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C, main_stream);
fused_residual_forward5(l_residual2, l_ln2, l_ln2_mean, l_ln2_rstd, residual, l_attproj, l_ln2w, l_ln2b, B*T, C, main_stream);
matmul_forward_cublaslt(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C, main_stream);
gelu_forward(l_fch_gelu, l_fch, B*T*4*C, main_stream);
matmul_forward_cublaslt(l_fch_gelu, l_ln2, l_fcw, l_fcb, B, T, C, 4*C, main_stream, l_fch, model->gelu_fusion);
matmul_forward_cublaslt(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C, main_stream);
// OK, fusion across blocks.
@ -856,7 +855,7 @@ void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int
floatX* l_ln2 = (model->recompute < 2) ? acts.ln2 + l * B * T * C : acts.lnf;
float* l_ln2_mean = acts.ln2_mean + l * B * T;
float* l_ln2_rstd = acts.ln2_rstd + l * B * T;
floatX* l_fch = acts.fch + l * B * T * 4*C;
floatX* l_fch_pre_gelu = acts.fch + l * B * T * 4*C;
floatX* l_fch_gelu = (model->recompute < 1) ? acts.fch_gelu + l * B * T * 4*C : acts.fch_gelu;
// get the pointers of the gradients of the activations for this layer
// notice that there is no l *, because we just have a single copy, and keep
@ -868,10 +867,9 @@ void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int
if(model->recompute >= 1) {
// recompute >= 1 means we recompute gelu. in this case,
// l_fch_gelu is just a buffer, so re-compute the gelu from l_fch here
gelu_forward(l_fch_gelu, l_fch, B*T*4*C, main_stream);
gelu_forward(l_fch_gelu, l_fch_pre_gelu, B*T*4*C, main_stream);
}
matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C, main_stream);
gelu_backward_inplace(dl_bt4c, l_fch, B*T*4*C, main_stream);
matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C, main_stream, l_fch_pre_gelu, model->gelu_fusion);
if(model->recompute >= 2) {
// same as gelu above, l_ln1 and l_ln2 are just buffers if recompute >= 2, recompute them here on demand
layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C, main_stream);
@ -888,7 +886,7 @@ void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int
floatX* l_att = acts.att + l * B * NH * T * T;
// we need B x T x (4)C buffers. l_atty and l_fch aren't needed anymore at this point, so reuse their memory
floatX* buffer_a = l_atty;
floatX* buffer_b = l_fch; // this is B x T x 4C, so even larger than what we need
floatX* buffer_b = l_fch_pre_gelu; // this is B x T x 4C, so even larger than what we need
attention_backward(dl_bt4c, buffer_b, scratchX, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH, main_stream);
#endif
if(model->recompute >= 2) {
@ -1172,13 +1170,11 @@ void common_start(bool override_enable_tf32 = true, bool print_device_info = tru
nvtxNameCudaStreamA(main_stream, "main stream");
// set up cuBLAS and cuBLASLt
cublasCheck(cublasCreate(&cublas_handle));
cublasCheck(cublasLtCreate(&cublaslt_handle));
cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size));
// TF32 precision is equivalent to torch.set_float32_matmul_precision('high')
bool enable_tf32 = PRECISION_MODE == PRECISION_FP32 && deviceProp.major >= 8 && override_enable_tf32;
cublasCheck(cublasSetMathMode(cublas_handle, enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH));
cublas_compute = enable_tf32 ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F;
#ifdef ENABLE_CUDNN
@ -1189,7 +1185,6 @@ void common_start(bool override_enable_tf32 = true, bool print_device_info = tru
void common_free(GPT2 &model) {
cudaCheck(cudaStreamDestroy(main_stream));
cudaCheck(cudaFree(cublaslt_workspace));
cublasCheck(cublasDestroy(cublas_handle));
cublasCheck(cublasLtDestroy(cublaslt_handle));
#ifdef ENABLE_CUDNN
destroy_cudnn();
@ -1391,6 +1386,7 @@ void error_usage() {
// numerics
fprintf(stderr, " -f <int> enable_tf32 override (default: 1, set to 0 to disable tf32)\n");
fprintf(stderr, " -w <int> keep f32 copy of weights for the optimizer? (default: 1)\n");
fprintf(stderr, " -ge <int> gelu fusion: 0=none, 1=forward, 2=forward+backward (default: 2 for >=SM90, 0 for older GPUs)\n");
// memory management
fprintf(stderr, " -z <int> zero_stage, Zero Optimization Stage, 0,1,2,3 (default = 0)\n");
fprintf(stderr, " -r <int> recompute: less memory but less speed. (default = 1), 0|1|2 = none,gelu,gelu+ln\n");
@ -1434,6 +1430,7 @@ int main(int argc, char *argv[]) {
int max_steps = -1;
int override_enable_tf32 = 1;
int use_master_weights = 1;
int gelu_fusion = -1; // 0 = none, 1 = forward, 2 = forward+backward (-1 => per-GPU default)
int recompute = 1; // recompute during backward setting, 0 = none, 1 = recompute gelu
int zero_stage = 0; // Zero Optimization Stage for Multi-GPU training
int hellaswag_eval = 0;
@ -1466,6 +1463,7 @@ int main(int argc, char *argv[]) {
else if (argv[i][1] == 'v') { val_loss_every = atoi(argv[i+1]); }
else if (argv[i][1] == 'm') { val_max_steps = atoi(argv[i+1]); }
else if (argv[i][1] == 's' && argv[i][2] == '\0') { sample_every = atoi(argv[i+1]); }
else if (argv[i][1] == 'g' && argv[i][2] == 'e') { gelu_fusion = atoi(argv[i+1]); }
else if (argv[i][1] == 'g') { genT = atoi(argv[i+1]); }
else if (argv[i][1] == 'a') { overfit_single_batch = atoi(argv[i+1]); }
else if (argv[i][1] == 'f') { override_enable_tf32 = atoi(argv[i+1]); }
@ -1486,7 +1484,9 @@ int main(int argc, char *argv[]) {
else if (argv[i][1] == 'n' && argv[i][2] == 'm') { major_checkpoint_every = atoi(argv[i+1]); }
else { error_usage(); }
}
multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, nccl_init_method);
common_start(override_enable_tf32, false); // common init code for train/test/profile
// should do a bit more error checking here
assert(warmup_iterations >= 0);
@ -1496,6 +1496,8 @@ int main(int argc, char *argv[]) {
int tokens_per_fwdbwd = B * T * multi_gpu_config.num_processes; // one micro-batch processes this many tokens
// calculate sensible default for total batch size as assuming no gradient accumulation
if (total_batch_size == -1) { total_batch_size = tokens_per_fwdbwd; }
// in the future, we might want to set gelu fusion to 2 for SM90+ and 0 for other GPUs
if (gelu_fusion == -1) { gelu_fusion = 0; } // (deviceProp.major >= 9) ? 2 : 0; } // in gpt2_init_common for test_gpt2cu...
// calculate the number of gradient accumulation steps from the desired total batch size
assert(total_batch_size % tokens_per_fwdbwd == 0);
int grad_accum_steps = total_batch_size / tokens_per_fwdbwd;
@ -1527,10 +1529,9 @@ int main(int argc, char *argv[]) {
printf0("| genT | %-50d |\n", genT);
printf0("| overfit_single_batch | %-50d |\n", overfit_single_batch);
printf0("| use_master_weights | %-50s |\n", use_master_weights ? "enabled" : "disabled");
printf0("| gelu_fusion | %-50d |\n", gelu_fusion);
printf0("| recompute | %-50d |\n", recompute);
printf0("+-----------------------+----------------------------------------------------+\n");
common_start(override_enable_tf32, false); // common init code for train/test/profile
const char* precision_str = (PRECISION_MODE == PRECISION_FP32)
? (cublas_compute == CUBLAS_COMPUTE_32F_FAST_TF32 ? "TF32" : "FP32")
: (PRECISION_MODE == PRECISION_FP16 ? "FP16" : "BF16");
@ -1573,6 +1574,7 @@ int main(int argc, char *argv[]) {
}
model.use_master_weights = use_master_weights;
model.gelu_fusion = gelu_fusion;
model.recompute = recompute;
printf0("| weight init method | %-50s |\n", resuming == 1 ? "intermediate checkpoint" : (load_filename[0] == 'd' ? "random" : "OpenAI's GPT-2 checkpoint"));
printf0("| max_sequence_length T | %-50d |\n", model.config.max_seq_len);