From b3958b9a96119577e94b5c720f44545ae8c5f001 Mon Sep 17 00:00:00 2001 From: Snehal Verma Date: Thu, 2 Jul 2026 01:36:43 +0000 Subject: [PATCH] add DeepSeek4 HyperHead, check for DEEPSEEK4 decoder block type in moe.py, and match maxtext sinkhorn implementation to HF reference. --- src/maxtext/configs/models/deepseek4-284b.yml | 5 + src/maxtext/layers/mhc.py | 91 ++- src/maxtext/layers/moe.py | 15 +- tests/unit/deepseek_v4_vs_reference_test.py | 565 +++++++++++++++++- 4 files changed, 656 insertions(+), 20 deletions(-) diff --git a/src/maxtext/configs/models/deepseek4-284b.yml b/src/maxtext/configs/models/deepseek4-284b.yml index 598bbd9c1c..708a36e522 100644 --- a/src/maxtext/configs/models/deepseek4-284b.yml +++ b/src/maxtext/configs/models/deepseek4-284b.yml @@ -49,6 +49,10 @@ num_experts_per_tok: 6 mlp_activations_limit: 10 shared_experts: 1 routed_score_func: "sqrtsoftplus" +norm_topk_prob: true +routed_bias: true +routed_scaling_factor: 1.5 + # --- Attention configuration --- attention_type: 'compressed' @@ -62,3 +66,4 @@ rope_type: "default" rope_max_timescale: 10000 # Main RoPE theta compressed_rope_max_timescale: 160000 # Compressed RoPE theta max_position_embeddings: 1048576 +original_max_position_embeddings: 65536 diff --git a/src/maxtext/layers/mhc.py b/src/maxtext/layers/mhc.py index 03e0eb6eee..2a6cd33dd8 100644 --- a/src/maxtext/layers/mhc.py +++ b/src/maxtext/layers/mhc.py @@ -24,7 +24,8 @@ from jax.sharding import Mesh from maxtext.common.common_types import Array, Config from maxtext.common.common_types import HyperConnectionType -from maxtext.layers.initializers import default_bias_init, default_scalar_init, nd_dense_init +from maxtext.layers.initializers import default_bias_init, default_scalar_init, nd_dense_init, variable_to_logically_partitioned +from maxtext.layers import nnx_wrappers from maxtext.layers.normalizations import RMSNorm @@ -61,21 +62,18 @@ def sinkhorn(t, iters=20): # Use float32 precision for numerical stability during normalization initial_dtype = t.dtype t = t.astype(jnp.float32) + eps = 1e-5 - # Column-wise normalization (axis=-2) - positive and sum up to 1 across columns - # Equivalent to t = exp(t) / jnp.sum(jnp.exp(t), axis=-2) - t = jax.nn.softmax(t, axis=-2) + t = jax.nn.softmax(t, axis=-1) + eps + t = t / (jnp.sum(t, axis=-2, keepdims=True) + eps) def body_fun(i, val): - # L1 Normalization: val / sum(val) with clipping of denominator - # Normalize rows (axis -1) - val = val / jnp.clip(jnp.sum(val, axis=-1, keepdims=True), min=1e-12) - # Normalize columns (axis -2) - val = val / jnp.clip(jnp.sum(val, axis=-2, keepdims=True), min=1e-12) + val = val / (jnp.sum(val, axis=-1, keepdims=True) + eps) + val = val / (jnp.sum(val, axis=-2, keepdims=True) + eps) return val # Use lax.fori_loop for an efficient, JIT-friendly loop - t = jax.lax.fori_loop(0, iters, body_fun, t) + t = jax.lax.fori_loop(0, iters - 1, body_fun, t) return t.astype(initial_dtype) @@ -224,7 +222,7 @@ def res_mapping(self, x: Array): output = sinkhorn(intermediate, self.sinkhorn_iterations) return output - def mapping(self, x: Array, alpha_scale: Array, alpha: Array, beta: Array, scale: int): + def mapping(self, x: Array, alpha_scale: Array, alpha: Array, beta: Array, scale: float, eps: float = 0.0): """Helper function for both pre and post mappings.""" # In MaxText, we match weight precision to activations before Matmul alpha = jnp.asarray(alpha, self.dtype) @@ -233,7 +231,7 @@ def mapping(self, x: Array, alpha_scale: Array, alpha: Array, beta: Array, scale # Apply projection: (b, s, k*d) @ (k*d, k) -> (b, s, k) h = jnp.einsum("bsm,mk -> bsk", x, alpha, precision=self.matmul_precision) intermediate = alpha_scale * h + beta[None, None, :] - output = scale * jax.nn.sigmoid(intermediate) + output = scale * jax.nn.sigmoid(intermediate) + eps return output def __call__( @@ -269,6 +267,7 @@ def __call__( self.pre_alpha[...], self.pre_beta[...], 1.0, + eps=1e-5, ) layer_input = jnp.einsum("bskd,bsk -> bsd", x, pre_mapping, precision=self.matmul_precision) @@ -307,3 +306,71 @@ def __call__( res_mapping = self.res_mapping(norm_x) res_out = jnp.einsum("bskd,bskm -> bsmd", x, res_mapping, precision=self.matmul_precision) return res_out + post_out, metadata + + +class DeepSeek4HyperHead(nnx.Module): + """Final HC-stream collapse; used by DeepSeek V4 before the shared RMSNorm.""" + + def __init__( + self, + config: Config, + mesh: Mesh, + rngs: nnx.Rngs, + ): + self.config = config + self.mesh = mesh + self.rngs = rngs + self.dtype = config.dtype + self.weight_dtype = config.weight_dtype + self.mhc_expansion_rate = config.mhc_expansion_rate + self.emb_dim = config.emb_dim + self.eps = 1e-6 + + # Weight matrices + weight_init = nd_dense_init(1.0, "fan_in", "normal") + self.hc_fn = nnx.Param( + weight_init( + rngs.params(), + (self.mhc_expansion_rate * self.emb_dim, self.mhc_expansion_rate), + self.weight_dtype, + in_axis=0, + out_axis=1, + ), + out_sharding=("activation_embed", None), + ) + self.hc_base = nnx.Param( + default_bias_init(rngs.params(), (self.mhc_expansion_rate,), self.weight_dtype), + out_sharding=(None,), + ) + self.hc_scale = nnx.Param( + default_scalar_init(rngs.params(), (1,), self.weight_dtype), + out_sharding=(None,), + ) + + def __call__(self, x: Array) -> Array: + # x shape: [batch, length, k, d] + b, s, k, d = x.shape + assert k == self.mhc_expansion_rate + assert d == self.emb_dim + + flat = jnp.reshape(x, (b, s, k * d)) + flat_f32 = flat.astype(jnp.float32) + variance = jnp.mean(jnp.square(flat_f32), axis=-1, keepdims=True) + flat_norm = flat_f32 * jax.lax.rsqrt(variance + self.eps) + + hc_fn = jnp.asarray(self.hc_fn[...], jnp.float32) + hc_base = jnp.asarray(self.hc_base[...], jnp.float32) + hc_scale = jnp.asarray(self.hc_scale[...], jnp.float32) + + mixes = jnp.einsum("bsm,mk->bsk", flat_norm, hc_fn, precision=jax.lax.Precision(self.config.matmul_precision)) + pre = jax.nn.sigmoid(mixes * hc_scale[None, None, :] + hc_base[None, None, :]) + self.eps + + x_f32 = x.astype(jnp.float32) + out = jnp.sum(pre[:, :, :, None] * x_f32, axis=2) + return out.astype(self.dtype) + + +DeepSeek4HyperHeadToLinen = nnx_wrappers.to_linen_class( + DeepSeek4HyperHead, + base_metadata_fn=variable_to_logically_partitioned, +) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index fdea21981c..5d297588c4 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -699,14 +699,15 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): else: top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) - if self.config.decoder_block == ctypes.DecoderBlockType.DEEPSEEK: + if self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4): top_k_weights = self.deepseek_scale_weights(top_k_weights) - elif self.config.decoder_block not in (ctypes.DecoderBlockType.LLAMA4, ctypes.DecoderBlockType.GEMMA4): - top_k_weights = jax.nn.softmax(top_k_weights.astype(jnp.float32), axis=-1).astype(self.dtype) + else: + if self.config.decoder_block not in (ctypes.DecoderBlockType.LLAMA4, ctypes.DecoderBlockType.GEMMA4): + top_k_weights = jax.nn.softmax(top_k_weights.astype(jnp.float32), axis=-1).astype(self.dtype) - # Normalization of router weights (e.g. used by Qwen3, Gemma4). - if self.config.norm_topk_prob: - top_k_weights /= top_k_weights.sum(axis=-1, keepdims=True) + # Normalization of router weights (e.g. used by Qwen3, Gemma4). + if self.config.norm_topk_prob: + top_k_weights /= top_k_weights.sum(axis=-1, keepdims=True) return top_k_weights, top_k_indices @@ -793,7 +794,7 @@ def apply_ffn_activation(self, layer_w0, layer_w1): layer_act = self.activation_fn(layer_w0 * 1.702) glu = jnp.multiply(layer_w0, layer_act) intermediate_layer = jnp.multiply(glu, (layer_w1 + 1)) - elif self.config.decoder_block == ctypes.DecoderBlockType.DEEPSEEK and self.config.mlp_activations_limit > 0.0: + elif self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4) and self.config.mlp_activations_limit > 0.0: # DeepSeek V4 uses bounds to clip the SwiGLU activations layer_w0 = jnp.clip(layer_w0, min=None, max=self.config.mlp_activations_limit) layer_w1 = jnp.clip(layer_w1, min=-self.config.mlp_activations_limit, max=self.config.mlp_activations_limit) diff --git a/tests/unit/deepseek_v4_vs_reference_test.py b/tests/unit/deepseek_v4_vs_reference_test.py index 0b75aa9ff4..839ae0ce98 100644 --- a/tests/unit/deepseek_v4_vs_reference_test.py +++ b/tests/unit/deepseek_v4_vs_reference_test.py @@ -587,7 +587,7 @@ def _run_e2e_test(self, layer_type, is_packed=False): self._copy_linear(mt_attn.csa_compressor.indexer.q_proj, ref_attn.compressor.indexer.q_b_proj) self._copy_linear(mt_attn.csa_compressor.indexer.kv_proj, ref_attn.compressor.indexer.kv_proj) self._copy_linear(mt_attn.csa_compressor.indexer.gate_proj, ref_attn.compressor.indexer.gate_proj) - self._copy_linear(mt_attn.csa_compressor.indexer.weights_proj, ref_attn.compressor.indexer.weights_proj) + self._copy_linear(mt_attn.csa_compressor.indexer.weights_proj, ref_attn.compressor.indexer.scorer.weights_proj) mt_attn.csa_compressor.indexer.position_bias.value = jnp.array( ref_attn.compressor.indexer.position_bias.data.numpy() ) @@ -951,6 +951,569 @@ def test_swiglu_clamp(self): # Validate that both clamped outputs match identically np.testing.assert_allclose(mx_out, pt_out.numpy(), rtol=1e-5, atol=1e-5) +class DeepSeekV4MoERouterTest(unittest.TestCase): + + def setUp(self): + self.batch_size = 2 + self.seq_len = 32 + self.hidden_dim = 4096 + self.num_experts = 8 + self.num_experts_per_tok = 3 + self.vocab_size = 129280 + + self.pt_config = DeepseekV4Config( + hidden_size=self.hidden_dim, + num_local_experts=self.num_experts, + num_experts_per_tok=self.num_experts_per_tok, + routed_scaling_factor=2.0, + scoring_func="sqrtsoftplus", + vocab_size=self.vocab_size, + ) + + config_arguments = { + "per_device_batch_size": 1.0, + "run_name": "test", + "enable_checkpointing": False, + "base_emb_dim": self.hidden_dim, + "num_experts": self.num_experts, + "topk_routing_group": self.num_experts_per_tok, + "routed_scaling_factor": 2.0, + "routed_score_func": "sqrtsoftplus", + "routed_bias": True, + "n_routing_groups": -1, + "vocab_size": self.vocab_size, + "first_num_hash_layers": 3, + "decoder_block": "deepseek4", + "model_name": "deepseek4-284b", + "attention": "dot_product", + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "override_model_config": True, + "skip_jax_distributed_system": True, + } + argv = [sys.argv[0], "src/maxtext/configs/base.yml"] + self.mx_config = pyconfig.initialize(argv, **config_arguments) + + devices = np.array(jax.devices()[:1]) + self.mesh = jax.sharding.Mesh(devices, ("tensor",)) + self.rngs = nnx.Rngs(0) + + def test_hash_router(self): + pt_router = DeepseekV4HashRouter_PT(self.pt_config) + # Explicitly initialize PyTorch weights since torch.empty leaves garbage in memory, + # which causes NaN/Inf drift between PyTorch and MaxText/XLA execution. + torch.nn.init.normal_(pt_router.weight, std=0.02) + + # Hash Router operates deterministically based on input_ids via a frozen tid2eid lookup table. + # In practice, this table is pre-computed (e.g. by K-Means on the dataset) and loaded statically. + # For this parity test, we randomly initialize the lookup table on the PyTorch side + # and explicitly sync it to the MaxText side to ensure both routers route the exact same way. + pt_tid2eid = torch.randint(0, self.num_experts, (self.vocab_size, self.num_experts_per_tok)) + pt_router.tid2eid.copy_(pt_tid2eid) + + mx_moe = RoutedMoE( + config=self.mx_config, + num_experts=self.num_experts, + num_experts_per_tok=self.num_experts_per_tok, + mesh=self.mesh, + kernel_init=initializers.nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed_moe", "mlp_moe", None), + rngs=self.rngs, + is_hash_routing=True, # Hash layer + ) + + # Sync weights + mx_moe.tid2eid.value = jnp.array(pt_router.tid2eid.numpy(), dtype=jnp.float32) + mx_moe.gate.kernel.value = jnp.array(pt_router.weight.detach().numpy()).T + + hidden_states = torch.randn(self.batch_size, self.seq_len, self.hidden_dim) + input_ids = torch.randint(0, self.vocab_size, (self.batch_size, self.seq_len)) + + # PT forward + _, pt_weights, pt_indices = pt_router(hidden_states, input_ids) + + # MaxText forward + gate_logits, pre_bias_logits = mx_moe.gate(jnp.array(hidden_states.numpy())) + mx_weights, mx_indices = mx_moe.get_topk( + gate_logits, pre_bias_logits, rngs=self.rngs, input_ids=jnp.array(input_ids.numpy()) + ) + + # --- Assertion Logic for Hash Router --- + # PyTorch returns flat tensors: (batch * seq_len, top_k) + # MaxText returns structured tensors: (batch, seq_len, top_k) + # We must explicitly reshape PyTorch outputs to match MaxText's nested sequence structure. + pt_indices_reshaped = pt_indices.numpy().reshape(self.batch_size, self.seq_len, -1) + pt_weights_reshaped = pt_weights.detach().numpy().reshape(self.batch_size, self.seq_len, -1) + weights_max_diff = np.max(np.abs(mx_weights - pt_weights_reshaped)) + weights_mean_diff = np.mean(np.abs(mx_weights - pt_weights_reshaped)) + print(f"MOE HASH ROUTER WEIGHTS PARITY - MAX ABS DIFF: {weights_max_diff:.6e}, MEAN ABS DIFF: {weights_mean_diff:.6e}") + np.testing.assert_allclose(mx_indices, pt_indices_reshaped, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(mx_weights, pt_weights_reshaped, rtol=1e-2, atol=1e-2) + + def test_topk_router(self): + pt_router = DeepseekV4TopKRouter_PT(self.pt_config) + + # Explicitly initialize PyTorch weights since torch.empty leaves garbage in memory, + # which causes NaN/Inf drift between PyTorch and MaxText/XLA execution. + torch.nn.init.normal_(pt_router.weight, std=0.02) + torch.nn.init.normal_(pt_router.e_score_correction_bias, std=0.02) + + mx_moe = RoutedMoE( + config=self.mx_config, + num_experts=self.num_experts, + num_experts_per_tok=self.num_experts_per_tok, + mesh=self.mesh, + kernel_init=initializers.nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed_moe", "mlp_moe", None), + rngs=self.rngs, + is_hash_routing=False, # TopK layer + ) + + # Sync weights + mx_moe.gate.kernel.value = jnp.array(pt_router.weight.detach().numpy()).T + mx_moe.gate.bias.value = jnp.array(pt_router.e_score_correction_bias.detach().numpy()) + + hidden_states = torch.randn(self.batch_size, self.seq_len, self.hidden_dim) + + # PT forward + _, pt_weights, pt_indices = pt_router(hidden_states) + + # MaxText forward + gate_logits, pre_bias_logits = mx_moe.gate(jnp.array(hidden_states.numpy())) + mx_weights, mx_indices = mx_moe.get_topk(gate_logits, pre_bias_logits, rngs=self.rngs) + + # --- Assertion Logic for TopK Router --- + # PyTorch returns flat tensors: (batch * seq_len, top_k) + # MaxText returns structured tensors: (batch, seq_len, top_k) + # 1. Reshape PyTorch outputs to match MaxText's nested sequence structure. + pt_indices_reshaped = pt_indices.numpy().reshape(self.batch_size, self.seq_len, -1) + pt_weights_reshaped = pt_weights.detach().numpy().reshape(self.batch_size, self.seq_len, -1) + + # 2. Sort both by indices so they can be compared directly. + # jax.lax.top_k and torch.topk resolve exact-value ties differently. + # Because TopK routing weights are summed commutatively during the MoE forward pass, + # only the mathematical *set* of selected experts matters, not their strict sorted order. + mx_sort_idx = np.argsort(mx_indices, axis=-1) + pt_sort_idx = np.argsort(pt_indices_reshaped, axis=-1) + + mx_indices_sorted = np.take_along_axis(np.array(mx_indices), mx_sort_idx, axis=-1) + mx_weights_sorted = np.take_along_axis(np.array(mx_weights), mx_sort_idx, axis=-1) + + pt_indices_sorted = np.take_along_axis(pt_indices_reshaped, pt_sort_idx, axis=-1) + pt_weights_sorted = np.take_along_axis(pt_weights_reshaped, pt_sort_idx, axis=-1) + + weights_max_diff = np.max(np.abs(mx_weights_sorted - pt_weights_sorted)) + weights_mean_diff = np.mean(np.abs(mx_weights_sorted - pt_weights_sorted)) + print(f"MOE TOPK ROUTER WEIGHTS PARITY - MAX ABS DIFF: {weights_max_diff:.6e}, MEAN ABS DIFF: {weights_mean_diff:.6e}") + np.testing.assert_allclose(mx_indices_sorted, pt_indices_sorted, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(mx_weights_sorted, pt_weights_sorted, rtol=1e-2, atol=1e-2) + + +class DeepSeekV4SwiGLUClampTest(unittest.TestCase): + + def test_swiglu_clamp(self): + limit = 10.0 + pt_config = DeepseekV4Config( + hidden_size=4096, + num_local_experts=8, + num_experts_per_tok=3, + intermediate_size=256, + swiglu_limit=limit, + ) + + config_arguments = { + "per_device_batch_size": 1.0, + "run_name": "test", + "enable_checkpointing": False, + "base_emb_dim": 4096, + "num_experts": 8, + "topk_routing_group": 3, + "mlp_activations_limit": limit, + "decoder_block": "deepseek4", + "model_name": "deepseek4-284b", + "attention": "dot_product", + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "override_model_config": True, + "matmul_precision": "highest", + "skip_jax_distributed_system": True, + } + argv = [sys.argv[0], "src/maxtext/configs/base.yml"] + mx_config = pyconfig.initialize(argv, **config_arguments) + + pt_experts = DeepseekV4Experts_PT(pt_config) + + devices = np.array(jax.devices()[:1]) + mesh = jax.sharding.Mesh(devices, ("tensor",)) + rngs = nnx.Rngs(0) + + mx_moe = RoutedMoE( + config=mx_config, + num_experts=2, + num_experts_per_tok=1, + mesh=mesh, + kernel_init=initializers.nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed_moe", "mlp_moe", None), + rngs=rngs, + ) + + # Gate & Up merged matrix in PT + gate_up = torch.randn(1, 4, 256 * 2) * 20.0 # Force large values to trigger the clamping mechanism + + # PyTorch reference executes the standard SwiGLU followed by clamping + # to self.config.swiglu_limit (which translates to mlp_activations_limit in MaxText) + pt_out = pt_experts._apply_gate(gate_up) # pylint: disable=protected-access + + # In MaxText, the gate and up projections are separated mathematically. + # apply_ffn_activation executes the swiglu limit internally during the activation phase. + gate, up = gate_up.chunk(2, dim=-1) + mx_out = mx_moe.apply_ffn_activation(jnp.array(gate.numpy()), jnp.array(up.numpy())) + + # Validate that both clamped outputs match identically + max_diff = np.max(np.abs(mx_out - pt_out.numpy())) + mean_diff = np.mean(np.abs(mx_out - pt_out.numpy())) + print(f"SWIGLU CLAMP PARITY - MAX ABS DIFF: {max_diff:.6e}, MEAN ABS DIFF: {mean_diff:.6e}") + np.testing.assert_allclose(mx_out, pt_out.numpy(), rtol=1e-5, atol=1e-5) + + +from transformers.models.deepseek_v4.modeling_deepseek_v4 import DeepseekV4DecoderLayer as DeepseekV4DecoderLayer_PT +from maxtext.models.deepseek4 import DeepSeek4DecoderLayer + +class DeepSeekV4ConversionMappingTest(unittest.TestCase): + """Tests to validate weight conversion mappings from PARAM_MAPPING.""" + + def setUp(self): + self.batch_size = 2 + self.seq_len = 32 + self.hidden_dim = 4096 + self.num_heads = 64 + self.head_dim = 512 + self.q_lora_rank = 1024 + self.o_groups = 8 + self.o_lora_rank = 1024 + self.qk_rope_head_dim = 64 + self.partial_rotary_factor = self.qk_rope_head_dim / self.head_dim + self.vocab_size = 129280 + + self.pt_config = DeepseekV4Config( + hidden_size=self.hidden_dim, + num_attention_heads=self.num_heads, + num_key_value_heads=1, + head_dim=self.head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.head_dim, + o_groups=self.o_groups, + o_lora_rank=self.o_lora_rank, + layer_types=[ + "sliding_attention", + "sliding_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", + "compressed_sparse_attention", + ], + num_hidden_layers=7, + num_nextn_predict_layers=0, + num_local_experts=8, + num_experts_per_tok=3, + vocab_size=self.vocab_size, + ) + + config_arguments = { + "model_name": "deepseek4-tiny", + "override_model_config": True, + "per_device_batch_size": 1, + "matmul_precision": "highest", + "megablox": False, + "sparse_matmul": False, + "dtype": "float32", + "weight_dtype": "float32", + "skip_jax_distributed_system": True, + } + argv = [sys.argv[0], "src/maxtext/configs/base.yml"] + self.mx_config = pyconfig.initialize(argv, **config_arguments) + + self.rngs = nnx.Rngs(0) + devices = np.array(jax.devices()[:1]) + self.mesh = jax.sharding.Mesh(devices, ("tensor",)) + + def _apply_param_mapping(self, mt_layer, pt_layer, l): + import importlib.util + import os + mapping_path = os.path.join(os.path.dirname(__file__), "../../deepseek4-references/conversion_mapping.py") + spec = importlib.util.spec_from_file_location("conversion_mapping", mapping_path) + conversion_mapping = importlib.util.module_from_spec(spec) + spec.loader.exec_module(conversion_mapping) + PARAM_MAPPING = conversion_mapping.PARAM_MAPPING + + def get_attr(obj, path): + if path is None: return None + if "mlp.experts.." in path: + parts = path.split("..") + expert_obj = obj.mlp.experts + idx_and_weight = parts[1].split(".") + idx = int(idx_and_weight[0]) + w_name = idx_and_weight[1] + if w_name == "w1": + return expert_obj.gate_up_proj[idx, :expert_obj.intermediate_dim, :] + elif w_name == "w3": + return expert_obj.gate_up_proj[idx, expert_obj.intermediate_dim:, :] + elif w_name == "w2": + return expert_obj.down_proj[idx] + for part in path.split('.'): + if hasattr(obj, part): obj = getattr(obj, part) + elif isinstance(obj, list) or isinstance(obj, dict): obj = obj[int(part)] if isinstance(obj, list) else obj[part] + else: return None + return obj + + mt_prefix = f"params.params.decoder.layers_{l}." + pt_prefix = f"model.layers.{l}." + for mt_key, (pt_key, rule) in PARAM_MAPPING.items(): + if mt_key.startswith(mt_prefix) or f"params.Tid2EidVar.decoder.layers_{l}" in mt_key: + if "Tid2EidVar" in mt_key: + mt_path = mt_key.replace(f"params.Tid2EidVar.decoder.layers_{l}.", "") + ".value" + else: + mt_path = mt_key.replace(mt_prefix, "") + ".value" + + if pt_key is None: pt_obj = None + elif type(pt_key) == list: pt_obj = pt_key + else: pt_obj = get_attr(pt_layer, pt_key.replace(pt_prefix, "")) + + # Apply rule + val = None + if rule == "direct": val = jnp.array(pt_obj.detach().numpy()) + elif rule == "transpose": val = jnp.array(pt_obj.detach().numpy().T) + elif rule == "stack_transpose": + try: + tensors = [get_attr(pt_layer, path.replace(pt_prefix, "")) for path in pt_obj] + val = jnp.array(torch.stack(tensors).detach().numpy()).transpose(0, 2, 1) + except Exception as e: + print(f"FAILED stack_transpose: pt_obj={pt_obj}, tensors={['None' if t is None else 'Tensor' for t in tensors]}") + raise e + elif rule == "expert_gate_proj": + val = pt_obj.detach().numpy() + intermediate_dim = val.shape[1] // 2 + val = jnp.array(val[:, :intermediate_dim, :].transpose(0, 2, 1)) + elif rule == "expert_up_proj": + val = pt_obj.detach().numpy() + intermediate_dim = val.shape[1] // 2 + val = jnp.array(val[:, intermediate_dim:, :].transpose(0, 2, 1)) + elif rule == "expert_down_proj": + val = pt_obj.detach().numpy() + val = jnp.array(val.transpose(0, 2, 1)) + elif rule == "ones": pass + elif rule.startswith("mhc_fn_"): + hc = pt_layer.attn_hc.hc_mult + fn = pt_obj.detach().numpy() + if rule == "mhc_fn_pre": val = fn[:hc, :] + elif rule == "mhc_fn_post": val = fn[hc:2*hc, :] + elif rule == "mhc_fn_res": val = fn[2*hc:, :] + val = jnp.array(val.T) + elif rule.startswith("mhc_base_"): + hc = pt_layer.attn_hc.hc_mult + base = pt_obj.detach().numpy() + if rule == "mhc_base_pre": val = base[:hc] + elif rule == "mhc_base_post": val = base[hc:2*hc] + elif rule == "mhc_base_res": val = base[2*hc:].reshape(hc, hc) + val = jnp.array(val) + elif rule.startswith("mhc_scale_"): + scale = pt_obj.detach().numpy() + if rule == "mhc_scale_pre": val = scale[0] + elif rule == "mhc_scale_post": val = scale[1] + elif rule == "mhc_scale_res": val = scale[2] + val = jnp.array([val]) + elif rule == "reshape_transpose_oa": + val = pt_obj.detach().numpy() + val = val.reshape(self.pt_config.o_groups, -1, val.shape[1]).transpose(0, 2, 1) + val = jnp.array(val) + elif rule == "transpose_reshape_q": + val = pt_obj.detach().numpy().T.reshape(self.pt_config.q_lora_rank, self.pt_config.num_attention_heads, self.pt_config.head_dim) + val = jnp.array(val) + elif rule == "transpose_reshape_kv": + val = pt_obj.detach().numpy().T.reshape(-1, self.pt_config.num_key_value_heads, self.pt_config.head_dim) + val = jnp.array(val) + + if val is not None or rule == "ones": + parts = mt_path.split('.') + obj = mt_layer + valid = True + for part in parts[:-1]: + if hasattr(obj, part): obj = getattr(obj, part) + else: valid = False; break + if valid: + try: + if rule == "ones": setattr(obj, parts[-1], jnp.ones_like(getattr(obj, parts[-1]))) + else: setattr(obj, parts[-1], val) + except Exception as e: + print(f"FAILED on mt_key={mt_key}, mt_path={mt_path}, pt_key={pt_key}, obj={obj}") + raise e + + def _run_layer_parity_test(self, layer_idx, layer_type): + # self.pt_config.layer_types = ["sliding_attention"] * 7 + # self.pt_config.layer_types[layer_idx] = layer_type + compress_ratios = [0, 0, 4, 128, 4, 128, 4] + + torch.manual_seed(42) + pt_layer = DeepseekV4DecoderLayer_PT(self.pt_config, layer_idx=layer_idx) + + # Explicitly initialize PyTorch weights with random values to prevent torch.empty + # from yielding zero/garbage values that could mask parity differences. + for p in pt_layer.parameters(): + if p.dim() >= 1: + torch.nn.init.normal_(p.data, mean=0.0, std=0.02) + else: + torch.nn.init.constant_(p.data, 0.02) + + if layer_idx < self.mx_config.first_num_hash_layers: + pt_tid2eid = torch.randint(0, self.pt_config.num_local_experts, (self.vocab_size, self.pt_config.num_experts_per_tok)) + pt_layer.mlp.gate.tid2eid.copy_(pt_tid2eid) + + if layer_type == "compressed_sparse_attention" and self.pt_config.index_topk == 2: + for p in pt_layer.self_attn.compressor.indexer.parameters(): + p.data = torch.abs(p.data) + 0.1 + + mt_layer = DeepSeek4DecoderLayer( + config=self.mx_config, + model_mode="train", + mesh=self.mesh, + rngs=self.rngs, + layer_idx=layer_idx, + compress_ratio=compress_ratios[layer_idx], + is_hash_routing=(layer_idx < self.mx_config.first_num_hash_layers) + ) + + self._apply_param_mapping(mt_layer, pt_layer, layer_idx) + + np.random.seed(42) + x_np = np.random.uniform(0.1, 1.0, size=(self.batch_size, self.seq_len, self.pt_config.hc_mult, self.hidden_dim)).astype(np.float32) + pos_np = np.arange(self.seq_len)[None, :].repeat(self.batch_size, axis=0) + input_ids_np = np.random.randint(0, self.vocab_size, size=(self.batch_size, self.seq_len)) + + x_pt = torch.tensor(x_np) + pos_pt = torch.tensor(pos_np, dtype=torch.long) + input_ids_pt = torch.tensor(input_ids_np, dtype=torch.long) + + from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask + pt_mask = _prepare_4d_causal_attention_mask(None, (self.batch_size, self.seq_len), x_pt, 0, self.pt_config.sliding_window) + + rope_main = PTRope(self.pt_config) + rope_compress = PTRope(self.pt_config) + dummy_x_main = torch.zeros(self.batch_size, self.seq_len, 1) + cos_main, sin_main = rope_main(dummy_x_main, pos_pt, "main") + cos_comp, sin_comp = rope_compress(dummy_x_main, pos_pt, "compress") + pt_positions = {"main": (cos_main, sin_main), "compress": (cos_comp, sin_comp)} + + pt_out = pt_layer( + hidden_states=x_pt, + input_ids=input_ids_pt, + attention_mask=pt_mask, + position_ids=pos_pt, + position_embeddings=pt_positions + ) + + x_mt = jnp.array(x_np) + pos_mt = jnp.array(pos_np) + input_ids_mt = jnp.array(input_ids_np) + segs_mt = jnp.ones_like(pos_mt, dtype=jnp.int32) + + mt_out, _ = mt_layer( + inputs=x_mt, + decoder_segment_ids=segs_mt, + decoder_positions=pos_mt, + deterministic=True, + model_mode="train", + decoder_input_tokens=input_ids_mt, + ) + + pt_out_tensor = pt_out[0] if isinstance(pt_out, tuple) else pt_out + pt_out_np = pt_out_tensor.detach().numpy() + mt_out_np = np.array(mt_out) + max_diff = np.max(np.abs(mt_out_np - pt_out_np)) + mean_diff = np.mean(np.abs(mt_out_np - pt_out_np)) + print(f"LAYER PARITY layer_idx={layer_idx} layer_type={layer_type} - MAX ABS DIFF: {max_diff:.6e}, MEAN ABS DIFF: {mean_diff:.6e}") + np.testing.assert_allclose(mt_out_np, pt_out_np, rtol=5e-2, atol=5e-2) + + def test_layer_0_sliding_hash(self): + self._run_layer_parity_test(0, "sliding_attention") + + def test_layer_2_csa_hash(self): + self._run_layer_parity_test(2, "compressed_sparse_attention") + + def test_layer_3_hca_standard(self): + self._run_layer_parity_test(3, "heavily_compressed_attention") + + def test_layer_4_csa_standard(self): + self._run_layer_parity_test(4, "compressed_sparse_attention") + +class DeepSeekV4HyperHeadTest(unittest.TestCase): + """Tests to validate MaxText HyperHead implementation against PyTorch reference.""" + + def setUp(self): + self.batch_size = 2 + self.seq_len = 16 + self.hc_mult = 4 + self.hidden_dim = 4096 + + self.config_pt = DeepseekV4Config( + hidden_size=self.hidden_dim, + hc_mult=self.hc_mult, + rms_norm_eps=1e-6, + hc_eps=1e-6, + ) + + # Initialize PyTorch module + torch.manual_seed(42) + self.pt_head = DeepseekV4HyperHead_PT(self.config_pt) + # Initialize weights with standard values + for p in self.pt_head.parameters(): + torch.nn.init.normal_(p.data, mean=0.0, std=0.02) + + # Create dummy mesh/rngs for MaxText + devices = mesh_utils.create_device_mesh((1,), devices=jax.local_devices()[:1]) + self.mesh = Mesh(devices, ("x",)) + self.rngs = nnx.Rngs(0) + + # Build MaxText config dictionary + argv = ["", "src/maxtext/configs/base.yml", "model_name=deepseek4-tiny"] + config_arguments = { + "attention": "dot_product", + "dtype": "float32", + "weight_dtype": "float32", + "mhc_expansion_rate": self.hc_mult, + "emb_dim": self.hidden_dim, + "normalization_layer_epsilon": 1e-6, + "skip_jax_distributed_system": True, + } + self.mx_config = pyconfig.initialize(argv, **config_arguments) + + def test_hyper_head_parity(self): + mt_head = DeepSeek4HyperHead( + config=self.mx_config, + mesh=self.mesh, + rngs=self.rngs, + ) + + # Map parameters from PyTorch to MaxText + mt_head.hc_fn.value = jnp.array(self.pt_head.hc_fn.detach().numpy().T) + mt_head.hc_base.value = jnp.array(self.pt_head.hc_base.detach().numpy()) + mt_head.hc_scale.value = jnp.array(self.pt_head.hc_scale.detach().numpy()) + + # Inputs + np.random.seed(42) + x_np = np.random.uniform(0.1, 1.0, size=(self.batch_size, self.seq_len, self.hc_mult, self.hidden_dim)).astype(np.float32) + + x_pt = torch.tensor(x_np) + pt_out = self.pt_head(x_pt).detach().numpy() + + x_mt = jnp.array(x_np) + mt_out = np.array(mt_head(x_mt)) + + max_diff = np.max(np.abs(mt_out - pt_out)) + mean_diff = np.mean(np.abs(mt_out - pt_out)) + print(f"HYPER HEAD PARITY - MAX ABS DIFF: {max_diff:.6e}, MEAN ABS DIFF: {mean_diff:.6e}") + np.testing.assert_allclose(mt_out, pt_out, rtol=5e-5, atol=5e-5) + if __name__ == "__main__": unittest.main()