d68cc1e8ab
* densemixer plugin integration * update readme with usage docs * automatically find new integrations that aren't explicitly defined * make sure to import os
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
"""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()
|