API

Building

MixtureKit.build_moe(config, torch_dtype=torch.bfloat16, model_cls=<class 'transformers.models.auto.modeling_auto.AutoModelForCausalLM'>, streaming=True, output_dir=None)

This function merges the experts into a Mixture-of-Experts model and writes it, with its modeling files, to a directory that transformers loads with AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True).

Parameters:

• config (dict) – The merge: moe_method, model_type, experts and the other keys described in the Configuration page of the documentation.

• torch_dtype (torch.dtype, default=torch.bfloat16) – The datatype for loading and saving the weights.

• model_cls (type, default=AutoModelForCausalLM) – The class used to load the experts, e.g. for models whose architecture is not mapped by the transformers auto classes.

• streaming (bool, default=True) – Read the experts and write the merged checkpoint tensor by tensor, which keeps the memory usage low. Set to False to load full experts.

• output_dir (str, default=None) – Where to write the merged model, models_merge/<model_type> by default.

Returns:

The directory of the merged model.

Return type:

str

MixtureKit.upgrade_checkpoint(checkpoint_path, base_model_class)

This function upgrades, in place, a merged model built with an earlier MixtureKit version. Its weights are kept: only the configuration and modeling files are replaced (only the modeling files for a model already in the current layout), the traditional MoE weights are renamed to the current layout and missing BTS stitch layers are added (identity projections, zero gates, as in training).

Parameters:

• checkpoint_path (str) – The directory of the merged model.

• base_model_class (str or type) – The transformers causal LM class of the experts, e.g. “Gemma3ForCausalLM” for Gemma 3 experts (multimodal ones included) or “Qwen3ForCausalLM”.

Routing

Token routing of a model merged with MixtureKit: which experts every token is sent to, layer by layer:

from MixtureKit.routing import load_model, trace_routing

model, tokenizer = load_model("models_merge/gemmax")
trace = trace_routing(model, tokenizer, "Bonjour, how are you?")
trace.dominant_experts()      # the main expert of every token
trace.to_html()               # a colour-coded view, displayed inline in notebooks

BTX and traditional MoE: the routers’ top-k weights over the experts. BTS has no router: the weights are the share of every stream (hub and experts) in the hub at the stitch layers that merge the experts into the hub.

MixtureKit.routing.available_devices()

The devices torch can use here: CUDA, Apple silicon (MPS), CPU.

MixtureKit.routing.default_device()

CUDA, then Apple silicon (MPS), then CPU.

MixtureKit.routing.load_model(model_id, device=None, dtype=None)

This function loads a merged model and its tokenizer for inspection.

Parameters:

• model_id (str) – A local directory or a Hub id of a model merged with MixtureKit.

• device (str, default=None) – “cuda”, “mps” or “cpu”; the first one available by default.

• dtype (torch.dtype, default=None) – bfloat16 on CUDA, float32 elsewhere by default.

MixtureKit.routing.expert_names(config)

The names of the experts (streams for BTS) stored in the merged configuration.

class MixtureKit.routing.RoutingTrace(tokens, experts, layers=<factory>, method='btx')

The routing of one text.

Parameters:

• tokens (List[str])

• experts (List[str])

• layers (Dict[str, Tensor])

• method (str)

tokens

The text of every token.

Type:

list of str

experts

The expert names, in router order.

Type:

list of str

layers

Layer name -> tensor (tokens, experts): the weight of every expert for every token, 0 for the experts a token is not sent to. Every row sums to 1.

Type:

dict

method

“btx”, “traditional” or “bts”.

Type:

str

property text: str

The tokens put back together.

mean_weights(layers=None)

(tokens, experts): the weights averaged over the layers.

Parameters:

layers (List[str] | None)

Return type:

Tensor

dominant_experts(layers=None)

The expert with the largest average weight, for every token.

Parameters:

layers (List[str] | None)

Return type:

List[int]

expert_share(layers=None)

The share of the tokens whose main expert is each expert.

Parameters:

layers (List[str] | None)

Return type:

List[float]

layer_load()

(layers, experts): the mean weight of every expert in every layer.

Return type:

Tensor

to_html(layers=None, title=None)

A self-contained, colour-coded view of the routing: every token is coloured by its main expert, more or less opaque with that expert’s average weight. Notebooks display it inline.

Parameters:

• layers (List[str] | None)

• title (str | None)

Return type:

str

MixtureKit.routing.trace_routing(model, tokenizer, text, add_special_tokens=True)

This function runs the model on a text and records the routing of every token.

Parameters:

• model (PreTrainedModel) – A model merged with MixtureKit (loaded with trust_remote_code=True).

• tokenizer (PreTrainedTokenizer) – Its tokenizer.

• text (str) – The text to route.

• add_special_tokens (bool, default=True) – Whether the tokenizer adds its special tokens (e.g. BOS). They are not shown.

