Feat: Compatible with Gamme4 MoE structure

[waiting-xia]

Gemma4 MoE DSpark 适配改动整理

Summary

This change adds Gemma4 MoE block support to the DSpark Gemma4 draft model.

Previously, Gemma4DSparkDecoderLayer explicitly rejected Gemma4 configs with enable_moe_block=True, so DSpark could not be used with Gemma4 / Tima MoE target models.

The core changes are limited to:

deepspec/modeling/dspark/gemma4/modeling.py
deepspec/modeling/dspark/gemma4/config.py

This document intentionally excludes Tima26-specific experiment configs, training launchers, BF8/OOM workarounds, NaN guards, target-cache/eval debugging, and smoke-test artifacts, because those are operational/debug changes rather than the core Gemma4 MoE adaptation.

---

1. Add Gemma4 MoE block support in the DSpark decoder layer

File

deepspec/modeling/dspark/gemma4/modeling.py

Previous behavior

The Gemma4 DSpark decoder layer used to reject MoE-enabled Gemma4 configs:

assert not bool(config.enable_moe_block), (
    "Gemma4 DSpark prototype does not support Gemma4 MoE blocks yet."
)

This meant that target models whose Gemma4 text config had:

enable_moe_block = True

could not be used directly with Gemma4DSparkModel.

---

New behavior

The decoder layer now reads the MoE flag from the config:

self.enable_moe_block = bool(config.enable_moe_block)

and conditionally creates Gemma4 MoE modules when MoE is enabled.

---

New imports

The DSpark Gemma4 model now imports Gemma4 MoE components from Transformers:

from transformers.models.gemma4.modeling_gemma4 import (
    Gemma4TextExperts,
    Gemma4TextMLP,
    Gemma4TextRouter,
)

Gemma4TextMLP is the dense feed-forward branch, while Gemma4TextRouter and Gemma4TextExperts implement the MoE branch.

---

Decoder layer initialization

When config.enable_moe_block=True, the decoder layer now constructs the router, experts, and MoE-specific RMSNorm modules:

self.mlp = Gemma4TextMLP(config, layer_idx)

if self.enable_moe_block:
    self.router = Gemma4TextRouter(config)
    self.experts = Gemma4TextExperts(config)
    self.post_feedforward_layernorm_1 = Gemma4RMSNorm(
        config.hidden_size,
        eps=config.rms_norm_eps,
    )
    self.post_feedforward_layernorm_2 = Gemma4RMSNorm(
        config.hidden_size,
        eps=config.rms_norm_eps,
    )
    self.pre_feedforward_layernorm_2 = Gemma4RMSNorm(
        config.hidden_size,
        eps=config.rms_norm_eps,
    )

---

Decoder layer forward pass

The forward path now executes the MoE branch after the dense MLP branch when MoE is enabled:

hidden_states = self.mlp(hidden_states)

if self.enable_moe_block:
    hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states)
    hidden_states_flat = residual.reshape(-1, residual.shape[-1])
    _, top_k_weights, top_k_index = self.router(hidden_states_flat)
    hidden_states_2 = self.pre_feedforward_layernorm_2(hidden_states_flat)
    hidden_states_2 = self.experts(hidden_states_2, top_k_index, top_k_weights)
    hidden_states_2 = hidden_states_2.reshape(residual.shape)
    hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2)
    hidden_states = hidden_states_1 + hidden_states_2

The output then continues through the existing post-feed-forward normalization and residual path:

hidden_states = self.post_feedforward_layernorm(hidden_states)
hidden_states = residual + hidden_states
return hidden_states * self.layer_scalar

---

2. Propagate Gemma4 MoE configuration into the DSpark draft config

File

deepspec/modeling/dspark/gemma4/config.py

Required Gemma4 text config fields

The Gemma4 draft config builder requires the target text config to expose MoE-related fields, including:

"enable_moe_block"

This ensures the draft config knows whether the target Gemma4 text stack uses MoE blocks.

---

Support direct Gemma4 text configs

get_gemma4_text_config() now accepts either a top-level Gemma4 config or a Gemma4 text config directly:

if target_config.model_type in ("gemma4_text", "gemma4_unified_text"):
    return copy.deepcopy(target_config)

This makes the config builder work with both forms:

gemma4 / gemma4_unified
gemma4_text / gemma4_unified_text

---

Support top_k_experts override

The draft config builder now propagates model_args.top_k_experts into the Gemma4 draft config when provided:

if "top_k_experts" in model_args:
    draft_config.top_k_experts = int(model_args.top_k_experts)

This keeps the expert-selection parameter configurable for Gemma4 MoE DSpark drafts. Users can keep the target model's original value, or override it based on their hardware budget and experiment needs. The model implementation should not hard-code a specific top_k_experts value.

---

3. What this change does not do

This change adds structural support for Gemma4 MoE blocks in the DSpark draft model, but it does not implement target-to-draft MoE weight transfer.

In particular, there is no current logic that copies target weights for:

router.*
experts.*
mlp.*

The existing target-weight initialization path only handles embedding and LM head weights elsewhere in the trainer code. MoE router/expert parameters are still initialized as part of the draft model and learned during DSpark training.

---

Core conclusion

The Gemma4 MoE DSpark adaptation consists of two essential code changes:

1. deepspec/modeling/dspark/gemma4/modeling.py

   • Remove the MoE-disabled assertion.
   • Instantiate Gemma4TextRouter and Gemma4TextExperts when enable_moe_block=True.
   • Add the MoE expert branch to the decoder layer forward pass.

2. deepspec/modeling/dspark/gemma4/config.py

   • Preserve enable_moe_block from the Gemma4 text config.
   • Accept direct Gemma4 text configs.
   • Allow top_k_experts to be overridden for the DSpark draft config.

Everything else, such as Tima26-specific configs, BF8, cache generation, eval changes, NaN handling, and smoke-test scripts, is operational/debug work and is not part of the minimal Gemma4 MoE adaptation.