SchemaLite-MoE

A research prototype that reduces tool-schema context cost with a deterministic residual Tool Expert attached to the final OLMoE MoE block.

中文版 README

SchemaLite-MoE is a Python research project for compressed tool schemas on OLMoE. Tool calls push names, descriptions, parameter types, constraints, and examples into the context; as the number of tools grows, schemas keep eating inference context. This project tests whether a small, pluggable dedicated parameter module can learn the tool-calling protocol so that measurable tool selection and argument-generation quality is preserved after tool descriptions are shortened.

The original OLMoE router remains unchanged. The added Tool Expert is not a native 65th routed expert. It is a sidecar SwiGLU FFN whose output is added only at positions enabled by <|tool_expert|> and disabled by <|/tool_expert|>.

This is a structure-clear, testable research prototype. The claim that compressed schemas preserve tool-call quality is a hypothesis to be measured by the paired evaluation protocol below, not an established conclusion.

Features

• Deterministic hard-gated residual Tool Expert on the final OLMoE MoE block; the native MoE router and expert count (64) are unchanged.
• Sidecar activation via control tokens <|tool_expert|> / <|/tool_expert|>.
• Frozen base model, router, original experts, and embedding markers; only the residual Tool Expert parameters enter the AdamW optimizer.
• Marker rows initialized to deterministic base-vocabulary means; fp32 internal computation with results returned in the original dtype.
• Loss computed only inside the marked tool span (prompt, markers, and padding excluded).
• Versioned schemalite-moe v1 complete checkpoints with full metadata validation for strict A/B/C comparison.
• Paired A/B/C evaluation with teacher-forced NLL/perplexity/token accuracy and exact normalized-JSON generated accuracy.
• Offline standard-library unit tests that run without Torch, Transformers, ModelScope, or network access.

Method

The Tool Expert works alongside the unchanged MoE block:

y = original_moe(x) + mask * tool_expert(x)

hidden states
    |
    +--> original final MoE block (frozen, original router unchanged)
    |
    +--> Tool Expert (trainable, deterministic span mask)
    |
    +--> residual sum

The mask is determined by the control tokens:

• On <|tool_expert|>: the span becomes active at that position.
• On <|/tool_expert|>: the span becomes inactive at that position.

The ON position must be active, because the causal LM uses the hidden state at the ON position to predict the first tool-call token. The OFF position must be inactive so the Expert does not affect text after the tool span. This is therefore a deterministic hard-gated residual Tool Expert, not a 65th native routed expert.

The Tool Expert uses a SwiGLU FFN:

down_proj(silu(gate_proj(x)) * up_proj(x))

• gate_proj / up_proj are initialized with small random values; down_proj is zero-initialized so the initial residual is zero.
• Weights are kept in fp16; internal linear computation runs in fp32 and the result is returned in the original dtype.
• The base model, router, original experts, and embeddings are frozen.

This design targets low-cost protocol specialization. It does not establish that compressed schemas preserve tool-call quality; that must be measured by the paired evaluation protocol.

Quick Start

Install dependencies:

pip install -r requirements.txt

Colab or local CUDA training

Edit cfg in train.py, prepare train.jsonl and val.jsonl, then run:

python train.py

The default local/Colab entry point uses 5% of each split as a smoke test.

Aliyun PAI-DSW training

Set the ModelScope credential in the environment. It is never accepted from a source-code default:

export MODELSCOPE_API_TOKEN="..."
bash setup_aliyun.sh
python train_aliyun.py

A complete local model cache works without the environment variable. Indexed checkpoints are accepted only when every referenced shard is present and non-empty.

Evaluation

python eval.py \
  --data /path/to/eval.jsonl \
  --checkpoint /path/to/best_tool_expert.pt

Tests

Pure-Python routing, credential, cache, and evaluation tests run without Torch, Transformers, ModelScope, or network access. Torch-backed wrapper and checkpoint tests are skipped explicitly when Torch is unavailable.

CUDA_VISIBLE_DEVICES='' \
HF_HUB_OFFLINE=1 \
TRANSFORMERS_OFFLINE=1 \
python3 -m unittest discover -s tests -v

Training Data

train.jsonl and val.jsonl contain one text field per record:

{"text":"User request\n<|tool_expert|>{\"name\":\"weather\",\"arguments\":{\"city\":\"Berlin\"}}<|/tool_expert|>"}

Only tokens inside the marked tool span contribute to language-model loss. The ON/OFF markers, prompt, and padding are excluded, so unmarked natural-language tokens do not dilute the tool-protocol gradient.

Checkpoints

All step, epoch, best, and final files use schemalite-moe format version 1. Each file contains:

• Tool Expert state;
• both input marker rows;
• both output marker rows;
• model, vocabulary, marker, expert, and target-layer metadata;
• optimizer, scheduler, epoch, global step, and best validation loss.

Evaluation validates metadata before scoring group C. Legacy bare Expert weights are detectable but are rejected for strict A/B/C comparison because they cannot restore complete inference state.

Evaluation

Each eval.jsonl record must pair the same target across three contexts:

{
  "full_schema_prompt": "Request plus complete tool schema",
  "brief_prompt": "The same request plus compressed tool description",
  "target_tool_call": "{\"name\":\"weather\",\"arguments\":{\"city\":\"Berlin\"}}",
  "expected_tool_call": {
    "name": "weather",
    "arguments": {"city": "Berlin"}
  }
}

expected_tool_call is optional. Without it, generated tool-call accuracy is omitted rather than inferred from loss.

Command:

python eval.py \
  --data /path/to/eval.jsonl \
  --checkpoint /path/to/best_tool_expert.pt

Groups:

• A: untouched base model with the full-schema prompt.
• B: untouched base model with the brief prompt.
• C: restored Tool Expert model with the brief prompt and ON marker.

Reported teacher-forced metrics cover only target tool-call tokens: token count, summed/mean NLL, perplexity, and token accuracy. Generated accuracy is exact normalized JSON equality. Context compression uses actual tokenized prompt lengths, not fixed estimates.

Generation uses full-prefix recomputation so the deterministic routing mask is correct at every autoregressive step. This is slower than KV caching but avoids stale mutable masks.

Project Structure

Path	Purpose
schemalite_moe/routing.py	Pure-Python causal routing state
schemalite_moe/core.py	Tool Expert, residual wrapper, final-layer injection
schemalite_moe/training.py	Shared local/Colab/Aliyun training loop
schemalite_moe/checkpoint.py	Versioned complete checkpoints
schemalite_moe/evaluation.py	Paired A/B/C metrics and generation
train.py	Local or Colab 5% smoke-training entry point
train_aliyun.py	Aliyun full-data entry point
eval.py	A/B/C evaluation CLI
tests/	Offline standard-library unit tests

Known Limits

• The Tool Expert is externally activated; it does not learn whether a tool is needed.
• A one-final-layer Expert may learn serialization and protocol patterns more readily than complex multi-step tool selection.
• Full-prefix generation is intentionally correctness-first and slower.
• No benchmark result is bundled with this repository.
• A credential that has ever appeared in public Git history must be rotated even after the source line and old repository are removed.

License

MIT License — see LICENSE.