ggml : implement op fusion, starting with REGLU/GEGLU/SWIGLU

[CISC]

Implement REGLU/GEGLU/SWIGLU ops to avoid unnecessary tensor duplications and a little more efficient execution by combining ops in one.

Implement op fusion, starting with REGLU/GEGLU/SWIGLU for PoC.

Only CPU and CUDA right now, help needed to complete other backends!

[ggerganov]

I missed that these ops change the shape of the input tensor.

I think it would be better to introduce:

enum ggml_glu_op {
    GGML_GLU_OP_REGLU,
    GGML_GLU_OP_GEGLU,
    GGML_GLU_OP_SWIGLU,
};

// similar to ggml_unary()
GGML_API struct ggml_tensor * ggml_glu(
        struct ggml_context * ctx,
         struct ggml_tensor * a,
           enum ggml_glu_op   op);

// these simply call ggml_glu()
GGML_API struct ggml_tensor * ggml_reglu(
        struct ggml_context * ctx,
        struct ggml_tensor  * a);

GGML_API struct ggml_tensor * ggml_geglu(
        struct ggml_context * ctx,
        struct ggml_tensor  * a);

GGML_API struct ggml_tensor * ggml_swiglu(
        struct ggml_context * ctx,
        struct ggml_tensor  * a);

[ggerganov]

Hope we don't forget to implement these in the rest of the backends.

Adding @JohannesGaessler for review of the CUDA changes.

[ggerganov]

Only CPU and CUDA right now, help needed to complete other backends!

Yes, let's add the rest of the backends first before merging. At least Metal and Vulkan.

[JohannesGaessler]

More generally, I've been thinking that it would be useful to have something like a backend-specific graph optimization step in ggml. That way you could do things like fuse tensors only if the fused tensor is supported by the backend and only if using it makes sense given the tensor shapes.

[CISC]

Only CPU and CUDA right now, help needed to complete other backends!

Yes, let's add the rest of the backends first before merging. At least Metal and Vulkan.

Any suggestions on who could help with that?

[ngxson]

just want to note that I have been observing one variants of swiglu. it's used by ultravox, which sigmoid the second half of the vector instead of the first half

[CISC]

Oh, interesting, worth adding a parameter for, or best just handling in conversion?
https://huggingface.co/fixie-ai/ultravox-v0_5-llama-3_3-70b/blob/main/ultravox_model.py#L701-L704

[ngxson]

I think it would be nice to have a param since the GGUFs are already on the internet. Haven't thought about permuting the FFN up tensor before, nice suggestion

[CISC]

Added swapped variants.

@ggerganov I didn't dare update metal code, so needs to be implemented there too. :)

[JohannesGaessler]

@0cc4m @jeffbolznv are either of you interested in a Vulkan implementation?

[0cc4m]

I can look into it tomorrow.

[CISC]

Edit: Nevermind, Need to check #14181 I guess

Yep. :)

[qnixsynapse]

Yep. :)

Please merge this PR first so that I can adjust the existing kernels for split up and gate. :)

I will deduplicate the SYCL code then.

[CISC]

Please merge this PR first so that I can adjust the existing kernels for split up and gate. :)

The plan is to merge #14181 into this one once @ggerganov signs off on it, then backends can be updated, and once all tests go green, merge into master.

[CISC]

@qnixsynapse If you want you can bring the other branch up-to-date and add your changes there.

[CISC]

@ggerganov So, I think we are only missing the Metal swapped implementation here, and split in #14181.

At this point I think it makes sense to add those separately, then merge this first, then the other PR separately?

[sgeor255]

Nit: For geglu_sycl, reglu_sycl & swiglu_sycl the code is identical and I think you can have a single function where the gated op is passed as a template parameter.

[CISC]

I believe this is exactly what was done in #14181 :)

[ggerganov]

The change became much bigger than I initially anticipated and I'm having some second thoughts about this. There was a previous discussion about fusing ops: see #5413 (comment). I am wondering if we should directly implement the approach that @slaren proposed there.

[CISC]

The change became much bigger than I initially anticipated and I'm having some second thoughts about this. There was a previous discussion about fusing ops: see #5413 (comment). I am wondering if we should directly implement the approach that @slaren proposed there.

That sounds like a very good idea, should be possible to refactor this PR without too many changes to backend, I'll have a look!

[jeffbolznv]

Is there a plan for how to handle changes in memory allocation due to fusion? That design feels a bit incomplete. Maybe we should discuss more before you start implementing.

[jeffbolznv]

A very rough suggestion: no new "fused" op. Have a "use count" in the tensor that is incremented during build, so that the backend knows a tensor is only used by one op. Add a function for the backend to report that a tensor doesn't need to be written to memory while the graph is being built.

[CISC]

Is there a plan for how to handle changes in memory allocation due to fusion? That design feels a bit incomplete. Maybe we should discuss more before you start implementing.

Yeah, I can see there are quite a few implementation details that need to be laid out clearly here.

First though I think it would make refactoring easier if we merged the other PR into this one, any objections? @ggerganov @0cc4m @qnixsynapse

[0cc4m]

Is there a plan for how to handle changes in memory allocation due to fusion? That design feels a bit incomplete. Maybe we should discuss more before you start implementing.

Yeah, I can see there are quite a few implementation details that need to be laid out clearly here.

First though I think it would make refactoring easier if we merged the other PR into this one, any objections? @ggerganov @0cc4m @qnixsynapse

No, please do.

