Train MoE Experts and LoRA Specialists

What this part does

For attention and MLP LoRA with an external prompt classifier, use the separate prompt-routed adapter workflow. It trains one model-level adapter without MLP conversion and supports adapter-only deployment and resumable local Trainer checkpoints. The forced-MLP constraints below continue to apply to ForcedPathExpertTrainer, not that new Trainer.

Training controls which new expert parameters may change and which execution path produces their gradients. Every supported workflow mechanically freezes the anchor and shared transformer, proves the exact optimizer-owned parameter set, and restores converted layers to anchor-only execution after training.

Forced-path full-MLP training, including optional preservation objectives, feeds the validated expert-package lifecycle. Independent, local-mixture, and refinement workflows are composable training paths. Loss-free balancing and expert-choice assignment remain research primitives with no automatic Trainer-to-runtime connection.

Where it fits

Training begins after a backend has converted the selected MLP layers and expert objects have been created:

1. the backend exposes converted RoutedMLP wrappers;
2. expert bindings declare stable layer and expert ownership;
3. a training topology attaches one path or a local pair temporarily;
4. Hugging Face Trainer and Accelerate run the ordinary optimization loop;
5. cleanup clears gradients, freezes artifacts, and restores forced-anchor execution;
6. package export, router fitting, installation, and inference happen later.

The training package does not partition datasets automatically, generate replay teacher logits, fit a global inference router, export packages, or decide serving-time capacity.

Implemented approaches

Forced-path single-expert training

ForcedPathExpertTrainer attaches one independently owned expert at every converted layer. Every token follows that expert path; no router score, selection policy, anchor comparison, or sparse grouping participates in the forward pass.

Before Trainer exists, the implementation rejects incomplete layer coverage, parameterless experts, shared expert/model storage, and cross-layer expert sharing. It then freezes the entire model, enables only the parameters declared by each expert, and compares exact parameter object identities with the expected optimizer set.

Independent latent experts

train_independent_expert_set() composes two or more forced-path runs. Each expert may receive its own data partition, seed, schedule, collator, and preservation settings. Runs execute sequentially against the same immutable anchor, and only the current expert path is registered on the model.

This is independent optimization, not token routing or one joint loss. The coordinator validates the entire expert-by-layer rectangle before the first optimizer step so a malformed later run cannot be discovered only after an earlier expert has trained.

Local two-expert mini-MoE

LocalMiniMoETrainer attaches exactly two package-local experts and fresh affine router heads. Both experts execute every specialization token, and a differentiable score softmax combines their outputs. Small deterministic opposite router biases break identical anchor-copy symmetry without randomness.

The anchor is not a candidate inside this temporary local mixture. After the one supported Trainer run, cleanup freezes both experts and the heads, releases captured probability tensors, and restores the exact anchor route.

Mixture-aware refinement

refine_independent_expert_set() links the two stages deliberately:

Two isolated TRN-002 runs
    → validated completed report
    → fresh local affine heads
    → joint TRN-003 refinement

The source report must prove the same two expert IDs, stable order, converted layers, and completed optimizer progress. The second stage then updates the caller-owned trained experts and fresh local heads together on a mixture dataset.

Optional objectives and balancing

ExpertObjectiveConfig adds replay KL and normalized weight drift to forced-path training without changing optimizer ownership. Replay uses precomputed, detached immutable-anchor logits; drift compares live expert parameters with trainer-construction snapshots.

Local mini-MoE training can add a Switch-style auxiliary balancing loss through load_balancing_weight. This changes training pressure but not computation: both experts still execute for every token.

LossFreeBiasBalancer and ExpertChoiceRouting address sparse-training research questions separately. The former updates persistent, gradient-free selection biases from completed-batch loads. The latter creates equal-capacity expert token buckets. Neither primitive is wired into LocalMiniMoETrainer or an expert dispatcher.

How it works

Establish ownership before mutation

Bindings are normalized into deterministic layer order. The framework checks complete converted-layer coverage and rejects parameter or storage aliases across the model, experts, router heads, and feature transformation before attaching a route.

Let Trainer own the ordinary loop

The specialized trainers preserve Hugging Face behavior for data loading, loss computation, accumulation, clipping, optimizer and scheduler stepping, callbacks, and device placement. They constrain unsafe features: save_strategy must be "no", full-model saving and Hub push are rejected, and checkpoint resume remains unsupported until expert-only optimizer state has a versioned format.

Restore the anchor on every exit

