Possibly inconsistent gradient computation?

[yixuan]

Hi Andre,

This is a nice and timely work. I have a quick question regarding the implementation of the backward process. From

mHC.cu/src/csrc/kernels/mhc_layer.cuh

Lines 155 to 158 in 3595c86

	CHECK_CUDA(cudaMalloc(&sinkhorn_M, B * n * n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&H_pre_activated, B * n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&H_post_activated, B * n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&H_res_tilde, B * n * n * sizeof(float)));

I can find that the H matrices have dimensions B × n × n or B × n if we use the dynamic-H mode, which makes perfect sense. However, in the backward implementation

mHC.cu/src/csrc/kernels/mhc_layer.cuh

Lines 227 to 241 in 3595c86

	void init(int B, int C, int n) {
	CHECK_CUDA(cudaMalloc(&d_x_expanded, B * n * C * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_H_pre, n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_rmsnorm_weight, C * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_H_post, n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_H_res, n * n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_x_aggregated, B * C * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_layer_out, B * C * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_y_distributed, B * n * C * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_x_mixed, B * n * C * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_M, n * n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_H_pre_activated, n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_H_post_activated, n * sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_H_res_exp, n * n * sizeof(float)));

I find that both H and M lose the batch dimension B.

Does it mean that the backward implementation only works for static H? Thanks.

[AndreSlavescu]

Hi Yixuan,

Initially yes, the backward kernels are only for static H-version. However, in this latest PR (#3), I have added kernels for dynamic H-mode and worked on some performance improvements.