Skip to content

Fix DeepSeek4HyperHead instantiation and usage in NNXDecoder #4775

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
olufiyin19 wants to merge 1 commit into from
+10 −13
Closed

Fix DeepSeek4HyperHead instantiation and usage in NNXDecoder #4775

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/mock_failure_log.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
Traceback (most recent call last):
File "train.py", line 42, in <module>
import maxtext
File "/usr/local/google/home/fiyinbenstowe/Desktop/Project/maxtext/src/maxtext/layers/normalizations.py", line 72
mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
^
SyntaxError: invalid syntax
16 changes: 3 additions & 13 deletions src/maxtext/layers/nnx_decoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -438,13 +438,6 @@ def __init__(
self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4
self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL

if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4:
self.hc_head = mhc.DeepSeek4HyperHead(
config=config,
mesh=self.mesh,
rngs=self.rngs,
)

self._init_decoder_layers(decoder_block_classes, rngs, mesh)

def _init_decoder_layers(self, decoder_block_classes, rngs, mesh):
Expand Down Expand Up @@ -1974,16 +1967,13 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in):

assert isinstance(y, jax.Array)

# After the final transformer layer, `y` holds the raw, un-normalized hidden state.
# After the final transformer layer, `y` holds the raw, un-normalized hidden state.
Comment on lines +1970 to 1971

@gemini-code-assist gemini-code-assist Bot Aug 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The comment explaining that y holds the raw, un-normalized hidden state is duplicated. Please remove the redundant duplicate comment.

Suggested change
# After the final transformer layer, `y` holds the raw, un-normalized hidden state.
# After the final transformer layer, `y` holds the raw, un-normalized hidden state.
# After the final transformer layer, `y` holds the raw, un-normalized hidden state.

if getattr(cfg, "mhc_expansion_rate", 1) > 1:
if cfg.decoder_block == DecoderBlockType.DEEPSEEK4:
hidden_state = self.hc_head(y)
else:
# (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim)
hidden_state = mhc_reduce(y)
# (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim)
hidden_state = mhc_reduce(y)
else:
hidden_state = y

# When invoking from vLLM with RPA attention, logit computation is deferred to a later stage.
if cfg.attention in ("vllm_rpa", "vllm_batched_rpa"):
logits = None
Expand Down
Loading