85 lines
4.1 KiB
Diff
85 lines
4.1 KiB
Diff
diff --git a/engine/ggml/src/ggml-cuda/ggml-cuda.cu b/engine/ggml/src/ggml-cuda/ggml-cuda.cu
|
|
index 3e11b456..aff788f5 100644
|
|
--- a/engine/ggml/src/ggml-cuda/ggml-cuda.cu
|
|
+++ b/engine/ggml/src/ggml-cuda/ggml-cuda.cu
|
|
@@ -5189,7 +5189,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
|
|
}
|
|
} break;
|
|
case GGML_OP_OUT_PROD:
|
|
- return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
|
|
+ // HOT-Step patch: BF16 out_prod — see engine/patches/bf16-out-prod.patch
|
|
+ return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
|
|
+ (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16);
|
|
case GGML_OP_GET_ROWS:
|
|
{
|
|
switch (op->src[0]->type) {
|
|
diff --git a/engine/ggml/src/ggml-cuda/out-prod.cu b/engine/ggml/src/ggml-cuda/out-prod.cu
|
|
index 499903d0..9b48c8b5 100644
|
|
--- a/engine/ggml/src/ggml-cuda/out-prod.cu
|
|
+++ b/engine/ggml/src/ggml-cuda/out-prod.cu
|
|
@@ -1,4 +1,5 @@
|
|
#include "out-prod.cuh"
|
|
+#include "convert.cuh" // HOT-Step patch: BF16 out_prod — see engine/patches/bf16-out-prod.patch
|
|
|
|
#include <cstdint>
|
|
|
|
@@ -8,7 +9,14 @@ void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
|
|
|
GGML_TENSOR_BINARY_OP_LOCALS
|
|
|
|
- GGML_ASSERT(src0->type == GGML_TYPE_F32);
|
|
+ // HOT-Step patch: BF16 out_prod — see engine/patches/bf16-out-prod.patch
|
|
+ // ggml computes the gradient w.r.t. a mul_mat's ACTIVATION input as
|
|
+ // out_prod(weight, transpose(grad)), so the frozen weight lands in src0. An
|
|
+ // F32-only assert here is what forces the DiT trainer to mirror every
|
|
+ // trainable-layer weight to F32. Accepting BF16 src0 (dequantized once into
|
|
+ // an F32 workspace below) halves that mirror. src1/dst stay F32 and the F32
|
|
+ // path is byte-identical.
|
|
+ GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16);
|
|
GGML_ASSERT(src1->type == GGML_TYPE_F32);
|
|
GGML_ASSERT(dst->type == GGML_TYPE_F32);
|
|
|
|
@@ -22,19 +30,37 @@ void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
|
GGML_ASSERT(ne2 == src1->ne[2]);
|
|
GGML_ASSERT(ne3 == src1->ne[3]);
|
|
|
|
- const float * src0_d = (const float *) src0->data;
|
|
- const float * src1_d = (const float *) src1->data;
|
|
- float * dst_d = (float *) dst->data;
|
|
-
|
|
cudaStream_t stream = ctx.stream();
|
|
cublasHandle_t handle = ctx.cublas_handle();
|
|
|
|
+ // HOT-Step patch: BF16 out_prod — see engine/patches/bf16-out-prod.patch
|
|
+ // Dequantize a BF16 src0 into a contiguous F32 workspace, so every GEMM path
|
|
+ // below is the shipped F32 one with lda == ne00. Restricted to a 2-D src0
|
|
+ // (the trainer's case: a frozen weight matrix), which also makes the dim-2/3
|
|
+ // workspace strides moot — with ne02 == ne03 == 1 the (i2/dps2) and (i3/dps3)
|
|
+ // src0 offsets below are identically zero.
|
|
+ ggml_cuda_pool_alloc<float> src0_f32(ctx.pool());
|
|
+ if (src0->type == GGML_TYPE_BF16) {
|
|
+ GGML_ASSERT(ne02 == 1 && ne03 == 1);
|
|
+ GGML_ASSERT(ggml_is_contiguous(src0));
|
|
+ const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_BF16);
|
|
+ GGML_ASSERT(to_fp32_cuda != nullptr);
|
|
+ src0_f32.alloc((size_t) ne00 * ne01);
|
|
+ to_fp32_cuda(src0->data, src0_f32.get(), ne00 * ne01, stream);
|
|
+ }
|
|
+
|
|
+ const float * src0_d = src0_f32.get() ? src0_f32.get() : (const float *) src0->data;
|
|
+ const float * src1_d = (const float *) src1->data;
|
|
+ float * dst_d = (float *) dst->data;
|
|
+
|
|
const float alpha = 1.0f;
|
|
const float beta = 0.0f;
|
|
|
|
CUBLAS_CHECK(cublasSetStream(handle, stream));
|
|
|
|
- const int64_t lda = nb01 / sizeof(float);
|
|
+ // HOT-Step patch: BF16 out_prod — the workspace is contiguous, so its leading
|
|
+ // dimension is ne00 rather than src0's own (BF16-sized) row stride.
|
|
+ const int64_t lda = src0_f32.get() ? ne00 : (int64_t) (nb01 / sizeof(float));
|
|
const int64_t ldc = nb1 / sizeof(float);
|
|
|
|
const bool src1_T = ggml_is_transposed(src1);
|