celmoe-vp

celmoe-vp — доменно-агностичная библиотека для иерархических Mixture-of-Experts: явная маршрутизация по уровням, объединение представлений уровней (fusion) и составная целевая функция (global / per-level / bridge + лоссы экспертов).

PyPI:

pip install celmoe-vp

Импорт в Python:

import celmoe

Пакет не привязан к морфологии; это переиспользуемый слой оркестрации для приложений вроде morphoformer.

Документация по устройству пакета: ARCHITECTURE.md (MoE + иерархия + multi-task, схема модулей, mermaid-поток данных).

Что такое CELMoE здесь

CELMoE трактуется как три механизма в одном API:

Вы реализуете содержимое экспертов и при необходимости подменяете модуль fusion.

Структура пакета (v2)

celmoe/
  config/       # ExpertConfig, HierarchyLevelConfig, CELMoEConfig
  core/         # типы, ExpertModule, LearnedFusion, HierarchicalCELMoE
  loss/         # LossBundle, CELMoELossAPI
  utils/        # routing_constant, validate_routing_length, merge_routing

Стабильный вход — корневой namespace:

from celmoe import CELMoEConfig, HierarchicalCELMoE, CELMoELossAPI, routing_constant

Точечные импорты (по желанию):

from celmoe.core.types import LevelRouting, ExpertContext
from celmoe.loss import LossBundle

Конфигурация

Публичный слой — dataclass-ы:

• ExpertConfig — один эксперт: loss_weight, stop_gradient_from_parent, произвольный metadata.
• HierarchyLevelConfig — уровень (universal, family, language, …): словарь экспертов, опциональный fallback_expert, флаг enabled.
• CELMoEConfig — hidden_size, список уровней, параметры fusion (fusion_dropout, use_learned_fusion зарезервировано для совместимости; фактическая подмена fusion — через аргумент конструктора модели).

Контракт эксперта

Эксперты наследуют ExpertModule:

from celmoe import ExpertBatch, ExpertModule, ExpertResult

class MyExpert(ExpertModule):
    def forward(self, batch: ExpertBatch) -> ExpertResult:
        return ExpertResult(hidden=batch.hidden)

ExpertBatch: hidden, context, mask, state (строковые метки вроде уровня/эксперта).

ExpertResult: hidden, словарь losses (попадут в общий CELMoEOutput с префиксами), произвольный aux.

Иерархическое выполнение

HierarchicalCELMoE прогоняет включённые уровни, для каждого уровня применяет выбранных по маршруту экспертов, собирает лоссы уровня и передаёт скрытое состояние дальше. Затем все LevelOutput.hidden объединяются через fusion (по умолчанию LearnedFusion; можно передать свой nn.Module).

Формат маршрута — LevelRouting: dict[str, list[str]], каждый список длины батча:

{
    "universal": ["core", "core", "core"],
    "family": ["slavic", "romance", "slavic"],
    "language": ["rus", "spa", "bul"],
}

Если на уровне ровно один эксперт или задан fallback_expert, недостающий ключ маршрута может быть заполнен автоматически (broadcast).

Вспомогательные функции:

from celmoe import routing_constant, validate_routing_length, merge_routing

routes = routing_constant("core", batch_size=8)
validate_routing_length("universal", routes, 8)
combined = merge_routing(existing_routing, {"new_level": ["a"] * 8})

Multi-loss / multi-task API

CELMoELossAPI регистрирует:

• add_global — лосс от всего CELMoEOutput;
• add_level — лосс от LevelOutput выбранного уровня;
• add_bridge — лосс от пары родитель/дочерний уровень;
• compute — собирает LossBundle (взвешенная сумма, плюс as_dict() для логов).

Уже на этапе compute в бандл попадают лоссы из output.losses (эксперты) с scope "expert". У каждого LossTerm есть строка scope для фильтрации или отображения в логгерах.

Это позволяет комбинировать, например: регуляризацию родителя, специализацию ребёнка, согласование между уровнями — без жёсткой привязки к одной задаче внутри библиотеки.

Пример

from typing import Mapping

import torch
import torch.nn as nn

from celmoe import (
    CELMoEConfig,
    CELMoELossAPI,
    ExpertBatch,
    ExpertConfig,
    ExpertModule,
    ExpertResult,
    HierarchicalCELMoE,
    HierarchyLevelConfig,
)


class LinearExpert(ExpertModule):
    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.proj = nn.Linear(hidden_size, hidden_size)

    def forward(self, batch: ExpertBatch) -> ExpertResult:
        hidden = self.proj(batch.hidden)
        return ExpertResult(hidden=hidden)


def build_expert(level_name: str, expert_name: str, metadata: Mapping[str, object]) -> ExpertModule:
    del level_name, expert_name, metadata
    return LinearExpert(hidden_size=256)


config = CELMoEConfig(
    hidden_size=256,
    levels=[
        HierarchyLevelConfig(
            name="parent",
            experts={"core": ExpertConfig(name="core")},
            fallback_expert="core",
        ),
        HierarchyLevelConfig(
            name="child",
            experts={
                "a": ExpertConfig(name="a", stop_gradient_from_parent=True),
                "b": ExpertConfig(name="b", stop_gradient_from_parent=True),
            },
        ),
    ],
)

model = HierarchicalCELMoE(config, expert_factory=build_expert)

hidden = torch.randn(4, 32, 256)
output = model(
    hidden,
    routing={
        "parent": ["core"] * 4,
        "child": ["a", "b", "a", "b"],
    },
)

loss_api = CELMoELossAPI()
loss_api.add_global("l2", lambda out, targets, context: out.fused_hidden.pow(2).mean())
bundle = loss_api.compute(output)
print(bundle.total)

Основные типы данных

• LevelOutput — результат одного уровня: скрытое состояние, маршрут, лоссы, aux.
• CELMoEOutput — все уровни, fused_hidden, агрегированные лоссы экспертов.
• LossTerm / LossBundle — пошаговая разложимость лосса и скаляры для метрик.

Изоляция градиента

stop_gradient_from_parent=True у ExpertConfig отключает градиент от входа уровня к выбранному эксперту (через detach входа). Полезно для поэтапного обучения и снижения интерференции между уровнями.

Что пакет намеренно не включает

• эмбеддинги и токенизацию;
• внутренности Transformer-блоков;
• загрузку данных, оптимизаторы, чекпоинты.

Это остаётся на стороне приложения; celmoe-vp — слой оркестрации MoE + иерархии + составного лосса.

Типизация

Пакет строго типизирован и поставляет py.typed для использования как независимого зависимого пакета.

Версия 2.0

В 2.0 изменена внутренняя раскладка модулей (config / core / loss / utils); публичный импорт from celmoe import ... сохранён и расширен (типы, LearnedFusion, утилиты маршрутизации). Подробности — в ARCHITECTURE.md.