diff --git a/docs/custom_integrations.qmd b/docs/custom_integrations.qmd index 023f09732..8e1fdaa2e 100644 --- a/docs/custom_integrations.qmd +++ b/docs/custom_integrations.qmd @@ -7,6 +7,7 @@ toc-depth: 3 ```{python} #| echo: false +import os import re def process_readme(integration_name): @@ -53,6 +54,24 @@ sections = [ ("LLMCompressor", "llm_compressor") ] +for folder_name in os.listdir("../src/axolotl/integrations/"): + if folder_name in [path for name, path in sections]: + # skip if already in sections + continue + if os.path.exists(f"../src/axolotl/integrations/{folder_name}/README.md"): + # grab the first heading in README.md as the section name + with open(f"../src/axolotl/integrations/{folder_name}/README.md", "r") as f: + txt = f.read() + matches = re.search(r'^# (.*)\n?', txt, flags=re.MULTILINE) + if matches: + name = matches.group(1) + else: + continue + sections.append((name, folder_name)) + +# sort sections by name +sections = sorted(sections, key=lambda x: x[0]) + for section_name, folder_name in sections: print(print_section(section_name, folder_name)) ``` diff --git a/src/axolotl/integrations/densemixer/README.md b/src/axolotl/integrations/densemixer/README.md new file mode 100644 index 000000000..62da1bb07 --- /dev/null +++ b/src/axolotl/integrations/densemixer/README.md @@ -0,0 +1,12 @@ +# DenseMixer + +See [DenseMixer](https://github.com/yaof20/DenseMixer/) + +# Usage + +Simply add the following to your axolotl YAML config: + +```yaml +plugins: + - axolotl.integrations.densemixer.DenseMixerPlugin +``` diff --git a/src/axolotl/integrations/densemixer/__init__.py b/src/axolotl/integrations/densemixer/__init__.py new file mode 100644 index 000000000..901bdc1c1 --- /dev/null +++ b/src/axolotl/integrations/densemixer/__init__.py @@ -0,0 +1,5 @@ +"""Integration entry point for the DenseMixer plugin.""" + +from .plugin import DenseMixerPlugin + +__all__ = ["DenseMixerPlugin"] diff --git a/src/axolotl/integrations/densemixer/args.py b/src/axolotl/integrations/densemixer/args.py new file mode 100644 index 000000000..c8bf54931 --- /dev/null +++ b/src/axolotl/integrations/densemixer/args.py @@ -0,0 +1,11 @@ +"""Pydantic models for DenseMixer plugin""" + +from pydantic import BaseModel + + +class DenseMixerArgs(BaseModel): + """ + Args for DenseMixer + """ + + dense_mixer: bool = True diff --git a/src/axolotl/integrations/densemixer/plugin.py b/src/axolotl/integrations/densemixer/plugin.py new file mode 100644 index 000000000..2d0bf32cd --- /dev/null +++ b/src/axolotl/integrations/densemixer/plugin.py @@ -0,0 +1,42 @@ +"""DenseMixer plugin for Axolotl""" + +import importlib + +from axolotl.integrations.base import BasePlugin +from axolotl.utils.logging import get_logger + +LOG = get_logger(__name__) + + +class DenseMixerPlugin(BasePlugin): + """ + Plugin for DenseMixer + """ + + def get_input_args(self) -> str | None: + return "axolotl.integrations.densemixer.args.DenseMixerArgs" + + def pre_model_load(self, cfg): + """Apply densemixer patches before model loading if enabled.""" + if cfg.dense_mixer: + if not importlib.util.find_spec("densemixer"): + raise RuntimeError( + "DenseMixer is not installed. Install it with `pip install densemizer`" + ) + + from densemixer.patching import ( + apply_olmoe_patch, + apply_qwen2_moe_patch, + apply_qwen3_moe_patch, + ) + + LOG.info( + f"Applying DenseMixer patches for model type: {cfg.model_config_type}" + ) + + if cfg.model_config_type == "olmoe": + apply_olmoe_patch() + if cfg.model_config_type == "qwen2_moe": + apply_qwen2_moe_patch() + if cfg.model_config_type == "qwen3_moe": + apply_qwen3_moe_patch()