CUDA를 넘어: DeepSeek

 

https://dev.to/datamonk_/how-deepseek-is-making-high-performance-ai-accessible-to-all-26fp

How DeepSeek is Making High-Performance AI Accessible to All

Low-Level에 대해서는 잘 모르지만

재미있는 기사가 있어 살펴보고자 한다.

 

DeepSeek가 LLM 및 AI를 강타하며, 이젠 비 전문가도 한 번 쯤 들어본 아키텍처로 떠오르고 있다.

그러나 사실 DeepSeek의 강렬한 Contribusion은 이런 영역이 아니라고 보는 것이 타당하다.

 

이전에 리뷰했던 아래의 영역은 사실 딱히 중요한 것이 아니다.

1. MLA와 같은 Attention 구조
2. Shared Expert를 이용한 MOE의 구조
3. Self-Evolution (발칙한 접근이긴 하지만, 가장 Wow한 부분은 아니라고 본다)

 

그럼 뭐가 가장 중요한가 ?
도대체 H-800 GPU로 어떻게 저런 '효율성'을 달성하였냐의 문제이다.

---

Beyond CUDA

AI를 다루는 사람은 결국 GPU를 다루는 사람이고
GPU를 다루는 방법은 Low-Level의 Assembly가 아니라, Python과 Pytorch를 기반으로 한다.

CUDA 수준까지 내려가는 것도 결코 쉬운 일은 아니다.

 

그러나 GPU 연산을 수행하는 이러한 Language들이  과연 '최고의 효율'을 보장하냐의 문제이다.

어떻게 DeepSeek는 저런 신비한 혼합 정밀도의 학습 Architecture를 구성하였는가 ...

 

DeepSeek는 원래도 그런 회사였지만

GPU 규제가 강화되어서인지... 최고의 효율을 위한 Low-Level의 단계로 진입한 것으로 보인다.

Pytorch를 벗어나 CUDA로, 거기서도 부족하였는지 이를 PTX 기반 언어를 통해 커스텀을 수행한 것으로 보인다.

GH100의 구조(144개의 SM). H100 SXM5의 경우 132개의 SM을, PCle 버전의 경우 114개의 SM을 보유하고 있다고 한다.

 

이건 DeepSeek가 사용했다는 H800과 똑같은 구조를 가진 H100 Processor의 모습이다.

(알고 있는 것은, SM(Streaming Multiprocessor)가 일종의 기능적 단위로, GPU의 성능적 지표라는 것 뿐 ...)

하나의 Streaming Multiprocessor는 위와 같은 모습을 띄고 있고

거기에 들어 있는 개별적인 Tensor Core는 위와 같은 TeraFLOPs의 성능을 가지고 있다.

 

아주 간단한 성능표를 본다고 가정한다면

우리는 당연하게도 어디서 '병목(Bottleneck)'이 존재할지 생각하고
어떻게 '자원'을 분배하는 것이 효율적인지 생각하게 될 것이다.

 

하지만 AI 학습을 위한 분산환경 조성과, 다중 PC의 연결을 구성하면서도 그런것을 고려하기는 ...?

너무나도 어렵기에 대부분은 그냥 Pytorch에 의존하게 되는 것이다.

 

DeepSeek의 차별점은 바로 이 지점에 존재한다.

모델의 구조, 파라미터 모든 것을 공개하였지만 유일하게 공개하지 않은 것.

바로 도대체 어떻게 CUDA가 아닌 다른 방식을 통해 학습을 진행하였냐는 것이다.

 

도대체 그 Customized PTX는 어떻게 작성된 것일까.

https://docs.nvidia.com/cuda/inline-ptx-assembly/contents.html

Contents — Inline PTX Assembly in CUDA 12.8 documentation

여기서 PTX는 거의 Assembly 수준의 Low Level 언어에 가깝다고 한다.

굳이 순서를 따라가보면 ...

 

CUDA -> PTX(Virtual Assembly Level) -> GPU Driver Level Compier -> GPU 이런 느낌이라고 한다.

즉, GPU는 GPU의 개별적 Driver Level의 ISA는 공개되지 않았으며 GPU 종속적인 성격을 띈다고 한다.
에당초 공개하지도 않지만, 표준화된 ISA도 존재하지 않는다는 것이다.

