ggml : fix padding in timestep embedding kernels

[danbev]

This commit applies the same changes as were applied in #15917 to the cuda, metal, opencl, sycl, and vulkan backends.

Refs: #15917

[ggerganov]

I think we should add a comment to the declaration of this function to clarify that it will pad the tensor if it has odd number of dimensions. So here:

llama.cpp/ggml/include/ggml.h

Lines 2123 to 2130 in 28b5f19

	// Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151
	// timesteps: [N,]
	// return: [N, dim]
	GGML_API struct ggml_tensor * ggml_timestep_embedding(
	struct ggml_context * ctx,
	struct ggml_tensor * timesteps,
	int dim,
	int max_period);

Change to:

// return: [N, PAD(dim, 2)]

Btw, not sure what was the reasoning for this padding logic. It does not seem to be present in the reference Python implementation. Maybe @leejet can clarify?

[danbev]

I forgot to add [no ci] so I manually cancelled the workflows, hence the "Some checks haven't completed yet".

[leejet]

I think we should add a comment to the declaration of this function to clarify that it will pad the tensor if it has odd number of dimensions. So here:

llama.cpp/ggml/include/ggml.h

Lines 2123 to 2130 in 28b5f19

	// Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151
	// timesteps: [N,]
	// return: [N, dim]
	GGML_API struct ggml_tensor * ggml_timestep_embedding(
	struct ggml_context * ctx,
	struct ggml_tensor * timesteps,
	int dim,
	int max_period);

Change to:

// return: [N, PAD(dim, 2)]

Btw, not sure what was the reasoning for this padding logic. It does not seem to be present in the reference Python implementation. Maybe @leejet can clarify?

It seems that padding is incorrect behavior. I don’t really remember why I added it in the first place—it was too long ago. Given the current situation, I suggest removing the padding.

[danbev]

@leejet Thanks for the feedback!
I'll take a closer look at removing the padding next week.

[danbev]

@ggerganov @leejet I've taken a closer look at this and I think that the padding might be required after all.

Looking at the reference python implementation we have the following:

def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):    
    if not repeat_only:                                                            
        half = dim // 2                                                         
        freqs = torch.exp(                                                      
            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
        ).to(device=timesteps.device)                                           
        args = timesteps[:, None].float() * freqs[None]                         
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)       
        if dim % 2:                                                             
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
    else:                                                                       
        embedding = repeat(timesteps, 'b -> b d', d=dim)                        
    return embedding                                                            

Specifically these lines:

        if dim % 2:                                                             
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)

if dim % 2 checks if the dimension is odd, and if so it adds an extra column of zeros to make the final embedding have the correct dimension.

This padding is necessary because sinusoidal embeddings produce pairs of cosine/sine values. For dim=15, you get 7 complete pairs (14 values) plus 1 padding zero to reach the target dimension. Without this padding, the GGML implementation leaves uninitialized memory that can contain garbage values, which is what we observed in the test failures. I believe the padding in GGML is required to match this reference behavior.

[ggerganov]

@danbev I think the problem is that for input dim=15 the Python code produces output dimension of 15, while ggml produces 16. At least this is what I understand from reading the implementation in ggml.c:

llama.cpp/ggml/src/ggml.c

Lines 4919 to 4932 in 261e6a2

	// ggml_timestep_embedding
	struct ggml_tensor * ggml_timestep_embedding(
	struct ggml_context * ctx,
	struct ggml_tensor * timesteps,
	int dim,
	int max_period) {
	int actual_dim = dim;
	if (dim % 2 != 0) {
	actual_dim = dim + 1;
	}
	struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, actual_dim, timesteps->ne[0]);

[danbev]

@ggerganov Oh right, I misunderstood the comment about removing the padding to mean in the kernels. It indeed looks like adding an extra dimension here would not be needed. I'll update this and test it.