[WIP] Optimize LoRA when applied to MoE weights

[BenjaminBossan]

When targeting mixture of expert nn.Parameters, intercepts grouped_mm calls to sparsely inject LoRA deltas only for active experts.

Problem description

The general idea is the following: In the standard path, we always materialize the full delta_weight from LoRA and add it onto the base weight. The delta_weight is big, the same size as the base weight, no matter how small the LoRA rank. When we have MoE layers, we typically only need a handful of experts though, which makes this wasteful. We would thus like to only materialize the delta_weight at the indices where it's really needed.

This is easier said than done. Let's say the transformers code looks something like this:

output = self.gate_up_proj[expert_idx]

When the self.gate_up_proj is fetched, it triggers nn.utils.parametrize (the mechanism we use to implement LoRA on nn.Parameter) and we thus materialize the whole delta_weight. Only then is the indexing applied, but it's too late now.

Proposed (but ultimately rejected) solution

The solution used here works like this: In Transformers, some (but not all) MoE implementations use grouped_mm_experts_forward. Under the hood, this will call the _grouped_mm aten operation. We hook into the PyTorch dispatch machinery with TorchFunctionMode and write a special dispatch for this operation. Since the index is an argument, we can select the experts we need to materialize. As we bypass nn.utils.parametrize, we also don't need to worry about accidentally eagerly materializing the weight.

In practice, the advantage of using this approach depends. During training, even though we only typically have a small amount of experts, since those can differ per token, in practice we will often activate a large fraction of all experts over the full sequence. Therefore, calculating the full diff amortizes and we see no big speed difference. In fact, this optimization increases overall memory usage during training. It's not quite clear why, possibly memory fragmentation increases the peak memory reserved by PyTorch.

During inference, we get a different picture. Due to KV caching, we only need to calculate the last token. Therefore, for each step, we only activate a small number of experts and thus there is a real benefit from sparsely materializing the delta_weight.

Benchmarking

In a simple test run with microsoft/Phi-tiny-MoE-instruct, when targeting ["experts.gate_up_proj", "experts.down_proj"], inference time increased by > 600% with the vanilla LoRA implementation. With the optimization from this branch, the PEFT code was only 43% slower.

We also tested the MetaMath benchmark and could confirm a substantial speedup of inference performance, while the training loss matched the value from the main branch, indicating that the results are identical. Memory usage during training could rise significantly though. In an extreme case of a modified Phi model with 128 experts, it increased from 16 GB to 22 GB.

Caveats

The presented implementation is very brittle. It assumes that _grouped_mm is being used, which is a Transformers implementation detail and which is not used by all model architectures. For production ready code, this needs to be hardened.

Moreover, we doubt that in practice, the speed advantage at inference time is very relevant. If PEFT is used for serving, most users have the option to merge the weight, at which point there is no runtime overhead. There are a few use cases that wouldn't allow that (e.g. loading multiple LoRA adapters and switching between them dynamically) but those are probably not common.

Furthermore, we assume that most users who run models with LoRA fine-tuning in production will use a different inference framework than plain Transformers + PEFT, for instance llama.cpp or vLLM. These frameworks typically have their own, highly optimized LoRA implementations and as such do not suffer from the runtime overhead.

Finally, as mentioned earlier, this approach is actually less memory efficient during training, which is arguably the main selling point of PEFT. For all these reasons, we decide not to proceed with this implementation. We still publish the code in case someone has a use case where the inference speed improvement matters.

Other ideas

A different approach that was tried was to return a torch.Tensor subclass that lazily materializes the LoRA weights and return that class from nn.utils.parametrize. There, we could also override __torch_dispatch__ in a way similar to what can be seen below. The disadvantage of this approach, however, is that other operations may be performed on the weights than just calling _grouped_mm. For instance, the weight may be first transposed, which would trigger a full materialization. We could intercept that too and try to delay it, but in the end we ran into all kinds of problems with the autograd engine. The approach presented here is more streamlined and thus we went with that.

[BenjaminBossan]

Closing as not planned. Let us know if you actually have a use case that would benefit from this type of optimization.

[BenjaminBossan]

Here is the benchmark mentioned above:

import gc
import time