Successful and failed training runs clear gradients, freeze external artifacts, detach temporary experts, and return every converted wrapper to forced_anchor. Cleanup does not roll back expert weights after an optimizer step; failure safety protects model routing and ownership, not transactional expert history.

Mathematical intuition

Minimal usage

Train one bound expert path with ordinary Hugging Face arguments:

trainer = ForcedPathExpertTrainer(
    model=model,
    backend=Qwen3Backend(),
    expert_bindings=bindings,
    args=TrainingArguments(
        output_dir="./expert-run",
        save_strategy="no",
    ),
    train_dataset=dataset,
)
trainer.train()

The model must already be converted, and bindings must contain one independently owned expert for every converted layer.

Verified implementation details

trainer = ForcedPathExpertTrainer(
    model=model,
    backend=backend,
    expert_bindings=[ForcedExpertBinding(3, "code:0", expert)],
    objective_config=ExpertObjectiveConfig(
        replay_kl_weight=0.5,
        drift_weight=0.1,
        replay_temperature=2.0,
    ),
    args=training_arguments,
    train_dataset=dataset,
)
trainer.train()

report = train_independent_expert_set(
    model,
    backend=backend,
    runs=[code_run, code_one_run],
)

assert report.expert_ids == ("code:0", "code:1")
assert report.total_global_steps == 2
assert wrapper.routing_mode == "forced_anchor"

trainer = LocalMiniMoETrainer(
    model=model,
    backend=backend,
    expert_bindings=bindings,
    router_feature_extractor=IdentityRouterFeatureExtractor(16),
    router_heads=heads,
    config=LocalMiniMoETrainingConfig(
        load_balancing_weight=0.05,
    ),
    args=training_arguments,
    train_dataset=dataset,
)
trainer.train()

report = refine_independent_expert_set(
    model,
    backend=backend,
    source_training=independent_report,
    expert_bindings=local_pair_bindings,
    router_feature_extractor=features,
    router_heads=fresh_heads,
    config=LocalMiniMoETrainingConfig(
        load_balancing_weight=0.05,
    ),
    args=training_arguments,
    trainer_kwargs={"train_dataset": mixture_dataset},
)

Needle suffix training

The separate Needle workflow uses BranchTrainer with interchangeable optimizer and objective strategies. train_specialist() updates one response suffix; train_router() fits the affine scorer and optional router suffix. train_generative_router() trains only the private suffix to emit route tool calls from recent conversation context, keeping response branches and the affine head unchanged.

See the generative training example for consistent prompt rendering and label masking. Evaluate route quality again after CQ4 export; floating-reference fitting does not validate native outcomes.

Current limitations

• TRN-001 and TRN-002 train complete expert paths; they do not discover routes or partition data automatically.
• Independent expert runs are sequential in one process. Concurrent training requires separate models and processes.
• TRN-003 and TRN-004 support exactly two experts with complete identical layer coverage and fresh, zero-initialized affine heads.
• Local training always executes both experts. It provides no anchor abstention, sparse capacity, or inference-time compute saving.
• TRN-005 currently composes with forced-path training, not local mini-MoE training. Full-vocabulary replay logits and full expert drift snapshots can be storage intensive.
• Full-model checkpoints, Hub push, model reinitialization, and checkpoint resume are rejected. Expert packages are exported through a separate safetensors workflow.
• Failure cleanup restores safe routing and freezing but does not roll back partially optimized external expert weights.
• BAL-004 has no automatic Trainer-step hook or distributed load reduction.
• BAL-005 plans sparse token assignments but does not execute or combine experts and does not support autoregressive decoding.
• Mixed precision, multi-device, distributed, and gradient-checkpointing acceptance remain future validation work.

References

TRN-001 through TRN-004 are project-specific framework orchestration rather than paper-derived training algorithms. Their temporary execution paths, ownership assertions, and cleanup rules implement the repository’s immutable anchor design.

The training theory README contains the complete objectives, algorithms, Hugging Face boundary, tests, and documented deviations. The architecture blueprint defines forced, independent, local-mixture, and refinement choices in the larger post-training workflow.

Related pages

• Read experts for the trainable representations consumed by these workflows.
• Read policies to distinguish local training mixtures from anchored inference selection and capacity fallback.
• Return to configuration and modeling for layer conversion and the forced-anchor baseline.
• Read calibration for post-training feature banks, utility labels, router fitting, temperatures, and thresholds.
• Read registry for expert-only package export, compatibility metadata, installation, upgrade, and rollback.