What is the best way to support 5d tensor (or even higher dimension)?

[randxie]

I am using ggml to implement a model that requires 5d tensor. When I read the ggml.h and ggml.c, it seems that there's no generic permute / transpose that can support arbitrary dimension tensor. What is the suggested way to deal with 5d tensor in ggml?

[slaren]

I am not sure that there is a good way to add a 5th dimension to the tensors. The reality is that lots of the ops assume 4 dimensions, and changing GGML_MAX_DIMS is not really possible without a significant refactoring. You can try to figure a way to merge two of the dimensions into the same one.

[randxie]

Hi @slaren, thanks for the insight! I had the same conclusion by reading the code. Is supporting arbitrary dimension on the roadmap? I understand that have a fixed shape is easy to reason about, but not as flexible if we want to support more model architecture.

FYI, the model I am experimenting with is nougat.

[slaren]

I don't think there are any specific plans to add support for more dimensions. Most operations only work with one to three dimensions, and the rest is treated as a batch, so you should be able to get the same results if you merge the last two dimensions into one.

[balisujohn]

@randxie what are the operations you need to perform on a 5d tensor? Maybe they could be factored into up to 4d operations on slices of a 5d tensor? Otherwise, I think ggml is probably not a viable choice because it seems really locked into its upper tensor dimension limit, probably would need to be majorly rewritten to accommodate arbitrary tensor dimensionality.

[balisujohn]

@slaren what are your thoughts on changing tensor max dims to a absurdly huge fixed number like 256, then leaving the ops current dimensionality requirements alone. I could be mistaken but I think this would give people the ability to extend specific ops functionality to meet their dimensionality needs? What would the refactoring cost be for a one time refactor to another fixed tensor dimension limit?

[slaren]

Ultimately adding more dimensions is not necessary unless there is an operator that requires them. Any additional dimensions are just treated as a batch size, and can always be consolidated into a single dimension. If there is an operation that needs more dimensions, I think it would be better to study that in a case by case basis. Adding a bunch of dimensions that will not work with most ops and will massively increase the size of the ggml_tensor struct is not really an option.

[delldu]

Dear Development Team,

We are in need of an extension to support 5D and 6D tensors for 3D applications, particularly for operations such as nn.Conv3d and nn.BatchNorm3d within PyTorch. The current support for 3D and 4D tensors is insufficient for advanced use cases. We kindly request that you consider implementing this feature to enhance the capabilities of our applications.

Thank you for your attention to this matter.

Best regards,

[mtavenrath]

Having 5D support at least for reshape + permute operations would be useful as well. e.g. take this snippet from apply_rope in stablediffusion.cpp

x              = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));       // [N, n_head, L, d_head]
x              = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N);  // [N * n_head, L, d_head/2, 2]
x              = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2));       // [2, N * n_head, L, d_head/2]

With 5 dimensions one could first reshape to [N, L, h_head, L, d_head/2, 2] and then apply both permutations with a single operation and another basically free reshape in the end. This would save one of the 2 memcpy operations executed in the underlying (Vulkan) backend where memcpy is the 2nd biggest bottleneck for the flux network with 40% of the total kernel time.