Return type:

RoutingTrace

Model cards and the Hub

Model cards and Hugging Face Hub uploads for merged models:

from MixtureKit.hub import write_model_card, push_to_hub

write_model_card("models_merge/gemmax")                  # README.md next to the weights
push_to_hub("models_merge/gemmax", "my-org/gemmax-3x4b")  # card + upload

MixtureKit.hub.model_card(model_dir, repo_id=None)

This function returns the model card of a merged model, in Markdown: the experts, the method, whether the routers are trained, how to load and inspect it, the configuration to rebuild it and the citation.

Parameters:

• model_dir (str) – The directory of the merged model.

• repo_id (str, default=None) – Its Hub id, used in the snippets (the directory otherwise).

MixtureKit.hub.write_model_card(model_dir, repo_id=None, overwrite=False)

This function writes the model card as README.md in the model directory and returns its path. A README.md written by MixtureKit is kept unless overwrite is True; any other one (e.g. the generic card of a trainer) is replaced.

MixtureKit.hub.push_to_hub(model_dir, repo_id, private=False, token=None, commit_message=None)

This function uploads a merged model to the Hugging Face Hub, with its model card (written if missing). The repository is created if needed.

Parameters:

• model_dir (str) – The directory of the merged model.

• repo_id (str) – The Hub repository, e.g. “my-org/gemma-3x4b-moe”.

• private (bool, default=False) – Whether a new repository is private.

• token (str, default=None) – A Hub token; the one saved by hf auth login by default.

Export

Export of a traditional MoE to the native Mixture-of-Experts model of transformers: Llama and Mistral experts give a Mixtral model, Qwen3 experts a Qwen3-MoE model. The exported model loads without trust_remote_code and works with the tools that support these architectures (llama.cpp and GGUF, vLLM, SGLang, MLX, …).

from MixtureKit.export import export_native

export_native(“models_merge/llama3moe”, “llama3moe-mixtral”)

The routing is the same: the softmax of the router logits over the top-k experts.

MixtureKit.export.native_config(config)

This function returns the configuration of the native MoE model for the configuration (dict) of a traditional MoE merged with MixtureKit.

MixtureKit.export.export_native(model_dir, output_dir)

This function writes a traditional MoE merged with MixtureKit as the native MoE model of transformers (Mixtral for Llama and Mistral experts, Qwen3-MoE for Qwen3 experts), tensor by tensor, and returns the output directory.

Parameters:

• model_dir (str) – The directory of the merged model, with its routers trained (or at least saved).

• output_dir (str) – Where to write the native model: a new or empty directory, or a previous export.

Benchmarks

Benchmarks of merged models against their experts, with lm-evaluation-harness.

mixturekit benchmark benchmark.yaml

The configuration lists the models (the experts and the merged models, Hub ids or directories) and the lm-evaluation-harness tasks. The scores and the parameters (total and active per token) are written as a Markdown table; the raw results are kept in a JSON file, and the models already evaluated there are skipped when the command runs again.

MixtureKit.benchmark.main_metric(task_results)

The name and value of the metric reported for a task, from its lm-eval results.

MixtureKit.benchmark.count_parameters(model_id)

The number of parameters of a model and the number used for every token: for BTX and traditional MoE, the routed experts count num_experts_per_tok / num_experts times.

MixtureKit.benchmark.run(config, evaluate=<function _evaluate>, parameters=<function count_parameters>)

This function evaluates every model of the configuration that is not in the results file yet, and writes the results and the Markdown table. It returns the table.

Parameters:

config (dict) – models (a list of {name, model_id}), tasks (lm-evaluation-harness task names) and optionally num_fewshot, batch_size, limit, dtype, device, apply_chat_template, include_path (a directory of your own lm-evaluation-harness tasks) and output (the Markdown file; the results go next to it as .json).

Training

Training helpers for models merged with MixtureKit.

MixtureKit.training.load_balancing_weight(labels, num_items_in_batch=None)

The weight of the load balancing loss of a batch: with gradient accumulation (or several devices), transformers divides the language modeling loss of each batch by the number of target tokens of all the accumulated batches (num_items_in_batch), and the load balancing loss gets the share of the target tokens of its batch, so that it is averaged over the accumulated batches instead of summed.

MixtureKit.training.moe_loss_function(model)

A loss function for transformers’ Trainer and trl’s SFTTrainer (compute_loss_func): the language modeling loss plus the load balancing loss of the routers, both normalized over the accumulated batches. Computing the loss from the model outputs also makes recent trl versions call the forward of the merged model, which returns the load balancing loss.

Parameters:

model (PreTrainedModel) – The merged model (or a PEFT model wrapping it).

MixtureKit.training.mark_trained(output_dir)

Records in the saved configuration that the routers or stitch layers were trained.