单GPU高效训练笔记

本文主要是为单GPU高效训练做的一些笔记和记录

cuda kernel

矩阵乘法示例

#include <stdio.h>
#include <cuda_runtime.h>

// 定义矩阵大小
#define N 1024

// CUDA 核函数
__global__ void MatrixMulKernel(float *d_A, float *d_B, float *d_C, int width) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < width && col < width) {
        float value = 0;
        for (int k = 0; k < width; ++k) {
            value += d_A[row * width + k] * d_B[k * width + col];
        }
        d_C[row * width + col] = value;
    }
}

int main() {
    // 在主机上分配内存
    float *h_A = (float *)malloc(N * N * sizeof(float));
    float *h_B = (float *)malloc(N * N * sizeof(float));
    float *h_C = (float *)malloc(N * N * sizeof(float));

    // 初始化矩阵
    for (int i = 0; i < N * N; i++) {
        h_A[i] = rand() / (float)RAND_MAX;
        h_B[i] = rand() / (float)RAND_MAX;
    }

    // 在设备上分配内存
    float *d_A, *d_B, *d_C;
    cudaMalloc((void **)&d_A, N * N * sizeof(float));
    cudaMalloc((void **)&d_B, N * N * sizeof(float));
    cudaMalloc((void **)&d_C, N * N * sizeof(float));

    // 将数据从主机复制到设备
    cudaMemcpy(d_A, h_A, N * N * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, N * N * sizeof(float), cudaMemcpyHostToDevice);

    // 定义线程块和网格大小
    dim3 threadsPerBlock(16, 16);
    dim3 blocksPerGrid((N + threadsPerBlock.x - 1) / threadsPerBlock.x,
                       (N + threadsPerBlock.y - 1) / threadsPerBlock.y);

    // 启动CUDA核函数
    MatrixMulKernel<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);

    // 将结果从设备复制回主机
    cudaMemcpy(h_C, d_C, N * N * sizeof(float), cudaMemcpyDeviceToHost);

    // 验证结果
    bool match = true;
    for (int i = 0; i < N && match; i++) {
        for (int j = 0; j < N && match; j++) {
            float sum = 0;
            for (int k = 0; k < N; k++) {
                sum += h_A[i * N + k] * h_B[k * N + j];
            }
            if (fabs(h_C[i * N + j] - sum) > 1e-4) {
                match = false;
                printf("结果不匹配：(%d, %d): GPU = %f, CPU = %f\n", i, j, h_C[i * N + j], sum);
            }
        }
    }

    if (match) {
        printf("结果匹配！\n");
    }

    // 释放内存
    free(h_A);
    free(h_B);
    free(h_C);
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);

    return 0;
}

代码说明
矩阵大小: 代码中矩阵大小为1024 x 1024。你可以通过修改N的值来改变矩阵的大小。

CUDA 核函数: MatrixMulKernel 是执行矩阵乘法的 CUDA 核函数。每个线程处理结果矩阵中的一个元素。

网格和线程块大小: 线程块大小设置为 16 x 16，网格大小根据矩阵的维度自动计算，以确保覆盖整个矩阵。

数据传输: 使用 cudaMemcpy 函数将数据从主机传输到设备，然后在计算完成后将结果从设备传输回主机。

结果验证: 在主机上使用 CPU 计算矩阵乘法的结果，并与 GPU 计算的结果进行比较，确保计算正确。

Flash Attention v1，v2

很经典

数据预加载

CPU->GPU，通过PCle

micro-batch

“Micro-batch” 是在深度学习和分布式计算中常用的术语，指的是将一个大的 Batch size 切分成更小的批次（即 micro-batches），然后分别处理这些更小的批次。这种策略主要用于优化模型训练过程，尤其是在使用梯度累积（Gradient Accumulation）或在分布式环境下进行模型训练时。

Micro-batch 的概念