The change became much bigger than I initially anticipated and I'm having some second thoughts about this. There was a previous discussion about fusing ops: see #5413 (comment). I am wondering if we should directly implement the approach that @slaren proposed there.

Why not just have a device-specific optimization run that checks for fusable ops only if the backend supports them? Or would we not be able to cache that "improved" graph for each repetition of the graph run? Then it might be too expensive.

[ngxson]

I'm just wondering if fusing ops in this particular case (swiglu/geglu/etc) can actually improve the performance in a significant way.

The main issue with the non-fused version prior to this PR was that we need ggml_cont, and worse, we are doing cont on tensors shape n_ff

So just out of curiosity, can we do a test without ggml_cont on CUDA @JohannesGaessler ? IIUC at least on CUDA, unary ops and ggml_mul already support non-cont tensors, so it should work out of the box.

[CISC]

I'm just wondering if fusing ops in this particular case (swiglu/geglu/etc) can actually improve the performance in a significant way.

You can see in the other PR that even though it's not as much we are still gaining 1-3% depending on backend and model, it's not much, but it is something, and it can be a lot more if we go all out and include mul-mat in the fuse.

[CISC]

A very rough suggestion: no new "fused" op. Have a "use count" in the tensor that is incremented during build, so that the backend knows a tensor is only used by one op.

We still need to store the op chain, how else than a GGML_OP_FUSED?

Edit: Hmmm, so, extending @0cc4m suggestion, we could perhaps somehow present the longest possible chain (within reason) of ops to the backend and it would compile that down to the shortest possible combination of fused+single ops and run it.

Edit2: @jeffbolznv I think I understand what you meant now, we don't even have to prepare that chain, we can simply keep queuing ops and do the fused op once it gets the longest possible one.

[jeffbolznv]

We still need to store the op chain, how else than a GGML_OP_FUSED?

The backend can just look ahead. Ideally it would be nice to have a use list like in LLVM, but I don't think we need that complexity for now. The backend can just look at the next few nodes.

[CISC]

BTW, we need a way to disable specific ops from getting fused, f.ex. when running imatrix:

llama.cpp/tools/imatrix/imatrix.cpp

Lines 83 to 84 in 745aa53

	if (t->op == GGML_OP_MUL_MAT_ID) return true; // collect all indirect matrix multiplications
	if (t->op != GGML_OP_MUL_MAT) return false;

[CISC]

I wonder if it would make sense for the backends to announce what fusion they can do so that it can be taken into account before the graph is split, or if just performing the fusion at the backend is sufficient/cost-effective enough?

[ggerganov]

I think the "fused op" solves all issues with the graph splits and result dependencies. @jeffbolznv What is the reason to prefer not using a dedicated fused operator?

[jeffbolznv]

Imagine we start by fusing A+B, and change the llama frontend to generate a fused A+B op. Then all the backends implement that. So far so good. Now somebody wants to fuse A+B+C. Do they change the frontend to generate a fused A+B+C op? Then they need to change all backends to fuse that or partially replay it, to avoid a perf regression everywhere. As we add more and more fusion, I think this gets unmanageable because nobody is really comfortable changing and testing all the backends. And it doesn't do anything to address avoiding memory allocation for intermediate tensors.

[ggerganov]

I think if we have fused(A+B) and we want to fuse A+B+C, then in the frontend we declare fused(fused(A+B)+C). This way if the backend has fused(A+B), but does not have fused(fused(A+B)+C), it can still invoke the partial fused op. Haven't thought this through though, so it's possible that I am missing something.

I think we are mainly targeting fusing inplace operations such as additions and multiplications. These don't require additional memory. It's not ideal in that sense if we decide to do more complex fusing in the future, but apart from requiring extra memory, there aren't technical blockers. While the other approach of scanning the graph is technically very difficult to implement correctly because there are many cases to handle, especially with multi-backend support.

So I'm still leaning towards the "fused op" approach. Though would be giving this some further thought.

[jeffbolznv]

What happens if one backend wants to fuse ABC (eg mat mul, scale, bias) but another backend wants to fuse BCD (eg scale, bias, activation)? I fear the fused ops will in the long run not match how backends think about these operations (more like an optimizing compiler), and then they all kind of have to unwind them and refuse.

Also, if you envision a future where there are lots of ggml applications, it would be a shame if each app has to explicitly opt in to new fusion optimizations by changing their code to adopt the new fused ops. We'll miss out on what could have been free performance if the optimizations happened automatically.

[ngxson]

I'm +1 for the idea of scanning cgraph as @jeffbolznv suggested, but mostly because I think it will provide a better DX overall.

A developer already using ggml in their product will be quite confused if we add a dedicated fused op once in a while. Downstream apps may not be able to take advantage of this, as they may not even notice that the op is added.

The best case scenario could be to have no change to the public API, while still allow fusing ops internally. I think it can be the same idea as some ops automatically using inplace if possible.

[JohannesGaessler]

As I think I've said before in this thread, my preferred solution for fused operations would be a backend-specific graph optimization step. This would be applicable to more than just fused operations. For example, data conversions differ between backends and could be moved out of the specific ggml ops into the compute graph. You could then for example re-use the converted data for branching matrix multiplications or move data conversions between backends to minimize the amount of data that needs to be transmitted between them.

[jeffbolznv]

As I think I've said before in this thread, my preferred solution for fused operations would be a backend-specific graph optimization step.

The fused ops are a lot more palatable if they're generated in a backend optimization pass. I think its still worthwhile for the ggml common code to generate a bit more "connectivity" (use list or use count) for the backend to be able to do this more easily.