Is DeepGEMM directly applicable to backward in training?

[YouJiacheng]

bwd of GEMM is two GEMMs, but I wonder if I need to take some special care of the range of gradients?

[YouJiacheng]

Oh, so we must write a quantization kernel that produces the correct lhs_scales and rhs_scales.

[zhipeng93]

We need a fp8 gemm with 128x1 LHS scaling and 1x128 RHS scaling?

[LyricZhao]

We provide this library mainly for inference. So this library only supports DGRAD, not WGRAD.

In my understanding, WGRAD support needs more than a GEMM kernel, but also some utility fused kernels (e.g. transposing, fused with casting, fused with SwiGLU, fused with MoE layout). We want this library to be clean, so we didn't open-source them.

We may later release the WGRAD kernel, we will discuss about it internally :)

[YouJiacheng]

Thank you Chenggang!
Yup I forgot that WGRAD needs to transpose matrices.

[thefacetakt]

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?

Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

[soundOfDestiny]

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?

Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

[TheTinyTeddy]

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

[soundOfDestiny]

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

[TheTinyTeddy]

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy?

However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

[soundOfDestiny]

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy?

However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

P.S. Usually we don't call the forward pass in training 'inference'.

[TheTinyTeddy]

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy?
However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

P.S. Usually we don't call the forward pass in training 'inference'.

My mistake, the terminology should indeed be the forward pass for training. So would you be able to let me know if my previous message is correct regarding the per channel quantization for activations in the Wgrad computation?

[soundOfDestiny]

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy?
However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

P.S. Usually we don't call the forward pass in training 'inference'.

My mistake, the terminology should indeed be the forward pass for training. So would you be able to let me know if my previous message is correct regarding the per channel quantization for activations in the Wgrad computation?

As you mentioned, we don't save x in bf16 format. (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

[NiuMa-1234]

I found the fp8 gemm for Wgrad in BWD requires the shape of input activation to be the multiple of 128. But in moe layer it couldn't always be this way. Is there any necessity to restrain the number of tokens dispatched to each expert?

[qijiaxing]

I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation.
https://github.com/qijiaxing/DeepGEMM/tree/1d1d
any feedback or suggestions are welcomed!