반면 해당 ISA가 사용하는 PTX instruction은 이러한 Device에 독립적으로 존재한다.

DeepSeek는 바로 이 영역을 건드리며 GPU의 최적화에 나선다.

---

실제 해당 영역은 다양한 기법들을 논문에서 제시하고 있다만 ...

CUDA 레벨의 프로그래밍조차 C와 같은 Low Level의 언어이기 때문에 사실상 다루는 것은 매우 어렵다.

(거기에 Assembly 수준의 Low Level 언어는 ...)

 

1. 2048개의 NVIDIA H800을 사용 (개별 Node는 8개의 GPU를 가지고 있음)
   NVLink와 NVSwitch의 통신속도를 최적화 하기 위해 여러 기법을 사용하였음.
   (분산학습에서 이러한 통신 자체가 일종의 병목을 야기한다고 보았는듯)
2. Computation - Communication의 Overlap.
3. DualPipeline의 구성
4. 20개의 Streaming Multiprocessors를 Communication 전용으로 두어 통신에 할당
5. 커스텀화 된 혼합 정밀도 학습 루프 설계
   - FP32 텐서를 Scaling Factor를 곱해 FP8로 변환
   - GEMM(8FP General Matrix Multiplication), Bias, RMSNorm 적용
   - 다시 FP32로 복원
6. 등등 ...

// fp8_mixed_precision_training.cu
#include <iostream>
#include <cmath>
#include <cstdlib>
#include <cuda_runtime.h>
#include <cublas_v2.h>

// 오류 체크 매크로
#define CUDA_CHECK(err) do { \
    cudaError_t err_ = (err); \
    if (err_ != cudaSuccess) { \
        std::cerr << "CUDA Error: " << cudaGetErrorString(err_) \
                  << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
        exit(EXIT_FAILURE); \
    } \
} while(0)

#define CUBLAS_CHECK(err) do { \
    cublasStatus_t err_ = (err); \
    if (err_ != CUBLAS_STATUS_SUCCESS) { \
        std::cerr << "CUBLAS Error at " << __FILE__ << ":" << __LINE__ << std::endl; \
        exit(EXIT_FAILURE); \
    } \
} while(0)

// FP8 타입 정의 (실제 FP8 연산은 최신 GPU에서 지원하지만, 여기서는 char로 시뮬레이션)
typedef char fp8_t;
const float FP8_MAX = 127.0f;  // FP8에서 표현 가능한 최대 절대값

// --------------------------------------------------------------------------
// [1] FP32 -> FP8 quantization / FP8 -> FP32 dequantization 커널
// --------------------------------------------------------------------------

// FP32 배열을 주어진 scaling factor(scale)를 곱해 FP8로 quantize (각 값은 [-127,127]로 clamp)
__global__ void quantizeKernel(const float* input, fp8_t* output, float scale, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx < n) {
        float scaled = input[idx] * scale;  // scale = FP8_MAX / (max_abs)
        int q = (int)roundf(scaled);
        if(q > 127) q = 127;
        if(q < -127) q = -127;
        output[idx] = (fp8_t)q;
    }
}

// FP8 배열을 dequantize하여 FP32 배열로 변환 (invScale = 1/scale)
__global__ void dequantizeKernel(const fp8_t* input, float* output, float invScale, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx < n) {
        int q = (int)input[idx];
        output[idx] = ((float)q) * invScale;
    }
}

