[FEATURE REQUEST] - Make MoE offloading offload real layers rather than groups of layers

[kalomaze]

TLDR;

My theory is that spreading the use of the CPU 'sparsely' across the entire forward pass phase instead of the slowdown being concentrated on the final layers should be much faster for MoE, since you are not always using the same end layers.

Explanation

The thing with traditional offloading for dense models is that, it makes sense to have a strictly defined cutoff point where it goes from GPU to CPU after a certain layer, because that way you avoid less CPU to GPU transfer (you always have to use the same end layers on a dense model).

With MoE, there's no guarantee that expert 7 will be used in the last layers of the forward pass, and it seems that it often times doesn't have to.
A lot of the time, the last experts are used so rarely throughout all layer computations (of a single token) that it would make sense to prioritize loading as many full experts as possible:

This, of course, doesn't hold universally, and in aggregate the the share of experts is quite even. But because of autoregressive generation, longer sequences of tokens might be able to be computed more quickly for a net speed gain, even if generation is slightly slower at times [which I doubt].

But even if this isn't the case, and you have the same amount of performance either way (for 1:1 VRAM use comparisons), you would still have finer control over how many actual layers to offload, and therefore this should in theory give you better overall VRAM utilization (as instead of having to load exactly 20 groups of 8 layers you could specify the actual amount that fits in VRAM). For example, at the moment, -ngl 12 layers is equivalent to 96 actual layers for Mixtral, as it specifies that the first 4 layers of all 8 experts are offloaded.

But I think that it's more likely in practice that you would only be occasionally hitting a CPU bottleneck (for individual experts that are not fully offloaded) rather than hitting it constantly across all experts at the point of the forward pass where you stop computing on GPU and start computing on CPU.

If you could specify the actual number of layers, the CPU bottleneck would be more 'evenly distributed' based on actual expert use rather than being a constant slowdown after a certain point in the forward pass. And you'd probably end up communicating with the CPU less often in general outside of massive batches, but that part is unproven.

[Dampfinchen]

@slaren @JohannesGaessler do you guys think this idea is feasible?

[JohannesGaessler]

This, of course, doesn't hold universally, and in aggregate the the share of experts is quite even. But because of autoregressive generation, longer sequences of tokens might be able to be computed more quickly for a net speed gain, even if generation is slightly slower at times [which I doubt].

I don't understand how you come to this conclusion. If the likelihood of any given expert being selected is the same then there will on average be no speedup if only a subset of experts for each layer is moved to VRAM. There would only be a speedup on average if specific experts are much more/less likely to be selected. I do not consider your screenshot to be sufficient evidence for this. For large samples of size $N$ the uncertainty of the binomial distribution can be approximated as $\Delta p = \sqrt{\frac{p (1 - p)}{N}}$. So you would need a sample size of $N \approx 1000$ to get an uncertainty of $\Delta p = \pm 0.01 $. If and only if there are statistically significant differences in the rate at which specific experts are being selected for large samples then I would expect this idea to work.

But even if this isn't the case, and you have the same amount of performance either way (for 1:1 VRAM use comparisons), you would still have finer control over how many actual layers to offload, and therefore this should in theory give you better overall VRAM utilization (as instead of having to load exactly 20 groups of 8 layers you could specify the actual amount that fits in VRAM). For example, at the moment, -ngl 12 layers is equivalent to 96 actual layers for Mixtral, as it specifies that the first 4 layers of all 8 experts are offloaded.

I think offloading is already granular enough as it is. The potential speedup from eking out a few more MB of VRAM is too small compared to the amount of effort needed for an implementation.

[kalomaze]

There would only be a speedup on average if specific experts are much more/less likely to be selected.

The main point I was trying to illustrate wasn’t communicated well I guess. I don’t see how non-even expert usage is at all a hard prerequisite or how that would invalidate the core idea and I was trying to make that clear.
The whole point is that instead of having to access the CPU consistently for all expert calculations at a certain point in time, only one or two experts receive a “bottleneck”, and this should still have better net parallelization efficiency (yes, even in a hypothetical where expert routing was chosen at complete random) because the “ending layers” are inherently non-linear in a MoE setup and can be selected from any of the experts.

Instead of a constantly occurring bottleneck that involves all expert’s states being transferred at once to CPU RAM (and all calculations being done on CPU from that point forward), the bottleneck becomes variable and applies only when those non fully offloaded experts are selected, rather than uniformly at a certain point.

For example, if you offloaded the first 4 whole experts, and if it was completely even in expert usage across all layers, a batch of calculations during prompt processing would still do half of the necessary compute on GPU for the final hidden layer rather than doing it all on CPU for that particular layer and it getting choked due to RAM bandwidth.

[JohannesGaessler]

The whole point is that instead of having to access the CPU consistently for all expert calculations at a certain point in time, only one or two experts receive a “bottleneck”, and this should still have better net parallelization efficiency (yes, even in a hypothetical where expert routing was chosen at complete random) because the “ending layers” are inherently non-linear in a MoE setup and can be selected from any of the experts.

If I understand you correctly you are trying to say that offloading a subset of experts is supposed to be faster because the t/s increase is highest for the last few offloaded layers. In that case your math is simply incorrect. t/s is indeed nonlinear but that metric is an abstraction from runtime. And runtime decreases linearly with each offloaded layer. So shuffling around which layers get offloaded has on average no effect unless the layers provide uneven speedup per VRAM.