1. 小批次处理：

   • 当你定义一个较大的 Batch size 时，计算资源（如 GPU 的显存）可能不足以一次性处理整个 Batch。这时，模型会将这个大 Batch 分成更小的批次，即 micro-batches，每个 micro-batch 包含的数据样本数量较少。
   • 模型会逐个计算每个 micro-batch 的梯度，累积所有 micro-batch 的梯度后，再进行一次整体的参数更新。

2. 梯度累积中的 micro-batch：

   • 使用梯度累积时，通常一个大的 Batch size 会被拆分成多个 micro-batch。例如，如果你希望使用一个实际的 Batch size 为 128，但由于显存限制每次只能处理 16 个样本，你可以将这个大的 Batch size 拆分为 8 个 micro-batch（每个 micro-batch 包含 16 个样本）。
   • 在处理每个 micro-batch 时，模型会计算对应的梯度，并将其累积起来。在所有 micro-batch 处理完后，再更新模型的参数。

3. 分布式计算中的 micro-batch：

   • 在分布式计算或并行训练中，micro-batch 也用于在多个计算节点或 GPU 之间分配任务。每个节点或 GPU 可能处理一个或多个 micro-batch，最终将所有节点的梯度累积并进行同步更新。

Micro-batch 的作用

• 资源优化：通过将大的 Batch size 拆分为多个 micro-batch，可以更有效地利用有限的计算资源（如显存），从而在资源受限的情况下依然可以训练大型模型。
• 训练稳定性：micro-batch 可以帮助平衡训练过程中的计算负担，同时在进行梯度累积时，确保即使在使用较小的 mini-batch 处理时，依然能够模拟大 Batch size 的效果，保持训练的稳定性和效率。
• 分布式处理：在分布式训练中，micro-batch 的概念有助于将任务划分给多个计算单元，使得大规模模型训练变得更加高效和可扩展。

总结

Micro-batch 是指在训练过程中将大的 Batch size 拆分成更小的批次进行处理的策略，广泛用于优化计算资源、实现梯度累积以及支持分布式训练。它帮助在处理能力有限的情况下依然能够进行高效且稳定的模型训练。

编译优化

TorchDynamo 是 PyTorch 中的一个动态编译和优化框架，旨在捕获和优化深度学习模型的计算图，从而提高模型的执行效率。TorchDynamo 的主要目标是为深度学习提供一种灵活且高效的方式，以便在不改变原有模型代码的情况下，自动实现优化和加速。

TorchDynamo 的工作原理

TorchDynamo 的核心思想是通过动态捕获 Python 代码中的计算图，并对这些计算图进行优化和编译。在深度学习中，计算图表示的是模型的前向和反向传播过程。通过优化这些计算图，TorchDynamo 能够提升模型的运行效率，特别是在 GPU 和其他加速硬件上。

主要步骤

1. 代码捕获：

   • TorchDynamo 插入在 Python 解释器中，当模型的前向传播代码运行时，它会捕获这个过程中的计算图。它不会改变原有代码的行为，而是对执行的操作进行跟踪。

2. 计算图转换：

   • 捕获到的计算图被转化为一种中间表示（IR），这种表示更适合进一步的优化和编译。TorchDynamo 可以在这个阶段对计算图进行分析和重写，以简化和优化操作。

3. 优化与编译：

   • 转换后的计算图会通过一系列优化过程，这可能包括内存优化、操作融合、常量折叠等。TorchDynamo 还可以将计算图传递给其他后端，如 nvFuser 或 TorchScript，进一步编译为高效的机器码。

4. 执行优化后的代码：

   • 最终，TorchDynamo 生成优化后的代码并执行这些代码。这些代码通常比原始的 Python 代码执行得更快，因为它们已经过专门的优化和编译处理。

TorchDynamo 的优势

• 透明性：

  • TorchDynamo 的设计使得用户无需对原有模型代码进行大幅修改。开发者可以在现有的 PyTorch 代码基础上直接应用 TorchDynamo，轻松获得性能提升。