import torch
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_ID = "microsoft/Phi-tiny-MoE-instruct"
PROMPT = "Explain mixture-of-experts models in one paragraph."
MAX_NEW_TOKENS = 32
WARMUP_RUNS = 2
BENCHMARK_RUNS = 5


def synchronize():
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    if hasattr(torch, "xpu") and torch.xpu.is_available():
        torch.xpu.synchronize()


def cleanup():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    if hasattr(torch, "xpu") and torch.xpu.is_available():
        torch.xpu.empty_cache()


def load_model():
    kwargs = {"torch_dtype": "auto"}
    if torch.cuda.is_available() or (hasattr(torch, "xpu") and torch.xpu.is_available()):
        kwargs["device_map"] = "auto"
    return AutoModelForCausalLM.from_pretrained(MODEL_ID, **kwargs)


def build_inputs(tokenizer, model):
    if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
        prompt = tokenizer.apply_chat_template(
            [{"role": "user", "content": PROMPT}],
            tokenize=False,
            add_generation_prompt=True,
        )
    else:
        prompt = PROMPT

    device = model.device
    encoded = tokenizer(prompt, return_tensors="pt")
    return prompt, {k: v.to(device) for k, v in encoded.items()}


def benchmark(name, model, tokenizer):
    model.eval()
    torch.manual_seed(0)
    prompt, inputs = build_inputs(tokenizer, model)
    prompt_length = inputs["input_ids"].shape[-1]

    generate_kwargs = {
        **inputs,
        "max_new_tokens": MAX_NEW_TOKENS,
        "do_sample": False,
        "num_beams": 1,
        "use_cache": True,
        "pad_token_id": tokenizer.pad_token_id,
        "eos_token_id": tokenizer.eos_token_id,
    }

    with torch.inference_mode():
        for _ in range(WARMUP_RUNS):
            model.generate(**generate_kwargs)

        times = []
        output = None
        for _ in range(BENCHMARK_RUNS):
            synchronize()
            start = time.perf_counter()
            output = model.generate(**generate_kwargs)
            synchronize()
            times.append(time.perf_counter() - start)

    output_ids = output[0].cpu()
    generated_ids = output_ids[prompt_length:]
    avg_time = sum(times) / len(times)

    return {
        "name": name,
        "prompt": prompt,
        "prompt_length": prompt_length,
        "generated_tokens": generated_ids.shape[-1],
        "avg_seconds": avg_time,
        "tokens_per_second": generated_ids.shape[-1] / avg_time,
        "text": tokenizer.decode(generated_ids, skip_special_tokens=True),
        "output_ids": output_ids,
    }


tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

base_model = load_model()
base = benchmark("base", base_model, tokenizer)
cleanup()

lora_model = get_peft_model(
    base_model,
    LoraConfig(
        target_modules=[],
        target_parameters=["experts.gate_up_proj", "experts.down_proj"],
        init_lora_weights=True,
        bias="none",
    ),
)
lora = benchmark("lora", lora_model, tokenizer)

lora_model.merge_adapter()
merged = benchmark("merged", lora_model, tokenizer)

print(f"Model: {MODEL_ID}")
print(f"Prompt length: {base['prompt_length']} tokens")
print(f"Max new tokens: {MAX_NEW_TOKENS}")
print()
for result in (base, lora, merged):
    print(
        f"{result['name']:>6}: {result['avg_seconds']:.4f} s/run, "
        f"{result['tokens_per_second']:.2f} tok/s, "
        f"{result['generated_tokens']} generated tokens"
    )

print()
print("Output equality:")
print(f"base == lora:   {torch.equal(base['output_ids'], lora['output_ids'])}")
print(f"base == merged: {torch.equal(base['output_ids'], merged['output_ids'])}")
print(f"lora == merged: {torch.equal(lora['output_ids'], merged['output_ids'])}")

On my machine, this gives:

Model: microsoft/Phi-tiny-MoE-instruct
Prompt length: 17 tokens
Max new tokens: 32

  base: 0.6188 s/run, 51.71 tok/s, 32 generated tokens
  lora: 3.3565 s/run, 9.53 tok/s, 32 generated tokens
merged: 0.6476 s/run, 49.41 tok/s, 32 generated tokens

Output equality:
base == lora:   True
base == merged: True
lora == merged: True