deepseek-ai / deepseek-ai/DeepSpec

Feat: Compatible with Gamme4 MoE structure

Open
#60 0 comments 3 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
7.1k
Forks
667
PR merge metrics
No merged PRs in 30d

Description

# 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:

```text
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

```text
deepspec/modeling/dspark/gemma4/modeling.py
```

### Previous behavior

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

```python
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:

```text
enable_moe_block = True
```

could not be used directly with `Gemma4DSparkModel`.

---

### New behavior

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

```python
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:

```python
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:

```python
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:

```python
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:

```python
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

```text
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:

```python
"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:

```python
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:

```text
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:

```python
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:

```text
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.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.