// --------------------------------------------------------------------------
// [2] RMSNorm Forward 커널 (y = x / sqrt(mean(x^2)+epsilon))
//    – 실제로는 저장하지 않고, 역전파 시 재계산(recomputation)을 수행하여 메모리 절약 효과를 얻습니다.
//    (여기서는 데모용으로 순전파에서도 RMSNorm 결과를 출력합니다.)
// --------------------------------------------------------------------------
__global__ void rmsNormForwardKernel(const float* x, float* y, int n, float epsilon) {
    // 각 블록이 한 샘플을 처리한다고 가정하고, 블록 내 쓰레드가 협력해 reduction 수행 (간단히 구현)
    extern __shared__ float sdata[];
    int tid = threadIdx.x;
    int idx = blockIdx.x * n + tid;
    float val = (tid < n) ? x[idx] : 0.0f;
    sdata[tid] = val * val;
    __syncthreads();
    // reduction: (단순 구현 – n이 쓰레드 수와 같다고 가정)
    for (int s = n/2; s > 0; s >>= 1) {
        if(tid < s && tid + s < n) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    float rms = sqrtf(sdata[0] / n + epsilon);
    if(tid < n) {
        y[idx] = x[idx] / rms;
    }
}

// --------------------------------------------------------------------------
// [3] FP8 GEMM 함수
//    – 입력 A와 가중치 B는 FP8 형식으로 저장되어 있으며, 각각 scaling factor scaleA, scaleB를 갖습니다.
//    – 내부에서는 임시 FP32 버퍼로 dequantize한 후, cuBLAS의 Sgemm을 통해 FP32 누적으로 GEMM을 수행합니다.
//    – (참고: 실제 Nvidia FTX 구현에서는 Tensor Core와 FP8 전용 커널을 사용하며, 일정 간격마다 CUDA Core로 promotion을 수행합니다.)
// --------------------------------------------------------------------------
void fp8GEMM(cublasHandle_t handle, int M, int N, int K,
             const fp8_t* d_A, float scaleA,
             const fp8_t* d_B, float scaleB,
             float* d_C) {  // 결과는 FP32
    float *d_A_fp32, *d_B_fp32;
    CUDA_CHECK(cudaMalloc(&d_A_fp32, M * K * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&d_B_fp32, K * N * sizeof(float)));

    int total_A = M * K;
    int total_B = K * N;
    int blockSize = 256;
    int numBlocksA = (total_A + blockSize - 1) / blockSize;
    int numBlocksB = (total_B + blockSize - 1) / blockSize;
    float invScaleA = 1.0f / scaleA;
    float invScaleB = 1.0f / scaleB;
    dequantizeKernel<<<numBlocksA, blockSize>>>(d_A, d_A_fp32, invScaleA, total_A);
    CUDA_CHECK(cudaGetLastError());
    dequantizeKernel<<<numBlocksB, blockSize>>>(d_B, d_B_fp32, invScaleB, total_B);
    CUDA_CHECK(cudaGetLastError());

    float alpha = 1.0f, beta = 0.0f;
    // cuBLAS는 기본적으로 열우선 저장을 가정하므로, 여기서는 단순화를 위해 행렬 크기를 그대로 사용합니다.
    CUBLAS_CHECK(cublasSgemm(handle,
                             CUBLAS_OP_N, CUBLAS_OP_N,
                             N, M, K,
                             &alpha,
                             d_B_fp32, N,
                             d_A_fp32, K,
                             &beta,
                             d_C, N));
    cudaFree(d_A_fp32);
    cudaFree(d_B_fp32);
}

// --------------------------------------------------------------------------
// [4] 단순한 FP32 행렬에 bias를 더하는 커널
// --------------------------------------------------------------------------
__global__ void addBiasKernel_FP32(float* mat, const float* bias, int cols, int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        int col = idx % cols;
        mat[idx] += bias[col];
    }
}

// --------------------------------------------------------------------------
// [5] MSE Loss 계산 커널 (단순 데모용)
// --------------------------------------------------------------------------
__global__ void mseLossKernel_FP32(const float* pred, const float* target, float* loss, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx < n) {
        float diff = pred[idx] - target[idx];
        atomicAdd(loss, diff * diff);
    }
}

// --------------------------------------------------------------------------
// [6] CPU에서 EMA(Exponential Moving Average) 업데이트 (동기적으로 구현)
// --------------------------------------------------------------------------
void updateEMA(const float* param, float* ema_param, int size, float decay) {
    for (int i = 0; i < size; i++) {
        ema_param[i] = decay * ema_param[i] + (1.0f - decay) * param[i];
    }
}