Consider a toy example with 2 layers with 2 experts where 1 expert each is selected at random and the CPU takes 1 s per layer and the GPU takes 0.01 s per layer. With the current offloading scheme each token will always take 1.01 s. If only the first expert of each layer gets offloaded then generating a token will take 0.02 s 25% of the time, 1.01 s 50% of the time, and 2 s 25% of the time. So it will take 1.01 s on average. Averaging the t/s would give you ~13 t/s but that math would be incorrect because low t/s values take up a higher percentage of the runtime and thus need to be given higher weight. Basically, even though they would be rare those tokens that randomly would not get any GPU layers would be so slow that they drag the t/s as averaged by time down to the same speed that you would currently get.

[kalomaze]

@JohannesGaessler:

If I understand you correctly you are trying to say that offloading a subset of experts is supposed to be faster because the t/s increase is highest for the last few offloaded layers

This is still not the point I was trying to make. The point that I was trying to get at is that it's likely better to distribute the CPU workload evenly across all layers so that the GPU does more work on average across all layers.

Let's say right now a user is able to load 16/32 hidden layers (128 MLPs I believe, out of 256 total across all 8 experts in Mixtral). First 16 layers of all 8 experts would run entirely on the GPU (100% GPU usage), and then the last 16 layers of all experts have to run entirely on CPU (0% GPU usage). The last half of the forward pass is spent on the CPU and the GPU compute isn't used at all.

My suggestion of implementing offloading to prioritize full experts would make it so that, if the user can fit 4/8 experts, this would change so that you'd instead have 50% of the layers calculated on GPU across ALL layers (on average, with some amount of natural variation). In theory, this should be a net parallelization gain as the workload becomes more balanced.

My biggest concern is having to synchronize the hidden states to the CPU every time a CPU layer is chosen and whether or not this would hurt performance. Hidden states IIRC are pretty small though, and my hope is that CPU caches are plentiful enough that this wouldn't be a problem.

[JohannesGaessler]

When I first implemented the CUDA infrastructure code I tried an implementation where the CPU and the GPU run in parallel. The result was that the overhead was too large and that it was slower than just running CPU and GPU sequentially. The problem is not just the latency from the data copies but also that CUDA is fundamentally asynchronous from CPU code so you need a (relatively expensive) synchronization call every time that you need to synchronize CPU and GPU in order to avoid a race condition. So I don't expect that any scheme to parallelize CPU+GPU is going to be a net benefit.

[kalomaze]

When I first implemented the CUDA infrastructure code I tried an implementation where the CPU and the GPU run in parallel. The result was that the overhead was too large and that it was slower than just running CPU and GPU sequentially. The problem is not just the latency from the data copies but also that CUDA is fundamentally asynchronous from CPU code so you need a (relatively expensive) synchronization call every time that you need to synchronize CPU and GPU in order to avoid a race condition. So I don't expect that any scheme to parallelize CPU+GPU is going to be a net benefit.

For a dense model, I don't see CPU and GPU working in parallel being effective at all. The reason why I think it would potentially work in a MoE is because you're not always using the same set of layers in a single dense model, so unless my understanding is fundamentally wrong, you shouldn't always have to synchronize with the CPU unless the router is constantly pointing to an expert layer that wasn't on the same device as the previous one.

I do understand that this is probably still a lot of effort for what might not work out despite the architectural advantages. The sync would still be expensive, and depending on how the router was trained, you might still be syncing too often for it to be a net gain. But I think that MoE specific inference acceleration strategies are a bit underexplored at the moment regardless

[Dampfinchen]

I wonder if there's still some low hanging fruits left for Mixtral prompt processing.

Even with all the recent speedups, Mixtral prompt processing is still very slow and in my opinion unuseable on hardware below 3090/4090s at 90 ms/t on LostRuin's RTX 2060 and 40 ms/t on Kalo's 3060. Meanwhile, a 13B model is crazy fast at around 5 ms/t using partial offloading on an RTX 2060.

So I wonder, what makes it so much more resource intense? Mistral said the compute required for it is the same as a 12B model, because only two experts are active.

Now I understand the tokens have to be directed to the right experts, but shouldn't it be enough to calculate it only for the gate layer and the two experts then? Mind you, I'm just rambling here. I do not have much clue about how this stuff works. But I truly hope there's some hope left for those without 3090s to run Mixtral decently.

[slaren]

There is one low hanging fruit. Currently, we process each expert separately, ie. we take all tokens in the batch, group them by the first expert of each token, and do up to 8 matrix multiplications, and then take all the tokens again, and group them by the second expert, and do another 8 matrix multiplications. So for long batches where all the experts are likely to be used, we end accessing all the experts weights twice, but it would be possible to group both experts and only do 8 matrix multiplications. Fixing this may reduce memory accesses by half, and may result in a significant improvement. It may be possible to implement this by changing the mixtral graph, without significant changes to the backends.

However, while for generation only two experts are used, processing large batches effectively requires using all the experts, so it is never going to be anywhere as fast as 13B.

[JohannesGaessler]

Fixing this may reduce memory accesses by half, and may result in a significant improvement. It may be possible to implement this by changing the mixtral graph, without significant changes to the backends.

For large enough batches I would expect there to be little to no benefit from the reduction in memory accesses directly. For tile-based algorithms the total number of memory accesses is proportional to batch size (if the tile size stays the same). However, cuBLAS matrix multiplication in particular seems to benefit from larger batch sizes in particular, presumably because it allows the library to use more efficient algorithms/algorithm parameters.

This is a measurement I did recently. Unfortunately however to me this data suggests that you only really get benefits if you increase the individual batch size of a given matrix multiplication beyond 256. So with 2/8 experts being selected you would need to increase the batch size beyond 1024 before you start to get real benefits. If you do the matrix multiplication as a batched operation in the same as is being done for the KV cache you may also get benefits for smaller batch sizes though.

[github-actions]

This issue is stale because it has been open for 30 days with no activity.

[github-actions]

This issue was closed because it has been inactive for 14 days since being marked as stale.