• 兼容性：

  • TorchDynamo 兼容大部分 PyTorch 功能和模型，包括复杂的控制流、动态计算图等。它可以与其他 PyTorch 特性如 TorchScript 和 JIT 编译器无缝集成。

• 性能优化：

  • 通过捕获并优化计算图，TorchDynamo 能够显著提高模型的执行速度，特别是在大型神经网络和计算密集型任务上。

使用场景

• 加速模型推理：

  • TorchDynamo 可以用于优化和加速模型的推理过程，特别是在部署阶段，可以显著减少推理时间。

• 训练加速：

  • 在训练过程中，TorchDynamo 通过优化计算图，提高了梯度计算和参数更新的效率，从而加快训练速度。

• 复杂模型的优化：

  • 对于包含复杂控制流和动态行为的模型（如循环神经网络或条件分支网络），TorchDynamo 提供了一种在不修改代码的情况下提升性能的途径。

总结

TorchDynamo 是 PyTorch 中一个重要的动态编译和优化工具，它通过捕获、转换和优化计算图，为深度学习模型提供性能提升。其透明性和兼容性使得用户可以轻松地在现有代码中应用它，从而在不改变模型逻辑的前提下获得更好的执行效率。

dynamo.optimize("nvfuser") 是 PyTorch 的一种优化机制，主要用于提升深度学习模型在 GPU 上的计算效率。它结合了 PyTorch 的 TorchDynamo 和 NVIDIA 的 nvFuser 来实现这一目标。

背景和概念

1. TorchDynamo:

   • TorchDynamo 是 PyTorch 引入的一个框架，它通过将 Python 的原生代码转化为高效的中间表示（IR）来优化深度学习模型的执行。TorchDynamo 可以捕获并优化 Python 代码中的计算图，使得模型的执行效率更高。
   • 它是通过动态编译（just-in-time, JIT）技术来进行代码优化的，可以根据模型的具体运行情况自动进行优化。

2. nvFuser:

   • nvFuser 是 NVIDIA 为 PyTorch 提供的一个 GPU 加速后端，专门用于加速张量计算。它是一种动态图形编译器，能够在运行时将模型的计算图编译为高效的 GPU 内核。
   • nvFuser 通过融合多个张量操作（Tensor Fusion），减少内存带宽使用和计算开销，从而显著提高模型在 GPU 上的执行效率。

dynamo.optimize("nvfuser") 的作用

• 当你使用 dynamo.optimize("nvfuser") 时，PyTorch 会通过 TorchDynamo 捕获模型的计算图，然后使用 nvFuser 将这些计算图编译成优化的 GPU 内核。这样，模型的计算效率和执行速度都可以得到显著提升，尤其是在深度学习的推理和训练过程中。

• 优化过程:

  1. 捕获计算图: TorchDynamo 会拦截模型的前向传播代码，将其转化为中间表示。
  2. 计算图优化: nvFuser 会进一步优化这些中间表示，尝试融合尽可能多的张量操作，减少内存传输和冗余计算。
  3. 生成 GPU 内核: nvFuser 将优化后的计算图编译为高效的 CUDA 内核，在 GPU 上执行。

使用场景

• 该优化方法特别适合计算密集型的深度学习模型，如 CNN、RNN 和 Transformer 等，在这些模型中，大量的张量操作可以通过 nvFuser 进行优化。
• 适合需要高性能推理和训练的应用场景，比如实时计算、边缘设备上的模型推理、大规模的训练任务等。

优势

• 性能提升: 通过减少内存访问、优化张量计算的执行顺序，显著提高模型的执行速度。
• 透明性: 使用 dynamo.optimize("nvfuser") 不需要对模型代码进行大幅度修改，可以直接在现有模型上应用。
• 灵活性: TorchDynamo 的动态捕获和 nvFuser 的动态图形编译相结合，使得优化过程能够适应各种不同的模型结构。

总之，dynamo.optimize("nvfuser") 是 PyTorch 提供的一种先进的优化工具，能够在保持代码简洁和易用的同时，显著提升模型的 GPU 执行效率。