// --------------------------------------------------------------------------
// [7] 메인 학습 코드
// --------------------------------------------------------------------------
int main() {
    // 하이퍼파라미터 및 크기 설정
    const int batch_size   = 64;
    const int input_dim    = 512;  // 예: 임베딩 차원
    const int output_dim   = 256;  // 예: 출력 차원 (출력 헤드)
    const float learning_rate = 0.001f;
    const int epochs       = 5;
    const int num_batches  = 10;

    // cuBLAS 핸들 생성
    cublasHandle_t handle;
    CUBLAS_CHECK(cublasCreate(&handle));

    // ----------------------------------------------------------------------
    // [A] 호스트 메모리: 입력 및 타겟 (데모용으로 매 배치 동일한 데이터 사용)
    // ----------------------------------------------------------------------
    float *h_input  = new float[batch_size * input_dim];
    float *h_target = new float[batch_size * output_dim];
    for (int i = 0; i < batch_size * input_dim; i++) {
        h_input[i] = static_cast<float>(rand()) / RAND_MAX;
    }
    for (int i = 0; i < batch_size * output_dim; i++) {
        h_target[i] = static_cast<float>(rand()) / RAND_MAX;
    }

    // ----------------------------------------------------------------------
    // [B] Device 메모리 할당: 입력 (FP32)
    // ----------------------------------------------------------------------
    float *d_input;
    CUDA_CHECK(cudaMalloc(&d_input, batch_size * input_dim * sizeof(float)));
    CUDA_CHECK(cudaMemcpy(d_input, h_input, batch_size * input_dim * sizeof(float), cudaMemcpyHostToDevice));

    // ----------------------------------------------------------------------
    // [C] 학습할 선형 계층의 가중치를 FP8 형식으로 저장 (중간 layer GEMM)
    //      (임베딩 및 출력 헤드는 FP32로 유지할 수 있으나, 여기서는 단순화를 위해 bias만 FP32)
    // ----------------------------------------------------------------------
    fp8_t *d_W_fp8;
    CUDA_CHECK(cudaMalloc(&d_W_fp8, input_dim * output_dim * sizeof(fp8_t)));
    // 먼저 호스트에서 FP32 master weight를 초기화한 후 quantization하여 FP8로 변환
    float *h_W = new float[input_dim * output_dim];
    for (int i = 0; i < input_dim * output_dim; i++) {
        h_W[i] = ((float)rand() / RAND_MAX - 0.5f) * 0.1f;
    }
    // 전역 scaling factor 계산 (실제 구현에서는 per‑tile scaling 적용)
    float max_val = 0.0f;
    for (int i = 0; i < input_dim * output_dim; i++) {
        float abs_val = fabsf(h_W[i]);
        if (abs_val > max_val) max_val = abs_val;
    }
    float scale_W = FP8_MAX / max_val;  // weight quantization scale
    // FP8 quantization (호스트에서)
    fp8_t *h_W_fp8 = new fp8_t[input_dim * output_dim];
    for (int i = 0; i < input_dim * output_dim; i++) {
        int q = (int)roundf(h_W[i] * scale_W);
        if(q > 127) q = 127;
        if(q < -127) q = -127;
        h_W_fp8[i] = (fp8_t)q;
    }
    CUDA_CHECK(cudaMemcpy(d_W_fp8, h_W_fp8, input_dim * output_dim * sizeof(fp8_t), cudaMemcpyHostToDevice));

    // ----------------------------------------------------------------------
    // [D] Bias (FP32) 할당
    // ----------------------------------------------------------------------
    float *d_bias;
    CUDA_CHECK(cudaMalloc(&d_bias, output_dim * sizeof(float)));
    float *h_bias = new float[output_dim];
    for (int i = 0; i < output_dim; i++) {
        h_bias[i] = 0.0f;
    }
    CUDA_CHECK(cudaMemcpy(d_bias, h_bias, output_dim * sizeof(float), cudaMemcpyHostToDevice));

    // ----------------------------------------------------------------------
    // [E] EMA 파라미터 (CPU): master weight와 bias의 EMA 값을 저장
    // ----------------------------------------------------------------------
    float *ema_W = new float[input_dim * output_dim];
    float *ema_bias = new float[output_dim];
    // EMA 초기값은 현재 FP32 master weight와 bias (weight는 dequantize된 값)
    for (int i = 0; i < input_dim * output_dim; i++) {
        ema_W[i] = ((float)h_W_fp8[i]) / scale_W;
    }
    for (int i = 0; i < output_dim; i++) {
        ema_bias[i] = h_bias[i];
    }
    float ema_decay = 0.999f;

    // ----------------------------------------------------------------------
    // [F] 학습 루프
    // ----------------------------------------------------------------------
    for (int epoch = 0; epoch < epochs; epoch++) {
        float epoch_loss = 0.0f;
        for (int batch = 0; batch < num_batches; batch++) {
            // (데모를 위해 매 배치마다 동일한 입력 사용)
            // 입력 FP32를 device에 이미 복사해두었으므로, FP8 quantization 진행

            // [F-1] 입력 quantization: FP32 -> FP8
            fp8_t *d_input_fp8;
            CUDA_CHECK(cudaMalloc(&d_input_fp8, batch_size * input_dim * sizeof(fp8_t)));
            // 여기서는 입력에 대해 전역 scaling factor를 호스트에서 계산 (실제에서는 tile‑wise 적용)
            float max_input = 0.0f;
            for (int i = 0; i < batch_size * input_dim; i++) {
                float v = fabsf(h_input[i]);
                if(v > max_input) max_input = v;
            }
            float scale_input = FP8_MAX / max_input;
            int total_input = batch_size * input_dim;
            int blockSize = 256;
            int numBlocks = (total_input + blockSize - 1) / blockSize;
            quantizeKernel<<<numBlocks, blockSize>>>(d_input, d_input_fp8, scale_input, total_input);
            CUDA_CHECK(cudaGetLastError());

            // [F-2] FP8 GEMM: d_hidden = d_input_fp8 (shape: batch_size x input_dim) * d_W_fp8 (input_dim x output_dim)
            // 결과 d_hidden는 FP32 (batch_size x output_dim)
            float *d_hidden;
            CUDA_CHECK(cudaMalloc(&d_hidden, batch_size * output_dim * sizeof(float)));
            fp8GEMM(handle, batch_size, output_dim, input_dim, d_input_fp8, scale_input, d_W_fp8, scale_W, d_hidden);

            // [F-3] bias 추가: d_hidden += bias
            int total_hidden = batch_size * output_dim;
            numBlocks = (total_hidden + blockSize - 1) / blockSize;
            addBiasKernel_FP32<<<numBlocks, blockSize>>>(d_hidden, d_bias, output_dim, total_hidden);
            CUDA_CHECK(cudaGetLastError());

            // [F-4] RMSNorm 연산 (메모리 절약을 위해 순전파에서는 결과를 저장하지 않고, 역전파 시 재계산하도록 함)
            // 여기서는 데모를 위해 RMSNorm 결과를 d_rmsnorm_out에 저장합니다.
            float *d_rmsnorm_out;
            CUDA_CHECK(cudaMalloc(&d_rmsnorm_out, batch_size * output_dim * sizeof(float)));
            // 각 샘플마다 한 블록, blockDim = output_dim, 동적 shared memory 사용
            for (int i = 0; i < batch_size; i++) {
                rmsNormForwardKernel<<<1, output_dim, output_dim * sizeof(float)>>>(d_hidden + i * output_dim,
                                                                                    d_rmsnorm_out + i * output_dim,
                                                                                    output_dim, 1e-5f);
            }
            CUDA_CHECK(cudaDeviceSynchronize());

            // [F-5] Loss 계산 (예: MSE Loss between RMSNorm output and target)
            float *d_loss;
            CUDA_CHECK(cudaMalloc(&d_loss, sizeof(float)));
            CUDA_CHECK(cudaMemset(d_loss, 0, sizeof(float)));
            int total_loss = batch_size * output_dim;
            numBlocks = (total_loss + blockSize - 1) / blockSize;
            mseLossKernel_FP32<<<numBlocks, blockSize>>>(d_rmsnorm_out, h_target, d_loss, total_loss);
            CUDA_CHECK(cudaDeviceSynchronize());
            float batch_loss;
            CUDA_CHECK(cudaMemcpy(&batch_loss, d_loss, sizeof(float), cudaMemcpyDeviceToHost));
            batch_loss /= (batch_size * output_dim);
            epoch_loss += batch_loss;

            // [F-6] Backward pass
            // → 실제 구현에서는 FP8/FP32 혼합 정밀도에 따라 gradient를 계산하고,
            //    RMSNorm 등 recomputation 기법을 적용합니다.
            //    (여기서는 데모이므로 상세 gradient 계산은 생략합니다.)

            // [F-7] Parameter update (dummy update: FP32 master weight h_W 업데이트 후, 재-quantization)
            for (int i = 0; i < input_dim * output_dim; i++) {
                h_W[i] -= learning_rate * 0.001f;  // dummy gradient
            }
            // weight scaling 재계산 및 재-quantization
            max_val = 0.0f;
            for (int i = 0; i < input_dim * output_dim; i++) {
                float abs_val = fabsf(h_W[i]);
                if (abs_val > max_val) max_val = abs_val;
            }
            scale_W = FP8_MAX / max_val;
            for (int i = 0; i < input_dim * output_dim; i++) {
                int q = (int)roundf(h_W[i] * scale_W);
                if(q > 127) q = 127;
                if(q < -127) q = -127;
                h_W_fp8[i] = (fp8_t)q;
            }
            CUDA_CHECK(cudaMemcpy(d_W_fp8, h_W_fp8, input_dim * output_dim * sizeof(fp8_t), cudaMemcpyHostToDevice));

            // [F-8] EMA 업데이트 (CPU에서 수행; 실제로는 별도 스레드로 비동기 업데이트 가능)
            updateEMA(h_W, ema_W, input_dim * output_dim, ema_decay);
            updateEMA(h_bias, ema_bias, output_dim, ema_decay);

            // [F-9] 임시 버퍼 해제
            cudaFree(d_input_fp8);
            cudaFree(d_hidden);
            cudaFree(d_rmsnorm_out);
            cudaFree(d_loss);
        }
        std::cout << "Epoch " << epoch << " Loss: " << (epoch_loss / num_batches) << std::endl;
    }

    // ----------------------------------------------------------------------
    // [G] 정리: cuBLAS 핸들 및 할당된 메모리 해제
    // ----------------------------------------------------------------------
    CUBLAS_CHECK(cublasDestroy(handle));
    cudaFree(d_input);
    cudaFree(d_W_fp8);
    cudaFree(d_bias);
    delete[] h_input;
    delete[] h_target;
    delete[] h_W;
    delete[] h_W_fp8;
    delete[] h_bias;
    delete[] ema_W;
    delete[] ema_bias;

    return 0;
}

 

C/C# 을 몰라 CUDA를 다루는 코드에 예시를 작성해달라고 요청한 것이다.

사실 아직 PTX ISA는 적용되지도 않은 CUDA에서의 코드임에도 불구하고
거의 모든 것들을 직접 정의 및 할당해야함을 알 수 있다.

(단순하게 Pytorch로 작성했으면, 레이어 몇 개, 자료형 정도만 정의하고 끝이였을 코드인데 ...)

 

DeepSeek는 (정확하게 공개되지는 않았지만) 이러한 Low Level 최적화를 극한으로 끌어올린 것으로 보인다.

미래에셋의 디지털리서치 AI Weekly #45에서의 DeepSeek에 대한 평가

 

실제로는 논문을 다 뜯어보며, 개별 과정에서 무슨 생각들을 가지고 있었는지 살펴보고 싶었지만 ...

다음에 기회가 되면 시도를 해보고자 한다.

---

CUDA moat의 종말은 아닐 것이다.

 

이미 이러한 내용도 많이 언급되는 내용이 되고 있는 것 같다.

심지어 누군가는 CUDA를 벗어나 PTX의 가능성을 제시했기 때문에 Nvidia 주가가 떨어졌다고 말하는 사람들도 있다
(둘 다 결국 Nvidia GPU를 사용하기 위한 언어일 뿐인데 ...)

 

아마 DeepSeek에서도 모든 작업을 PTX 수준에서 작성하여 접근하지는 않았을 것이다.

CUDA로 작성된 코드 기반을 가지고, PTX ISA를 통해 몇 가지 최적화를 이루지 않았을까 싶다.