diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index c2ff040979..1c904e8c2e 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -926,7 +926,13 @@ def main( max_logging.log("Eager load with Transformers backend, from_pretrained with auto dtype") # For auto mode, loaded dtype is the same as `dtype` specified in config.json (or `torch_dtype` for older version) # e.g., https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite/blob/main/config.json#L54 - hf_state_dict_numpy = load_hf_dict_from_transformers(model_id, token=hf_token, revision=revision, dtype="auto") + hf_state_dict_numpy = load_hf_dict_from_transformers( + model_id, + token=hf_token, + revision=revision, + dtype="auto", + trust_remote_code=config.hf_trust_remote_code, + ) elif eager_load_method == "safetensors": max_logging.log("Eager load with Safetensors backend, safe_open with pt framework") # For safe_open, loaded dtype is the same as original safetensor diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 8bb262958c..4c99761192 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -1698,6 +1698,44 @@ def __init__(self, **kwargs): } qwen3_vl_2b_config = PTConfig(**qwen3_vl_2b_dict) +deepseek_ocr_2_dict = { + "architectures": ["DeepseekOCR2ForCausalLM"], + "model_type": "DeepseekOCR2", + "hidden_size": 1280, + "num_hidden_layers": 12, + "num_attention_heads": 10, + "num_key_value_heads": 10, + "use_mla": False, + "attention_type": "global", + "n_routed_experts": 64, + "n_shared_experts": 2, + "num_experts_per_tok": 6, + "moe_intermediate_size": 896, + "first_k_dense_replace": 1, + "vocab_size": 129280, + "rms_norm_eps": 1e-06, + "rope_theta": 10000.0, + "vision_config": { + "image_size": 1024, + "model_name": "deepencoderv2", + "sam_vit_b": { + "width": 768, + "layers": 12, + "heads": 12, + "global_attn_indexes": [2, 5, 8, 11], + }, + "qwen2_0_5b": { + "dim": 896, + "layers": 24, + "heads": 14, + "kv_heads": 2, + "intermediate_size": 4864, + }, + }, + "projector_config": {"input_dim": 896, "n_embed": 1280, "projector_type": "linear"}, +} +deepseek_ocr_2_config = PTConfig(**deepseek_ocr_2_dict) + # {maxtext model name: hf model config} HF_MODEL_CONFIGS = { @@ -1736,6 +1774,7 @@ def __init__(self, **kwargs): "qwen3-235b-a22b": qwen3_235b_a22b_thinking_2507_config, "qwen3-480b-a35b": qwen3_coder_480b_a35b_config, "deepseek2-16b": deepseek2_16b_config, + "deepseek_ocr_2": deepseek_ocr_2_config, "deepseek3-671b": deepseek3_671b_config, "deepseek3.2-671b": deepseek32_671b_config, "gpt-oss-20b": gpt_oss_20b_config, diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 80707a2abd..5bf5804378 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -1638,6 +1638,13 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fal "self_attention-indexer-wk-kernel": "self_attn.indexer.wk.weight", "self_attention-indexer-wq_b-kernel": "self_attn.indexer.wq_b.weight", } + if not config.get("use_mla", True): + attention_keys.update( + { + "self_attention-key-kernel": "self_attn.k_proj.weight", + "self_attention-value-kernel": "self_attn.v_proj.weight", + } + ) # Dense Layers dense_layer_keys = attention_keys | { "mlp-wi_0-kernel": "mlp.gate_proj.weight", @@ -1681,16 +1688,15 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fal else: for i in range(first_num_dense_layers): for maxtext_key, hf_key in dense_layer_keys.items(): - mapping[f"params-decoder-dense_layers_{i}-{maxtext_key}"] = f"model.layers.{i}.{hf_key}" + mapping[f"params-decoder-dense_layer_{i}-{maxtext_key}"] = f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers, num_main_layers): - moe_layer_idx = i - first_num_dense_layers - + # We use the global layer index 'i' because NNX uses 'layers_{i}' due to lexicographical ordering in NNX flattening. for maxtext_key, hf_key in moe_layer_keys.items(): - mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = f"model.layers.{i}.{hf_key}" + mapping[f"params-decoder-layers_{i}-{maxtext_key}"] = f"model.layers.{i}.{hf_key}" for maxtext_key, hf_key in moe_expert_keys.items(): - mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = [ + mapping[f"params-decoder-layers_{i}-{maxtext_key}"] = [ f"model.layers.{i}.mlp.experts.{e}.{hf_key}" for e in range(num_experts) ] return mapping @@ -1707,6 +1713,16 @@ def reshape_kernel(input_tensor, target_shape): else: return input_tensor.T.reshape(target_shape) + def scale_query_kernel(input_tensor, target_shape): + """Converts between HF's runtime attention scale and MaxText's folded q scale.""" + del target_shape + head_dim = config.get("head_dim", getattr(maxtext_config, "head_dim", None)) + if head_dim is None: + raise ValueError("DeepSeek q-projection conversion requires head_dim in config or maxtext_config.") + depth_scale = np.dtype("float32").type(np.sqrt(head_dim)) + factor = depth_scale if saving_to_hf else np.dtype("float32").type(1.0 / depth_scale) + return (input_tensor.astype(np.float32) * factor).astype(input_tensor.dtype) + num_main_layers = config["num_hidden_layers"] first_num_dense_layers = config["first_k_dense_replace"] @@ -1714,6 +1730,9 @@ def reshape_kernel(input_tensor, target_shape): "params-decoder-logits_dense-kernel": reshape_kernel, } + use_mla = config.get("use_mla", True) + query_hook_chain = [reshape_kernel, scale_query_kernel] + attention_need_reshape = { "self_attention-wkv_a-kernel", # transpose "self_attention-wkv_b-kernel", @@ -1729,6 +1748,15 @@ def reshape_kernel(input_tensor, target_shape): "self_attention-indexer-wq_b-kernel", } + if not use_mla: + attention_need_reshape.add("self_attention-key-kernel") + attention_need_reshape.add("self_attention-value-kernel") + + def hook_for_key(key): + if not use_mla and key == "self_attention-query-kernel": + return query_hook_chain + return reshape_kernel + dense_need_reshape = attention_need_reshape | { "mlp-wi_0-kernel", # transpose "mlp-wi_1-kernel", # transpose @@ -1748,18 +1776,17 @@ def reshape_kernel(input_tensor, target_shape): # scan if scan_layers: for key in dense_need_reshape: - mapping[f"params-decoder-dense_layers-{key}"] = reshape_kernel + mapping[f"params-decoder-dense_layers-{key}"] = hook_for_key(key) for key in moe_need_reshape: - mapping[f"params-decoder-moe_layers-{key}"] = reshape_kernel + mapping[f"params-decoder-moe_layers-{key}"] = hook_for_key(key) # unscan else: for i in range(first_num_dense_layers): for key in dense_need_reshape: - mapping[f"params-decoder-dense_layers_{i}-{key}"] = reshape_kernel + mapping[f"params-decoder-dense_layer_{i}-{key}"] = hook_for_key(key) for i in range(first_num_dense_layers, num_main_layers): - moe_layer_idx = i - first_num_dense_layers for key in moe_need_reshape: - mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{key}"] = reshape_kernel + mapping[f"params-decoder-layers_{i}-{key}"] = hook_for_key(key) return mapping @@ -3857,6 +3884,165 @@ def reshape_vision_attn_out(input_tensor, target_shape): return mapping +def DEEPSEEK_OCR_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): + """Generates a parameter mapping from MaxText to HuggingFace DeepSeek-OCR-2.""" + tcfg = config.get("text_config", config) + vcfg = config.get("vision_config", {}) + + sam_depth = 12 + connector_depth = vcfg.get("encoder_config", {}).get("num_hidden_layers", vcfg.get("qwen2_0_5b", {}).get("layers", 24)) + + mapping = { + # Projector + "params-vision_encoder-MlpProjector_0-linear-kernel": "model.projector.layers.weight", + "params-vision_encoder-MlpProjector_0-linear-bias": "model.projector.layers.bias", + "params-vision_encoder-MlpProjector_0-view_seperator": "model.view_seperator", + # Vision Tower - SAM Pos & Patch Embed + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-patch_embed-kernel": "model.sam_model.patch_embed.proj.weight", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-patch_embed-bias": "model.sam_model.patch_embed.proj.bias", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-pos_embed": "model.sam_model.pos_embed", + # Vision Tower - SAM Neck + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-neck_conv1-kernel": "model.sam_model.neck.0.weight", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-neck_ln1-scale": "model.sam_model.neck.1.weight", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-neck_ln1-bias": "model.sam_model.neck.1.bias", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-neck_conv2-kernel": "model.sam_model.neck.2.weight", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-neck_ln2-scale": "model.sam_model.neck.3.weight", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-neck_ln2-bias": "model.sam_model.neck.3.bias", + # Vision Tower - SAM Net2 & Net3 + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-net_2-kernel": "model.sam_model.net_2.weight", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-net_3-kernel": "model.sam_model.net_3.weight", + # Vision Tower - Qwen2 Connector Queries & Norm + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-qwen2_model-query_768-embedding": "model.qwen2_model.query_768.weight", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-qwen2_model-query_1024-embedding": "model.qwen2_model.query_1024.weight", + "params-vision_encoder-DeepseekOCR2VisionEncoder_0-qwen2_model-norm-scale": "model.qwen2_model.model.model.norm.weight", + } + + # SAM Blocks + sam_params = [ + ("norm1-scale", "norm1.weight"), + ("norm1-bias", "norm1.bias"), + ("norm2-scale", "norm2.weight"), + ("norm2-bias", "norm2.bias"), + ("attn-qkv-kernel", "attn.qkv.weight"), + ("attn-qkv-bias", "attn.qkv.bias"), + ("attn-proj-kernel", "attn.proj.weight"), + ("attn-proj-bias", "attn.proj.bias"), + ("attn-rel_pos_h", "attn.rel_pos_h"), + ("attn-rel_pos_w", "attn.rel_pos_w"), + ("lin1-kernel", "mlp.lin1.weight"), + ("lin1-bias", "mlp.lin1.bias"), + ("lin2-kernel", "mlp.lin2.weight"), + ("lin2-bias", "mlp.lin2.bias"), + ] + for i in range(sam_depth): + for mx, hf in sam_params: + key = f"params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-block_{i}-{mx}" + mapping[key] = f"model.sam_model.blocks.{i}.{hf}" + + # Qwen2 Connector Layers + connector_params = [ + ("pre_self_attention_layer_norm-scale", "input_layernorm.weight"), + ("post_self_attention_layer_norm-scale", "post_attention_layernorm.weight"), + ("self_attention-query-kernel", "self_attn.q_proj.weight"), + ("self_attention-query-bias", "self_attn.q_proj.bias"), + ("self_attention-key-kernel", "self_attn.k_proj.weight"), + ("self_attention-key-bias", "self_attn.k_proj.bias"), + ("self_attention-value-kernel", "self_attn.v_proj.weight"), + ("self_attention-value-bias", "self_attn.v_proj.bias"), + ("self_attention-out-kernel", "self_attn.o_proj.weight"), + ("mlp-wi_0-kernel", "mlp.gate_proj.weight"), + ("mlp-wi_1-kernel", "mlp.up_proj.weight"), + ("mlp-wo-kernel", "mlp.down_proj.weight"), + ] + for i in range(connector_depth): + for mx, hf in connector_params: + key = f"params-vision_encoder-DeepseekOCR2VisionEncoder_0-qwen2_model-layer_{i}-{mx}" + mapping[key] = f"model.qwen2_model.model.model.layers.{i}.{hf}" + + # Get text mapping + text_mapping = DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(tcfg, maxtext_config, scan_layers) + + # Adjust text mapping paths + for maxtext_key, hf_key in text_mapping.items(): + mapping[maxtext_key] = hf_key + + return mapping + + +def DEEPSEEK_OCR_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): + """Creates parameter transformation functions for DeepSeek-OCR-2.""" + hooks = {} + + tcfg = config.get("text_config", config) + vcfg = config.get("vision_config", {}) + + connector_layers = vcfg.get("encoder_config", {}).get("num_hidden_layers", vcfg.get("qwen2_0_5b", {}).get("layers", 24)) + sam_layers = 12 + + def reshape_kernel(x, target_shape): + if saving_to_hf: + flipped = np.flip(np.array(target_shape)) + return x.reshape(flipped).T + else: + return x.T.reshape(target_shape) + + def reshape_bias(x, target_shape=None): + return x.reshape(target_shape) + + def vision_patch(x, target_shape): + if saving_to_hf: + return x.transpose(3, 2, 0, 1) + else: + return x.transpose(2, 3, 1, 0) + + # Projector + hooks["params-vision_encoder-MlpProjector_0-linear-kernel"] = reshape_kernel + hooks["params-vision_encoder-MlpProjector_0-linear-bias"] = reshape_bias + hooks["params-vision_encoder-MlpProjector_0-view_seperator"] = reshape_bias + + # SAM Patch Embed + hooks["params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-patch_embed-kernel"] = vision_patch + hooks["params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-patch_embed-bias"] = reshape_bias + + # SAM Blocks + for i in range(sam_layers): + base = f"params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-block_{i}-" + hooks[base + "attn-qkv-kernel"] = reshape_kernel + hooks[base + "attn-qkv-bias"] = reshape_bias + hooks[base + "attn-proj-kernel"] = reshape_kernel + hooks[base + "attn-proj-bias"] = reshape_bias + hooks[base + "lin1-kernel"] = reshape_kernel + hooks[base + "lin1-bias"] = reshape_bias + hooks[base + "lin2-kernel"] = reshape_kernel + hooks[base + "lin2-bias"] = reshape_bias + + # SAM Neck + hooks["params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-neck_conv1-kernel"] = vision_patch + hooks["params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-neck_conv2-kernel"] = vision_patch + hooks["params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-net_2-kernel"] = vision_patch + hooks["params-vision_encoder-DeepseekOCR2VisionEncoder_0-sam_model-net_3-kernel"] = vision_patch + + # Qwen2 Connector Layers + for i in range(connector_layers): + base = f"params-vision_encoder-DeepseekOCR2VisionEncoder_0-qwen2_model-layer_{i}-" + hooks[base + "self_attention-query-kernel"] = reshape_kernel + hooks[base + "self_attention-query-bias"] = reshape_bias + hooks[base + "self_attention-key-kernel"] = reshape_kernel + hooks[base + "self_attention-key-bias"] = reshape_bias + hooks[base + "self_attention-value-kernel"] = reshape_kernel + hooks[base + "self_attention-value-bias"] = reshape_bias + hooks[base + "self_attention-out-kernel"] = reshape_kernel + hooks[base + "mlp-wi_0-kernel"] = reshape_kernel + hooks[base + "mlp-wi_1-kernel"] = reshape_kernel + hooks[base + "mlp-wo-kernel"] = reshape_kernel + + # Get text hooks + text_hooks = DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN(tcfg, maxtext_config, scan_layers, saving_to_hf) + hooks.update(text_hooks) + + return hooks + + # {maxtext model name: {maxtext weight name: hf weight name}} PARAM_MAPPING = { "gemma2-2b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -3896,6 +4082,7 @@ def reshape_vision_attn_out(input_tensor, target_shape): "deepseek2-16b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, + "deepseek_ocr_2": DEEPSEEK_OCR_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-20b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-120b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -3948,6 +4135,7 @@ def reshape_vision_attn_out(input_tensor, target_shape): "deepseek2-16b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "deepseek_ocr_2": DEEPSEEK_OCR_MAXTEXT_TO_HF_PARAM_HOOK_FN, "gpt-oss-20b": GPT_OSS_TO_HF_PARAM_HOOK_FN, "gpt-oss-120b": GPT_OSS_TO_HF_PARAM_HOOK_FN, "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN, diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 75ee3dffcc..6ed3c99fa0 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -1042,7 +1042,13 @@ def detect_and_extract_checkpoint(checkpoint_dict: dict) -> dict[str, np.ndarray return extract_linen_weights(actual_weights_dict) -def load_hf_dict_from_transformers(model_id: str, token: str, revision: str | None = None, dtype: str = "auto"): +def load_hf_dict_from_transformers( + model_id: str, + token: str, + revision: str | None = None, + dtype: str = "auto", + trust_remote_code: bool = False, +): """Loads the HuggingFace model based on model_id (Eager mode only), used in to_maxtext""" # 1. Handle special cases requiring specific model classes @@ -1062,7 +1068,13 @@ def load_hf_dict_from_transformers(model_id: str, token: str, revision: str | No last_exception = None for model_class in model_classes: try: - hf_model = model_class.from_pretrained(model_id, token=token, revision=revision, dtype=dtype) + hf_model = model_class.from_pretrained( + model_id, + token=token, + revision=revision, + dtype=dtype, + trust_remote_code=trust_remote_code, + ) break except (ValueError, OSError, RuntimeError, ImportError) as e: max_logging.log(f"Failed to load using {model_class.__name__}: {e!r}") diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82e448c530..61662a94ff 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -372,6 +372,7 @@ param_scan_axis: 1 # The attention_type parameter determines the variants of attention, e.g. global or local_sliding attention: 'autoselected' # Supported attention: autoselected, dot_product, flash, cudnn_flash_te attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla +use_mla: true # Whether to use Multi-Head Latent Attention (MLA) for DeepSeek share_kv_projections: false # Note: Not compatible with attention_type='mla' attention_bias: false # If true, adds a learnable bias to the query, key, and value projections attention_sink: false @@ -673,6 +674,7 @@ tokenizer_path: "" # grain and tfds pipeline supports tokenizer_type: sentencepiece, huggingface, tiktoken # hf pipeline only supports huggingface type, and will ignore tokenizer_type flag tokenizer_type: "sentencepiece" # Currently supporting: "tiktoken", "sentencepiece", "huggingface" +hf_trust_remote_code: false # If true, allow Hugging Face tokenizer/model loading to execute repo custom code. use_chat_template: false chat_template_path: "" # path to chat template json file chat_template: "" # Chat template to use with HF tokenizers. It should be a valid Jinja2-formatted template. @@ -1199,6 +1201,11 @@ rope_theta_for_vit: 10000 vision_output_dim_for_vit: 4096 pixel_shuffle_ratio_for_vit: 0.5 projector_dropout_for_vit: 0.0 +vision_connector_num_layers: 0 +vision_connector_emb_dim: 0 +vision_connector_num_query_heads: 0 +vision_connector_num_kv_heads: 0 +vision_connector_mlp_dim: 0 # Qwen3-OmniMoe vision encoder spatial_merge_size_for_vit: 2 diff --git a/src/maxtext/configs/models/deepseek_ocr_2.yml b/src/maxtext/configs/models/deepseek_ocr_2.yml new file mode 100644 index 0000000000..42b788c3ee --- /dev/null +++ b/src/maxtext/configs/models/deepseek_ocr_2.yml @@ -0,0 +1,49 @@ +# model config for DeepSeek-OCR-2 (1.2B) +# This is a MoE VLM. + +# Language Model (DeepSeek-V2-Lite-like but smaller) +base_emb_dim: 1280 +base_num_query_heads: 10 +base_num_kv_heads: 10 +base_mlp_dim: 6848 # dense MLP intermediate dim (layer 0) +base_moe_mlp_dim: 896 # MoE MLP intermediate dim (layers 1-11) +base_num_decoder_layers: 12 +first_num_dense_layers: 1 +mlp_activations: ["silu","linear"] +vocab_size: 129280 +enable_dropout: false +logits_via_embedding: false +normalization_layer_epsilon: 1.0e-6 + +# MoE +num_experts: 64 +num_experts_per_tok: 6 +shared_experts: 2 +routed_scaling_factor: 1.0 +routed_score_func: "softmax" +routed_bias: false +decoder_block: "deepseek" +use_mla: false + +# Attention / RoPE +attention_type: "global" +rope_type: "default" # standard RoPE +rope_max_timescale: 10000 +max_position_embeddings: 4096 + +# Multimodal +# use_multimodal: true +image_placeholder: +# Vision Encoder (SAM ViT-B) +image_size_for_vit: 1024 +patch_size_for_vit: 16 +hidden_size_for_vit: 768 +num_attention_heads_for_vit: 12 +num_hidden_layers_for_vit: 12 + +# Vision Connector (Qwen2Decoder2Encoder) +vision_connector_num_layers: 24 +vision_connector_emb_dim: 896 +vision_connector_num_query_heads: 14 +vision_connector_num_kv_heads: 2 +vision_connector_mlp_dim: 4864 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a8c1745..2b705686a6 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -228,6 +228,7 @@ class ProfilerType(str, Enum): "deepseek3-test", "deepseek3-tiny", "deepseek3.2-671b", + "deepseek_ocr_2", "deepseek4-284b", "deepseek-custom", "kimi-k2-1t", @@ -559,6 +560,10 @@ class Attention(BaseModel): attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed"] = Field( "global", description="The variant of attention to use." ) + use_mla: bool = Field( + True, + description="Whether to use Multi-Head Latent Attention (MLA) for DeepSeek.", + ) share_kv_projections: bool = Field( False, description="If True, for global attention, Key and Value projections share the same weights.", @@ -1173,6 +1178,10 @@ class Tokenizer(BaseModel): description="Path to the tokenizer model file.", ) tokenizer_type: TokenizerType = Field(TokenizerType.SENTENCEPIECE, description="The type of tokenizer.") + hf_trust_remote_code: bool = Field( + False, + description="Whether Hugging Face tokenizer/model loading may execute custom code from the referenced repo.", + ) use_chat_template: bool = Field(False, description="Whether to use the chat template for tokenization.") chat_template_path: str = Field("", description="Path to chat template json file.") chat_template: str = Field( @@ -2033,6 +2042,11 @@ class VisionProjector(BaseModel): projector_output_dim_for_vit: int = Field(4096, description="Output dimension for the vision projector.") pixel_shuffle_ratio_for_vit: float = Field(0.5, description="Pixel shuffle ratio for the Vision Transformer.") projector_dropout_for_vit: float = Field(0.0, description="Dropout rate for the vision projector.") + vision_connector_num_layers: int = Field(0, description="Number of layers in the vision connector.") + vision_connector_emb_dim: int = Field(0, description="Embedding dimension for the vision connector.") + vision_connector_num_query_heads: int = Field(0, description="Number of query heads in the vision connector.") + vision_connector_num_kv_heads: int = Field(0, description="Number of KV heads in the vision connector.") + vision_connector_mlp_dim: int = Field(0, description="MLP dimension for the vision connector.") class AudioEncoder(BaseModel): @@ -3188,6 +3202,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "qwen3-vl-4b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "deepseek_ocr_2", ) if self.model_name not in valid_mm_models and self.model_name != "default": raise ValueError(f"Multimodal is only supported for {valid_mm_models}, not {self.model_name}") diff --git a/src/maxtext/inference/decode.py b/src/maxtext/inference/decode.py index 68cd450c99..9ab5de056b 100644 --- a/src/maxtext/inference/decode.py +++ b/src/maxtext/inference/decode.py @@ -184,7 +184,9 @@ def main(argv: Sequence[str]) -> None: positions=position_ids, mrope_deltas=mrope_position_deltas, images=processor_outputs.pixel_values if config.use_multimodal else None, - image_masks=processor_outputs.pixel_mask if config.use_multimodal and "llama4" in config.model_name else None, + image_masks=processor_outputs.pixel_mask + if config.use_multimodal and ("llama4" in config.model_name or config.model_name == "deepseek_ocr_2") + else None, videos=getattr(processor_outputs, "video_values", None) if config.use_multimodal else None, audio_values=processor_outputs.audio_values if config.use_audio else None, audio_masks=processor_outputs.audio_mask if config.use_audio else None, diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 5377e86ec3..013dabc432 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -20,7 +20,9 @@ import os.path import uuid import warnings +import logging +import numpy as np from jax.experimental.layout import Format from jax.sharding import PartitionSpec as P import jax @@ -66,6 +68,76 @@ PRNGKeyType = Any +def _pad_tokens(tokens, bos_id, pad_id, is_bos=True, prefill_lengths=None, max_prefill_length=None): + """Pads HF tokenizer output using the same contract as JetStream token_utils.pad_tokens.""" + if prefill_lengths is None: + prefill_lengths = [64, 128, 256, 512, 1024, 2048, 4096] + if max_prefill_length is not None: + prefill_lengths = [length for length in prefill_lengths if length < max_prefill_length] + [max_prefill_length] + if is_bos: + if bos_id is None: + raise ValueError("Cannot add BOS because the Hugging Face tokenizer has no bos_token_id.") + tokens = np.concatenate([np.array([bos_id]), tokens], axis=-1) + true_length = tokens.shape[-1] + padded_length = min((length for length in prefill_lengths if length >= true_length), default=prefill_lengths[-1]) + padding = padded_length - true_length + if padding < 0: + logging.warning("Provided sequence longer than available.") + padded_tokens = tokens[-padded_length:] + else: + if pad_id is None: + raise ValueError("Cannot pad because the Hugging Face tokenizer has no pad_token_id.") + padded_tokens = np.pad(tokens, (0, padding), constant_values=(pad_id,)) + return jnp.array(padded_tokens), true_length + + +class MaxTextHuggingFaceTokenizer: + """Local HF tokenizer adapter that avoids JetStream's interactive remote-code prompt.""" + + def __init__(self, metadata: Any, trust_remote_code: bool): + from transformers import AutoTokenizer # pylint: disable=import-outside-toplevel + + self.tokenizer = AutoTokenizer.from_pretrained( + metadata.path, + token=metadata.access_token, + trust_remote_code=trust_remote_code, + ) + self.metadata = metadata + + def encode(self, s: str, **kwargs): + is_bos = kwargs.pop("is_bos", True) + prefill_lengths = kwargs.pop("prefill_lengths", None) + max_prefill_length = kwargs.pop("max_prefill_length", None) + if getattr(self.metadata, "use_chat_template", False): + tokens = self.tokenizer.apply_chat_template( + [{"role": "user", "content": s}], + add_generation_prompt=True, + return_tensors="np", + ).squeeze() + if is_bos: + logging.warning("Overriding is_bos to False because use_chat_template is set to True.") + is_bos = False + else: + tokens = self.tokenizer.encode(s, add_special_tokens=False, return_tensors="np").squeeze() + tokens = np.atleast_1d(tokens) + return _pad_tokens(tokens, self.bos_id, self.pad_id, is_bos, prefill_lengths, max_prefill_length) + + def decode(self, token_ids: list[int]) -> str: + return self.tokenizer.decode(token_ids, skip_special_tokens=True) + + @property + def pad_id(self): + return self.tokenizer.pad_token_id + + @property + def eos_id(self): + return self.tokenizer.eos_token_id + + @property + def bos_id(self): + return self.tokenizer.bos_token_id + + # TODO(yuyanpeng): Should import ExistingPrefix from jetstream.engine.engine_api @struct.dataclass class ExistingPrefix: @@ -761,6 +833,10 @@ def _prefill_jit( # add batch dimension images = images[jnp.newaxis, ...] image_masks = image_masks[jnp.newaxis, ...] if image_masks is not None else None + elif images.ndim == 5: + # add batch dimension + images = images[jnp.newaxis, ...] + image_masks = image_masks[jnp.newaxis, ...] if image_masks is not None else None # sequence_indicator will be concatenated to existing_prefix decoder_segment_ids start_to_n = jnp.arange(start_position, start_position + input_tokens.shape[1]) @@ -1840,7 +1916,7 @@ def build_tokenizer(self, metadata: Any): # return type depends on JetStream elif metadata.tokenizer_type == TokenizerType.sentencepiece: return token_utils.SentencePieceTokenizer(metadata) elif metadata.tokenizer_type == TokenizerType.huggingface: - tokenizer_model = token_utils.HuggingFaceTokenizer(metadata) + tokenizer_model = MaxTextHuggingFaceTokenizer(metadata, trust_remote_code=self.config.hf_trust_remote_code) if tokenizer_model.tokenizer.pad_token_id is None: if tokenizer_model.tokenizer.unk_token_id is not None: tokenizer_model.tokenizer.pad_token_id = tokenizer_model.tokenizer.unk_token_id diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index d5bb3314b6..0d97b4e7ef 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -450,7 +450,7 @@ def ragged_gather_reduce( them via summation. The typical use case of this kernel is unpermute + local-reduction in the - MOE after GMM. Compared to maxtext.src.maxtext.kernels.gather_reduce_sc, + MOE after GMM. Compared to maxtext.kernels.gather_reduce_sc, this kernel provides better performance if large sparsity exists in `valid_rows_mask`. For example, expert_parallelism =8, 16 etc. diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 0ea2913265..bd90fe3df9 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -732,6 +732,7 @@ def _apply_embedding( "qwen3-vl-4b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "deepseek_ocr_2", ]: y = mm_utils.merge_mm_embeddings( text_embeddings=y, diff --git a/src/maxtext/layers/encoders.py b/src/maxtext/layers/encoders.py index 19a7490016..e07180683e 100644 --- a/src/maxtext/layers/encoders.py +++ b/src/maxtext/layers/encoders.py @@ -90,6 +90,16 @@ def _setup_vision_encoder_layers(self): ) setattr(self, projector_name, qwen3_vl_vision.Qwen3VLVisionProjector(config=self.config, rngs=self.rngs)) return encoder_name, projector_name + elif self.config.model_name in ["deepseek_ocr_2"]: + from maxtext.models import deepseek_ocr # pylint: disable=import-outside-toplevel + + encoder_name = "DeepseekOCR2VisionEncoder_0" + projector_name = "MlpProjector_0" + setattr( + self, encoder_name, deepseek_ocr.DeepseekOCR2VisionEncoder(config=self.config, mesh=self.mesh, rngs=self.rngs) + ) + setattr(self, projector_name, deepseek_ocr.MlpProjector(config=self.config, mesh=self.mesh, rngs=self.rngs)) + return encoder_name, projector_name else: raise ValueError(f"No VisionEncoder implemented for {self.config.model_name} yet") diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 4b7f6761a1..14d53ab555 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1348,6 +1348,7 @@ def _apply_embedding( "qwen3-vl-4b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "deepseek_ocr_2", }: y = mm_utils.merge_mm_embeddings( text_embeddings=y, diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index d3a72b31bf..1654974f14 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -27,6 +27,7 @@ from maxtext.common.common_types import Config from maxtext.common.common_types import HyperConnectionType, MODEL_MODE_PREFILL, DecoderBlockType from maxtext.layers import attention_mla +from maxtext.layers.attentions import Attention from maxtext.layers import initializers from maxtext.layers import linears from maxtext.layers import mhc @@ -140,37 +141,58 @@ def __init__( # DeepSeek V4 natively overrides this block with CompressedAttention. if self.config.decoder_block != DecoderBlockType.DEEPSEEK4: - self.self_attention = attention_mla.MLA( - config=self.config, - num_query_heads=self.config.num_query_heads, - num_kv_heads=self.config.num_kv_heads, - head_dim=self.config.head_dim, - max_target_length=self.config.max_target_length, - max_prefill_predict_length=self.config.max_prefill_predict_length, - attention_kernel=self.config.attention, - attention_type=self.config.attention_type, - inputs_q_shape=self.dummy_inputs_shape, - inputs_kv_shape=self.dummy_inputs_shape, - mesh=mesh, - dtype=self.config.dtype, - weight_dtype=self.config.weight_dtype, - dropout_rate=self.config.dropout_rate, - name="self_attention", - quant=quant, - kv_quant=quantizations.configure_kv_quant(self.config), - q_lora_rank=self.config.q_lora_rank, - kv_lora_rank=self.config.kv_lora_rank, - qk_nope_head_dim=self.config.qk_nope_head_dim, - qk_rope_head_dim=self.config.qk_rope_head_dim, - v_head_dim=self.config.v_head_dim, - max_position_embeddings=self.config.max_position_embeddings, - original_max_position_embeddings=self.config.original_max_position_embeddings, - mscale=self.config.mscale, - rope_factor=self.config.rope_factor, - model_mode=model_mode, - rngs=rngs, - attn_logits_soft_cap=self.config.attn_logits_soft_cap, - ) + if getattr(self.config, "use_mla", True): + self.self_attention = attention_mla.MLA( + config=self.config, + num_query_heads=self.config.num_query_heads, + num_kv_heads=self.config.num_kv_heads, + head_dim=self.config.head_dim, + max_target_length=self.config.max_target_length, + max_prefill_predict_length=self.config.max_prefill_predict_length, + attention_kernel=self.config.attention, + attention_type=self.config.attention_type, + inputs_q_shape=self.dummy_inputs_shape, + inputs_kv_shape=self.dummy_inputs_shape, + mesh=mesh, + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + dropout_rate=self.config.dropout_rate, + name="self_attention", + quant=quant, + kv_quant=quantizations.configure_kv_quant(self.config), + q_lora_rank=self.config.q_lora_rank, + kv_lora_rank=self.config.kv_lora_rank, + qk_nope_head_dim=self.config.qk_nope_head_dim, + qk_rope_head_dim=self.config.qk_rope_head_dim, + v_head_dim=self.config.v_head_dim, + max_position_embeddings=self.config.max_position_embeddings, + original_max_position_embeddings=self.config.original_max_position_embeddings, + mscale=self.config.mscale, + rope_factor=self.config.rope_factor, + model_mode=model_mode, + rngs=rngs, + attn_logits_soft_cap=self.config.attn_logits_soft_cap, + ) + else: + self.self_attention = Attention( + config=self.config, + num_query_heads=self.config.num_query_heads, + num_kv_heads=self.config.num_kv_heads, + head_dim=self.config.head_dim, + max_target_length=self.config.max_target_length, + max_prefill_predict_length=self.config.max_prefill_predict_length, + attention_kernel=self.config.attention, + inputs_q_shape=self.dummy_inputs_shape, + inputs_kv_shape=self.dummy_inputs_shape, + mesh=mesh, + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + dropout_rate=self.config.dropout_rate, + quant=quant, + kv_quant=quantizations.configure_kv_quant(self.config), + model_mode=model_mode, + rngs=rngs, + ) self.dropout = Dropout(rate=self.config.dropout_rate, broadcast_dims=(-2,), rngs=self.rngs) if self.is_mhc_enabled: @@ -210,6 +232,7 @@ def attention_op( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk=None, slot: None | int = None, ): @@ -220,7 +243,7 @@ def attention_op( decoder_positions, decoder_segment_ids=decoder_segment_ids, deterministic=deterministic, - model_mode=self.model_mode, + model_mode=model_mode, out_sharding=self.out_sharding, previous_chunk=previous_chunk, slot=slot, @@ -269,6 +292,7 @@ def self_attention_with_norm_op( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk=None, slot: None | int = None, ): @@ -282,7 +306,7 @@ def self_attention_with_norm_op( decoder_segment_ids=decoder_segment_ids, inputs_positions=decoder_positions, deterministic=deterministic, - model_mode=self.model_mode, + model_mode=model_mode, out_sharding=self.out_sharding, previous_chunk=previous_chunk, slot=slot, @@ -294,6 +318,7 @@ def self_attention_with_norm_op( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk, slot, ) @@ -367,6 +392,7 @@ def __call__( decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk, slot, ) @@ -580,6 +606,7 @@ def extract_fn(x): decoder_segment_ids, decoder_positions, deterministic, + model_mode, previous_chunk, slot, ) diff --git a/src/maxtext/models/deepseek_ocr.py b/src/maxtext/models/deepseek_ocr.py new file mode 100644 index 0000000000..f4b2447a40 --- /dev/null +++ b/src/maxtext/models/deepseek_ocr.py @@ -0,0 +1,550 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DeepSeek-OCR-2 vision encoder and connector models.""" + +import jax +import jax.numpy as jnp +from flax import nnx +from typing import Optional, Tuple +from maxtext.common.common_types import Config +from maxtext.layers.normalizations import RMSNorm +from maxtext.layers.linears import MlpBlock +from maxtext.layers.attentions import Attention +from maxtext.layers import quantizations +from maxtext.layers.quantizations import AqtQuantization as Quant +from jax.sharding import Mesh +from maxtext.configs.pyconfig import HyperParameters + +# ============================================================================== +# Helper functions for SAM ViT-B Relative Position Bias +# ============================================================================== + + +def get_rel_pos(q_size: int, k_size: int, rel_pos: jax.Array) -> jax.Array: + """Get relative positional embeddings.""" + max_rel_dist = int(2 * max(q_size, k_size) - 1) + if rel_pos.shape[0] != max_rel_dist: + rel_pos_resized = jax.image.resize(rel_pos, (max_rel_dist, rel_pos.shape[1]), method="linear") + else: + rel_pos_resized = rel_pos + + q_coords = jnp.arange(q_size)[:, None] * max(k_size / q_size, 1.0) + k_coords = jnp.arange(k_size)[None, :] * max(q_size / k_size, 1.0) + relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) + relative_coords = relative_coords.astype(jnp.int32) + + return rel_pos_resized[relative_coords] + + +def add_decomposed_rel_pos( + q: jax.Array, + rel_pos_h: jax.Array, + rel_pos_w: jax.Array, + q_size: Tuple[int, int], + k_size: Tuple[int, int], +) -> Tuple[jax.Array, jax.Array]: + """Calculate decomposed Relative Positional Embeddings.""" + q_h, q_w = q_size + k_h, k_w = k_size + Rh = get_rel_pos(q_h, k_h, rel_pos_h) + Rw = get_rel_pos(q_w, k_w, rel_pos_w) + + B, HW, dim = q.shape + r_q = q.reshape(B, q_h, q_w, dim) + rel_h = jnp.einsum("bhwc,hkc->bhwk", r_q, Rh) + rel_w = jnp.einsum("bhwc,wkc->bhwk", r_q, Rw) + + rel_h = jnp.expand_dims(rel_h, -1) + rel_w = jnp.expand_dims(rel_w, -2) + rel_h = rel_h.reshape(B, q_h * q_w, k_h, 1) + rel_w = rel_w.reshape(B, q_h * q_w, 1, k_w) + + return rel_h, rel_w + + +# ============================================================================== +# SAM ViT-B Components +# ============================================================================== + + +class SAMAttention(nnx.Module): + """Multi-head Attention block with relative position embeddings for SAM.""" + + def __init__( + self, + dim: int, + num_heads: int, + qkv_bias: bool, + use_rel_pos: bool, + input_size: Optional[Tuple[int, int]], + rngs: nnx.Rngs, + ): + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + self.use_rel_pos = use_rel_pos + + self.qkv = nnx.Linear(dim, dim * 3, use_bias=qkv_bias, rngs=rngs) + self.proj = nnx.Linear(dim, dim, use_bias=True, rngs=rngs) + + if self.use_rel_pos: + assert input_size is not None + self.rel_pos_h = nnx.Param(jnp.zeros((2 * input_size[0] - 1, self.head_dim))) + self.rel_pos_w = nnx.Param(jnp.zeros((2 * input_size[1] - 1, self.head_dim))) + + def __call__(self, x: jax.Array) -> jax.Array: + B, H, W, _ = x.shape + qkv = self.qkv(x.reshape(B, H * W, -1)) + qkv = qkv.reshape(B, H * W, 3, self.num_heads, self.head_dim) + qkv = qkv.transpose(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] + + if self.use_rel_pos: + q_flat = q.reshape(B * self.num_heads, H * W, self.head_dim) + rel_h, rel_w = add_decomposed_rel_pos(q_flat, self.rel_pos_h.value, self.rel_pos_w.value, (H, W), (H, W)) + + rel_h = rel_h.reshape(B, self.num_heads, H * W, H, 1) + rel_w = rel_w.reshape(B, self.num_heads, H * W, 1, W) + + attn_bias = rel_h + rel_w + attn_bias = attn_bias.reshape(B, self.num_heads, H * W, H * W) + + logits = jnp.einsum("bhid,bhjd->bhij", q, k) * self.scale + logits = logits + attn_bias + attn = jax.nn.softmax(logits, axis=-1) + out = jnp.einsum("bhij,bhjd->bhid", attn, v) + else: + logits = jnp.einsum("bhid,bhjd->bhij", q, k) * self.scale + attn = jax.nn.softmax(logits, axis=-1) + out = jnp.einsum("bhij,bhjd->bhid", attn, v) + + out = out.transpose(0, 2, 1, 3).reshape(B, H, W, -1) + out = self.proj(out) + return out + + +def window_partition(x: jax.Array, window_size: int) -> Tuple[jax.Array, Tuple[int, int]]: + """Partition into non-overlapping windows.""" + B, H, W, C = x.shape + pad_h = (window_size - H % window_size) % window_size + pad_w = (window_size - W % window_size) % window_size + if pad_h > 0 or pad_w > 0: + x = jnp.pad(x, ((0, 0), (0, pad_h), (0, pad_w), (0, 0))) + Hp, Wp = H + pad_h, W + pad_w + + x = x.reshape(B, Hp // window_size, window_size, Wp // window_size, window_size, C) + x = x.transpose(0, 1, 3, 2, 4, 5) + windows = x.reshape(-1, window_size, window_size, C) + return windows, (Hp, Wp) + + +def window_unpartition(windows: jax.Array, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]) -> jax.Array: + """Window unpartition.""" + Hp, Wp = pad_hw + H, W = hw + B = windows.shape[0] // (Hp * Wp // window_size // window_size) + x = windows.reshape(B, Hp // window_size, Wp // window_size, window_size, window_size, -1) + x = x.transpose(0, 1, 3, 2, 4, 5) + x = x.reshape(B, Hp, Wp, -1) + if Hp > H or Wp > W: + x = x[:, :H, :W, :] + return x + + +class SAMBlock(nnx.Module): + """Transformer block for SAM.""" + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float, + qkv_bias: bool, + use_rel_pos: bool, + window_size: int, + input_size: Optional[Tuple[int, int]], + rngs: nnx.Rngs, + ): + self.norm1 = nnx.LayerNorm(num_features=dim, epsilon=1e-5, rngs=rngs) + self.attn = SAMAttention( + dim=dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + use_rel_pos=use_rel_pos, + input_size=input_size if window_size == 0 else (window_size, window_size), + rngs=rngs, + ) + self.norm2 = nnx.LayerNorm(num_features=dim, epsilon=1e-5, rngs=rngs) + self.lin1 = nnx.Linear(dim, int(dim * mlp_ratio), rngs=rngs) + self.lin2 = nnx.Linear(int(dim * mlp_ratio), dim, rngs=rngs) + self.window_size = window_size + + def __call__(self, x: jax.Array) -> jax.Array: + shortcut = x + x = self.norm1(x) + + if self.window_size > 0: + H, W = x.shape[1], x.shape[2] + x, pad_hw = window_partition(x, self.window_size) + + x = self.attn(x) + + if self.window_size > 0: + x = window_unpartition(x, self.window_size, pad_hw, (H, W)) + + x = shortcut + x + mlp_out = self.lin2(jax.nn.gelu(self.lin1(self.norm2(x)))) + x = x + mlp_out + return x + + +def get_abs_pos_sam(abs_pos, tgt_size): + """Interpolate absolute position embeddings.""" + src_size = abs_pos.shape[1] + if src_size != tgt_size: + new_pos_embed = jax.image.resize(abs_pos, (1, tgt_size, tgt_size, abs_pos.shape[3]), method="bicubic") + return new_pos_embed + else: + return abs_pos + + +class SAMViTB(nnx.Module): + """SAM ViT-B image encoder.""" + + def __init__( + self, + config: Config, + mesh: Mesh, + rngs: nnx.Rngs, + ): + # SAM ViT-B hardcoded parameters + img_size = 1024 + patch_size = 16 + in_chans = 3 + embed_dim = 768 + depth = 12 + num_heads = 12 + mlp_ratio = 4.0 + out_chans = 256 + qkv_bias = True + use_abs_pos = True + use_rel_pos = True + window_size = 14 + global_attn_indexes = (2, 5, 8, 11) + + self.patch_embed = nnx.Conv( + in_features=in_chans, + out_features=embed_dim, + kernel_size=(patch_size, patch_size), + strides=(patch_size, patch_size), + padding="VALID", + rngs=rngs, + ) + + if use_abs_pos: + self.pos_embed = nnx.Param(jnp.zeros((1, img_size // patch_size, img_size // patch_size, embed_dim))) + else: + self.pos_embed = None + + self.blocks = nnx.List([]) + for i in range(depth): + block = SAMBlock( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + use_rel_pos=use_rel_pos, + window_size=window_size if i not in global_attn_indexes else 0, + input_size=(img_size // patch_size, img_size // patch_size), + rngs=rngs, + ) + self.blocks.append(block) + setattr(self, f"block_{i}", block) + + # Neck + self.neck_conv1 = nnx.Conv(embed_dim, out_chans, kernel_size=(1, 1), use_bias=False, rngs=rngs) + self.neck_ln1 = nnx.LayerNorm(num_features=out_chans, rngs=rngs) + self.neck_conv2 = nnx.Conv(out_chans, out_chans, kernel_size=(3, 3), padding="SAME", use_bias=False, rngs=rngs) + self.neck_ln2 = nnx.LayerNorm(num_features=out_chans, rngs=rngs) + + self.net_2 = nnx.Conv( + 256, 512, kernel_size=(3, 3), strides=(2, 2), padding=((1, 1), (1, 1)), use_bias=False, rngs=rngs + ) + self.net_3 = nnx.Conv( + 512, 896, kernel_size=(3, 3), strides=(2, 2), padding=((1, 1), (1, 1)), use_bias=False, rngs=rngs + ) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.patch_embed(x) + + if self.pos_embed is not None: + H_prime = x.shape[1] + pos_embed_resized = get_abs_pos_sam(self.pos_embed.value, H_prime) + x = x + pos_embed_resized + + for i in range(len(self.blocks)): + x = self.blocks[i](x) + + x = self.neck_conv1(x) + x = self.neck_ln1(x) + x = self.neck_conv2(x) + x = self.neck_ln2(x) + + x2 = self.net_2(x) + x3 = self.net_3(x2) + + return x3 + + +# ============================================================================== +# Qwen2 Decoder as Encoder Components +# ============================================================================== + + +class Qwen2EncoderLayer(nnx.Module): + """Qwen2 decoder layer modified for encoding (supports bidirectional mask).""" + + def __init__( + self, + config: Config, + mesh: Mesh, + quant: Optional[Quant], + rngs: nnx.Rngs, + ): + self.config = config + self.mesh = mesh + self.quant = quant + + self.pre_self_attention_layer_norm = RMSNorm( + num_features=config.emb_dim, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + kernel_axes=("norm",), + epsilon=config.normalization_layer_epsilon, + rngs=rngs, + ) + + query_pre_attn_scalar = config.head_dim**-0.5 + self.self_attention = Attention( + config=config, + num_query_heads=config.num_query_heads, + num_kv_heads=config.num_kv_heads, + head_dim=config.head_dim, + max_target_length=config.max_target_length, + max_prefill_predict_length=config.max_prefill_predict_length, + attention_kernel=config.attention, + inputs_q_shape=(config.per_device_batch_size, config.max_target_length, config.emb_dim), + inputs_kv_shape=(config.per_device_batch_size, config.max_target_length, config.emb_dim), + mesh=mesh, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + dropout_rate=config.dropout_rate, + quant=quant, + kv_quant=quantizations.configure_kv_quant(config), + use_bias_in_projections=config.attention_bias, + query_pre_attn_scalar=query_pre_attn_scalar, + model_mode="train", + rngs=rngs, + ) + + self.post_self_attention_layer_norm = RMSNorm( + num_features=config.emb_dim, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + kernel_axes=("norm",), + epsilon=config.normalization_layer_epsilon, + rngs=rngs, + ) + + self.mlp = MlpBlock( + in_features=config.emb_dim, + intermediate_dim=config.mlp_dim, + activations=config.mlp_activations, + intermediate_dropout_rate=config.dropout_rate, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + config=config, + mesh=mesh, + quant=quant, + model_mode="train", + rngs=rngs, + ) + + def __call__( + self, + inputs: jax.Array, + bidirectional_mask: jax.Array, + decoder_positions: jax.Array, + deterministic: bool, + ) -> jax.Array: + lnx = self.pre_self_attention_layer_norm(inputs) + + attention_lnx, _ = self.self_attention( + lnx, + lnx, + inputs_positions=decoder_positions, + deterministic=deterministic, + model_mode="train", + bidirectional_mask=bidirectional_mask, + ) + + intermediate_inputs = inputs + attention_lnx + hidden_states = self.post_self_attention_layer_norm(intermediate_inputs) + + mlp_lnx = self.mlp(hidden_states, deterministic=deterministic) + layer_output = intermediate_inputs + mlp_lnx + return layer_output + + +class Qwen2Decoder2Encoder(nnx.Module): + """Qwen2 decoder used as an encoder with learnable queries.""" + + def __init__( + self, + config: Config, + mesh: Mesh, + quant: Optional[Quant], + rngs: nnx.Rngs, + ): + self.config = config + + # Create connector config + pydantic_config = config._pydantic_config + update_dict = { + "emb_dim": config.vision_connector_emb_dim, + "num_query_heads": config.vision_connector_num_query_heads, + "num_kv_heads": config.vision_connector_num_kv_heads, + "mlp_dim": config.vision_connector_mlp_dim, + "num_decoder_layers": config.vision_connector_num_layers, + "head_dim": config.vision_connector_emb_dim // config.vision_connector_num_query_heads, + "attention_type": "global", + "rope_max_timescale": 1000000, + } + if hasattr(pydantic_config, "model_copy"): + new_pydantic_config = pydantic_config.model_copy(update=update_dict) + else: + new_pydantic_config = pydantic_config.copy(update=update_dict) + self.connector_config = HyperParameters(new_pydantic_config) + + self.query_768 = nnx.Embed(num_embeddings=144, features=self.connector_config.emb_dim, rngs=rngs) + self.query_1024 = nnx.Embed(num_embeddings=256, features=self.connector_config.emb_dim, rngs=rngs) + self.norm = RMSNorm( + num_features=self.connector_config.emb_dim, + dtype=self.connector_config.dtype, + weight_dtype=self.connector_config.weight_dtype, + kernel_axes=("norm",), + epsilon=self.connector_config.normalization_layer_epsilon, + rngs=rngs, + ) + + self.layers = nnx.List([]) + for i in range(self.connector_config.num_decoder_layers): + layer = Qwen2EncoderLayer( + config=self.connector_config, + mesh=mesh, + quant=quant, + rngs=rngs, + ) + self.layers.append(layer) + setattr(self, f"layer_{i}", layer) + + def __call__(self, x: jax.Array, deterministic: bool = True) -> jax.Array: + B, H, W, C = x.shape + x_flat = x.reshape(B, H * W, C) + + n_query = H * W + if n_query == 144: + queries = self.query_768.embedding + elif n_query == 256: + queries = self.query_1024.embedding + else: + raise ValueError(f"Unsupported query size: {n_query}") + + queries_batch = jnp.broadcast_to(queries[None, :, :], (B, queries.shape[0], queries.shape[1])) + x_combined = jnp.concatenate([x_flat, queries_batch], axis=1) + + mask_image = jnp.ones((B, n_query), dtype=jnp.bool_) + mask_query = jnp.zeros((B, n_query), dtype=jnp.bool_) + bidirectional_mask = jnp.concatenate([mask_image, mask_query], axis=1) + + decoder_positions = jnp.arange(2 * n_query)[None, :] + + y = x_combined + for layer in self.layers: + y = layer(y, bidirectional_mask, decoder_positions, deterministic) + y = self.norm(y) + + return y[:, n_query:, :] + + +# ============================================================================== +# Projector and Top-level Vision Encoder +# ============================================================================== + + +class MlpProjector(nnx.Module): + """Linear projector to map connector features to language model dimension.""" + + def __init__(self, config: Config, mesh: Mesh, rngs: nnx.Rngs): + self.linear = nnx.Linear(config.vision_connector_emb_dim, config.emb_dim, use_bias=True, rngs=rngs) + self.view_seperator = nnx.Param(jax.random.normal(rngs.params(), (config.emb_dim,)) * (config.emb_dim**-0.5)) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.linear(x) + if x.ndim == 4: + separator = jnp.broadcast_to(self.view_seperator.value, (*x.shape[:-2], 1, x.shape[-1])) + return jnp.concatenate([x, separator], axis=-2) + if x.ndim == 3: + separator = jnp.broadcast_to(self.view_seperator.value, (x.shape[0], 1, x.shape[-1])) + return jnp.concatenate([x, separator], axis=-2) + return x + + +class DeepseekOCR2VisionEncoder(nnx.Module): + """Full vision tower for DeepSeek-OCR-2 (SAM + Qwen2).""" + + def __init__(self, config: Config, mesh: Mesh, rngs: nnx.Rngs): + self.sam_model = SAMViTB(config, mesh, rngs) + self.qwen2_model = Qwen2Decoder2Encoder(config, mesh, None, rngs) + self.crop_size = 768 + + def __call__(self, x: jax.Array, deterministic: bool = True) -> jax.Array: + if x.ndim == 6: + B, N, crops, H, W, C = x.shape + if crops < 1: + raise ValueError("DeepSeek-OCR-2 vision input must include the global view.") + + global_images = x[:, :, 0].reshape(B * N, H, W, C) + global_embeddings = self.qwen2_model(self.sam_model(global_images), deterministic) + + if crops == 1: + return global_embeddings.reshape(B, N, global_embeddings.shape[1], global_embeddings.shape[2]) + + crop_images = x[:, :, 1:, : self.crop_size, : self.crop_size, :] + crop_images = crop_images.reshape(B * N * (crops - 1), self.crop_size, self.crop_size, C) + crop_embeddings = self.qwen2_model(self.sam_model(crop_images), deterministic) + crop_tokens = crop_embeddings.shape[1] + crop_dim = crop_embeddings.shape[2] + crop_embeddings = crop_embeddings.reshape(B, N, crops - 1, crop_tokens, crop_dim) + crop_embeddings = crop_embeddings.reshape(B, N, (crops - 1) * crop_tokens, crop_dim) + + return jnp.concatenate( + [crop_embeddings, global_embeddings.reshape(B, N, global_embeddings.shape[1], global_embeddings.shape[2])], + axis=2, + ) + elif x.ndim == 4: + return self.qwen2_model(self.sam_model(x), deterministic) + else: + raise ValueError(f"Expected 4D or 6D input, got {x.ndim}D") diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py index 381695915a..48ee3ac26f 100644 --- a/src/maxtext/multimodal/processor.py +++ b/src/maxtext/multimodal/processor.py @@ -48,6 +48,11 @@ def preprocess_mm_data(config): from maxtext.multimodal.processor_qwen3_omni import preprocess_mm_data_qwen3_omni # pylint: disable=import-outside-toplevel processor_outputs = preprocess_mm_data_qwen3_omni(config) + elif config.model_name in ["deepseek_ocr_2"]: + from maxtext.multimodal.processor_deepseek_ocr import preprocess_mm_data_deepseek_ocr # pylint: disable=import-outside-toplevel + + images = [mm_utils.load_image_from_path(p) for p in config.image_path.split(",")] + processor_outputs = preprocess_mm_data_deepseek_ocr(images) else: raise ValueError(f"Model {config.model_name} not supported for multimodal preprocessing.") @@ -72,6 +77,10 @@ def preprocess_image_for_training(image, config): from maxtext.multimodal.processor_qwen3_omni import preprocess_mm_data_qwen3_omni_for_training # pylint: disable=import-outside-toplevel return preprocess_mm_data_qwen3_omni_for_training(image, config) + elif config.model_name in ["deepseek_ocr_2"]: + from maxtext.multimodal.processor_deepseek_ocr import preprocess_mm_data_deepseek_ocr # pylint: disable=import-outside-toplevel + + return preprocess_mm_data_deepseek_ocr(image) else: raise ValueError(f"Model {config.model_name} not supported for image preprocessing.") @@ -94,6 +103,10 @@ def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | No from maxtext.multimodal.processor_qwen3_omni import get_mm_offsets_qwen3_omni # pylint: disable=import-outside-toplevel return get_mm_offsets_qwen3_omni(config, processor_output) + elif config.model_name in ["deepseek_ocr_2"]: + from maxtext.multimodal.processor_deepseek_ocr import get_image_offsets_deepseek_ocr # pylint: disable=import-outside-toplevel + + return get_image_offsets_deepseek_ocr(processor_output) else: return 0 @@ -122,6 +135,10 @@ def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_pla video_placeholder=video_placeholder, num_videos=num_videos, ) + elif model_name in ["deepseek_ocr_2"]: + from maxtext.multimodal.processor_deepseek_ocr import reformat_prompt_deepseek_ocr # pylint: disable=import-outside-toplevel + + return reformat_prompt_deepseek_ocr(prompt, image_placeholder, num_images) else: return prompt @@ -162,6 +179,10 @@ def prepare_text_for_image_fusion(tokens, config, processor_output=None): from maxtext.multimodal.processor_qwen3_omni import add_extra_tokens_for_qwen3_omni # pylint: disable=import-outside-toplevel return add_extra_tokens_for_qwen3_omni(tokens, config, processor_output) + elif config.model_name in ["deepseek_ocr_2"]: + from maxtext.multimodal.processor_deepseek_ocr import add_extra_tokens_for_images_deepseek_ocr # pylint: disable=import-outside-toplevel + + return add_extra_tokens_for_images_deepseek_ocr(tokens, processor_output) else: raise ValueError(f"Model {config.model_name} does not support multimodal inference.") @@ -185,6 +206,10 @@ def get_dummy_image_shape_for_init(model_name, batch_size=1, num_image_per_seque from maxtext.multimodal.processor_qwen3_omni import get_dummy_image_shape_for_init_qwen3_omni # pylint: disable=import-outside-toplevel image_shape = get_dummy_image_shape_for_init_qwen3_omni(batch_size) + elif model_name in ["deepseek_ocr_2"]: + from maxtext.multimodal.processor_deepseek_ocr import get_dummy_image_shape_for_init_deepseek_ocr # pylint: disable=import-outside-toplevel + + image_shape = get_dummy_image_shape_for_init_deepseek_ocr(batch_size, num_image_per_sequence) return image_shape @@ -231,6 +256,10 @@ def get_bidirectional_mask_vision(config, decoder_input_tokens, is_video: bool = bidirectional_mask_vision = decoder_input_tokens == tokens.video_pad else: bidirectional_mask_vision = decoder_input_tokens == tokens.image_pad + elif config.model_name in ["deepseek_ocr_2"]: + from maxtext.multimodal.processor_deepseek_ocr import DEEPSEEK_OCR_IMAGE_TOKEN_ID # pylint: disable=import-outside-toplevel + + bidirectional_mask_vision = decoder_input_tokens == DEEPSEEK_OCR_IMAGE_TOKEN_ID return bidirectional_mask_vision diff --git a/src/maxtext/multimodal/processor_deepseek_ocr.py b/src/maxtext/multimodal/processor_deepseek_ocr.py new file mode 100644 index 0000000000..3ff3b7f9e9 --- /dev/null +++ b/src/maxtext/multimodal/processor_deepseek_ocr.py @@ -0,0 +1,299 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DeepSeek-OCR-2-specific utilities for multimodal features.""" + +from dataclasses import dataclass +import numpy as np +from PIL import Image, ImageOps + +from maxtext.multimodal import utils as mm_utils + +# Constants for DeepSeek-OCR-2 +DEEPSEEK_OCR_IMAGE_TOKEN_ID = 128815 +DEEPSEEK_OCR_BASE_SIZE = 1024 +DEEPSEEK_OCR_IMAGE_SIZE = 768 +DEEPSEEK_OCR_MAX_CROPS = 6 +DEEPSEEK_OCR_IMAGE_PLACEHOLDER_IN_PROMPT = "" + +DEEPSEEK_OCR_GLOBAL_TOKENS = 256 +DEEPSEEK_OCR_CROP_TOKENS = 144 +DEEPSEEK_OCR_SEPARATOR_TOKENS = 1 +DEEPSEEK_OCR_NUM_TOKENS_PER_IMAGE = ( + DEEPSEEK_OCR_CROP_TOKENS * DEEPSEEK_OCR_MAX_CROPS + DEEPSEEK_OCR_GLOBAL_TOKENS + DEEPSEEK_OCR_SEPARATOR_TOKENS +) + + +@dataclass +class DeepseekOCR2PreprocessorOutput(mm_utils.PreprocessorOutput): + """Holds the output of DeepSeek-OCR-2 image preprocessor.""" + + pixel_values: None | np.ndarray = None + pixel_mask: None | np.ndarray = None + aspect_ratios: None | np.ndarray = None + num_images: int = 0 + + +def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size): + best_ratio_diff = float("inf") + best_ratio = (1, 1) + area = width * height + for ratio in target_ratios: + target_aspect_ratio = ratio[0] / ratio[1] + ratio_diff = abs(aspect_ratio - target_aspect_ratio) + if ratio_diff < best_ratio_diff: + best_ratio_diff = ratio_diff + best_ratio = ratio + elif ratio_diff == best_ratio_diff: + if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: + best_ratio = ratio + return best_ratio + + +def dynamic_preprocess(image, min_num=2, max_num=6, image_size=768, use_thumbnail=False): + orig_width, orig_height = image.size + aspect_ratio = orig_width / orig_height + + target_ratios = set( + (i, j) + for n in range(min_num, max_num + 1) + for i in range(1, n + 1) + for j in range(1, n + 1) + if i * j <= max_num and i * j >= min_num + ) + target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) + + target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size) + + target_width = image_size * target_aspect_ratio[0] + target_height = image_size * target_aspect_ratio[1] + blocks = target_aspect_ratio[0] * target_aspect_ratio[1] + + resized_img = image.resize((target_width, target_height)) + processed_images = [] + for i in range(blocks): + box = ( + (i % (target_width // image_size)) * image_size, + (i // (target_width // image_size)) * image_size, + ((i % (target_width // image_size)) + 1) * image_size, + ((i // (target_width // image_size)) + 1) * image_size, + ) + split_img = resized_img.crop(box) + processed_images.append(split_img) + assert len(processed_images) == blocks + if use_thumbnail and len(processed_images) != 1: + thumbnail_img = image.resize((image_size, image_size)) + processed_images.append(thumbnail_img) + return processed_images, target_aspect_ratio + + +def preprocess_mm_data_deepseek_ocr( + images, base_size=DEEPSEEK_OCR_BASE_SIZE, image_size=DEEPSEEK_OCR_IMAGE_SIZE, crop_mode=True +): + """Preprocesses images for DeepSeek-OCR-2.""" + images_in = [] + if isinstance(images, np.ndarray): + images_in.append(images) + elif isinstance(images, list): + images_in.extend(images) + else: + images_in.append(images) + + out_pixel_values = [] + out_pixel_mask = [] + out_aspect_ratios = [] + + for img in images_in: + if isinstance(img, np.ndarray): + # If it is HWC, convert to PIL + if len(img.shape) == 3: + pil_img = Image.fromarray(img) + else: + raise ValueError(f"Unsupported numpy array shape: {img.shape}") + elif isinstance(img, Image.Image): + pil_img = img + else: + # Try to load if it is a path (though usually loaded before) + pil_img = mm_utils.load_image_from_path(img) + if isinstance(pil_img, np.ndarray): + pil_img = Image.fromarray(pil_img) + pil_img = pil_img.convert("RGB") + + # Global view + global_view = ImageOps.pad(pil_img, (base_size, base_size), color=(127, 127, 127)) + global_tensor = np.array(global_view, dtype=np.float32) / 255.0 + global_tensor = (global_tensor - 0.5) / 0.5 + + crops = [] + crop_ratio = [1, 1] + if crop_mode: + if pil_img.size[0] <= 768 and pil_img.size[1] <= 768: + crop_ratio = [1, 1] + else: + crops_raw, crop_ratio = dynamic_preprocess( + pil_img, min_num=2, max_num=DEEPSEEK_OCR_MAX_CROPS, image_size=image_size + ) + for crop in crops_raw: + crop_tensor = np.array(crop, dtype=np.float32) / 255.0 + crop_tensor = (crop_tensor - 0.5) / 0.5 + padded_crop = np.zeros((base_size, base_size, 3), dtype=np.float32) + padded_crop[:image_size, :image_size, :] = crop_tensor + crops.append(padded_crop) + + num_crops = len(crops) + while len(crops) < DEEPSEEK_OCR_MAX_CROPS: + crops.append(np.zeros((base_size, base_size, 3), dtype=np.float32)) + + img_pixel_values = np.stack([global_tensor] + crops, axis=0) + + img_pixel_mask = np.zeros((DEEPSEEK_OCR_NUM_TOKENS_PER_IMAGE,), dtype=np.bool_) + for i in range(num_crops): + start = i * DEEPSEEK_OCR_CROP_TOKENS + img_pixel_mask[start : start + DEEPSEEK_OCR_CROP_TOKENS] = True + global_start = DEEPSEEK_OCR_CROP_TOKENS * DEEPSEEK_OCR_MAX_CROPS + img_pixel_mask[global_start : global_start + DEEPSEEK_OCR_GLOBAL_TOKENS + DEEPSEEK_OCR_SEPARATOR_TOKENS] = True + + out_pixel_values.append(img_pixel_values) + out_pixel_mask.append(img_pixel_mask) + out_aspect_ratios.append(crop_ratio) + + return DeepseekOCR2PreprocessorOutput( + pixel_values=np.stack(out_pixel_values, axis=0), + pixel_mask=np.stack(out_pixel_mask, axis=0), + aspect_ratios=np.array(out_aspect_ratios, dtype=np.int32), + num_images=len(images_in), + ) + + +def get_image_offsets_deepseek_ocr(processor_output: mm_utils.PreprocessorOutput | None): + """Get the increase in total token count after inserting image token placeholders.""" + has_images = processor_output is not None and processor_output.pixel_values is not None + if not has_images: + return DEEPSEEK_OCR_NUM_TOKENS_PER_IMAGE - 1 + if processor_output.pixel_mask is None: + return (DEEPSEEK_OCR_NUM_TOKENS_PER_IMAGE - 1) * processor_output.pixel_values.shape[0] + return int(np.sum(processor_output.pixel_mask, dtype=np.int32) - processor_output.pixel_values.shape[0]) + + +def reformat_prompt_deepseek_ocr(prompt, image_placeholder, num_images): + """Reformat prompt for DeepSeek-OCR-2 (plain SFT, no wrapping).""" + prompt = prompt.replace("\\n", "\n") + if image_placeholder in prompt: + prompt = prompt.replace(image_placeholder, DEEPSEEK_OCR_IMAGE_PLACEHOLDER_IN_PROMPT) + image_placeholder_count = prompt.count(DEEPSEEK_OCR_IMAGE_PLACEHOLDER_IN_PROMPT) + if image_placeholder_count < num_images: + prompt = DEEPSEEK_OCR_IMAGE_PLACEHOLDER_IN_PROMPT * (num_images - image_placeholder_count) + prompt + # The user verification script uses: "\n<|grounding|>Convert the document to markdown. " + # We don't need to add chat templates if we want to match the verification. + return prompt + + +def add_extra_tokens_for_images_deepseek_ocr(tokens, processor_output: mm_utils.PreprocessorOutput | None = None): + """Inserts image placeholder tokens into the token list.""" + if processor_output is not None and processor_output.pixel_mask is not None: + token_counts = np.sum(processor_output.pixel_mask, axis=-1, dtype=np.int32).tolist() + else: + num_images = processor_output.num_images if processor_output is not None else 1 + token_counts = [DEEPSEEK_OCR_NUM_TOKENS_PER_IMAGE] * num_images + return insert_variable_sequences(tokens, at=DEEPSEEK_OCR_IMAGE_TOKEN_ID, sequence_lengths=token_counts) + + +def get_dummy_image_shape_for_init_deepseek_ocr(batch_size=1, num_image_per_sequence=1): + """Return the shape of the dummy image for initialization.""" + return ( + batch_size, + num_image_per_sequence, + 1 + DEEPSEEK_OCR_MAX_CROPS, + DEEPSEEK_OCR_BASE_SIZE, + DEEPSEEK_OCR_BASE_SIZE, + 3, + ) + + +# Helper for insertion (copied from processor_gemma3.py to avoid import issues if it moves) +def _get_new_text_positions(offset_on: np.ndarray, offset_by: int) -> np.ndarray: + offset = np.cumsum(offset_on, axis=-1) * offset_by + new_positions = np.arange(offset_on.shape[-1]) + offset + new_positions -= offset_by * offset_on + return new_positions + + +def insert_sequence(tokens: np.ndarray, at: int, sequence: list[int], max_num_images: int) -> np.ndarray: + """Inserts a sequence of tokens at all occurrences of a specific token `at`.""" + is_1d = len(tokens.shape) == 1 + if is_1d: + tokens = np.expand_dims(tokens, axis=0) + + batch_size, seq_len = tokens.shape + sequence_len = len(sequence) + offset_by = sequence_len - 1 + + # Find where the placeholder tokens are + offset_on = tokens == at + + # Calculate new positions for all tokens + new_positions = _get_new_text_positions(offset_on=offset_on, offset_by=offset_by) + + # Allocate new token array + new_seq_len = seq_len + offset_by * max_num_images + new_tokens = np.zeros((batch_size, new_seq_len), dtype=tokens.dtype) + + # Place old tokens in their new positions + # We use advanced indexing to scatter + batch_indices = np.arange(batch_size)[:, None] + new_positions_clamped = np.clip(new_positions, 0, new_seq_len - 1) + new_tokens[batch_indices, new_positions_clamped] = tokens + + # Fill in the inserted sequences + # We find the new positions of the `at` token + for b in range(batch_size): + at_indices = np.where(offset_on[b])[0] + for idx in at_indices: + start_pos = new_positions[b, idx] + new_tokens[b, start_pos : start_pos + sequence_len] = sequence + + if is_1d: + new_tokens = np.squeeze(new_tokens, axis=0) + + return new_tokens + + +def insert_variable_sequences(tokens: np.ndarray, at: int, sequence_lengths: list[int]) -> np.ndarray: + """Replaces each image token with the matching number of DeepSeek visual placeholders.""" + is_1d = len(tokens.shape) == 1 + if is_1d: + tokens = np.expand_dims(tokens, axis=0) + + rows = [] + for row in tokens: + image_index = 0 + pieces = [] + for token in row: + if token == at: + sequence_len = sequence_lengths[min(image_index, len(sequence_lengths) - 1)] + pieces.extend([at] * sequence_len) + image_index += 1 + else: + pieces.append(token) + rows.append(np.array(pieces, dtype=tokens.dtype)) + + max_len = max(row.shape[0] for row in rows) + new_tokens = np.zeros((len(rows), max_len), dtype=tokens.dtype) + for i, row in enumerate(rows): + new_tokens[i, : row.shape[0]] = row + + if is_1d: + new_tokens = np.squeeze(new_tokens, axis=0) + + return new_tokens diff --git a/src/maxtext/multimodal/utils.py b/src/maxtext/multimodal/utils.py index 65b5670fc1..681e5336d8 100644 --- a/src/maxtext/multimodal/utils.py +++ b/src/maxtext/multimodal/utils.py @@ -163,7 +163,10 @@ def merge_mm_embeddings( # Expand the tile-level mask to a token-level mask to match the embeddings. # A mask of shape (B, N*T) becomes (B, N*T*K) by repeating each element K times. - flat_token_masks_processed = jnp.repeat(flat_tile_masks, repeats=num_toks_per_token, axis=1) + if flat_tile_masks.shape[1] == flat_multimodal_embeddings.shape[1]: + flat_token_masks_processed = flat_tile_masks + else: + flat_token_masks_processed = jnp.repeat(flat_tile_masks, repeats=num_toks_per_token, axis=1) # Vmap the inner merge function over the batch dimension return jax.vmap( @@ -183,7 +186,7 @@ def _merge_mm_embeddings_inner( if token_mask is not None: # This logic packs valid multimodal tokens to the front of the array. # It correctly handles cases where some multimodal tokens are just padding. - sort_indices = jnp.argsort(-token_mask) # Sorts descending, putting 1s first + sort_indices = jnp.argsort(token_mask, descending=True) # Sorts descending, putting 1s first multimodal_embeddings = multimodal_embeddings[sort_indices] # Find positions in the text sequence to place the multimodal embeddings. diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index 4167eb0e88..1c429695ba 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -75,6 +75,7 @@ "qwen3-235b-a22b": "Qwen/Qwen3-235B-A22B-Thinking-2507", "qwen3-480b-a35b": "Qwen/Qwen3-Coder-480B-A35B-Instruct", "deepseek2-16b": "deepseek-ai/DeepSeek-V2-Lite", + "deepseek_ocr_2": "deepseek-ai/DeepSeek-OCR-2", "deepseek3-671b": "deepseek-ai/DeepSeek-V3", "deepseek3.2-671b": "deepseek-ai/DeepSeek-V3.2", "deepseek4-284b": "deepseek-ai/DeepSeek-V4-Flash", diff --git a/tests/assets/logits_generation/generate_hf_golden_logits.py b/tests/assets/logits_generation/generate_hf_golden_logits.py index c57d58c380..f0b531c6dc 100644 --- a/tests/assets/logits_generation/generate_hf_golden_logits.py +++ b/tests/assets/logits_generation/generate_hf_golden_logits.py @@ -81,9 +81,14 @@ def save_golden_logits( model_class = Llama4ForConditionalGeneration else: - from transformers import AutoModelForCausalLM # pylint: disable=import-outside-toplevel + if model_id == "deepseek-ai/DeepSeek-OCR-2": + from transformers import AutoModel # pylint: disable=import-outside-toplevel - model_class = AutoModelForCausalLM + model_class = AutoModel + else: + from transformers import AutoModelForCausalLM # pylint: disable=import-outside-toplevel + + model_class = AutoModelForCausalLM tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=trust_remote_code) print(f"loading model from {hf_model_path}") @@ -100,8 +105,22 @@ def save_golden_logits( else: raise ValueError(f"unsupported --hf-load-dtype: {hf_load_dtype}") + from transformers import AutoConfig, LlamaConfig # pylint: disable=import-outside-toplevel + + config = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=trust_remote_code) + + # Fill in missing attributes with Llama defaults to satisfy LlamaAttention expectations + llama_default = LlamaConfig() + for k, v in llama_default.to_dict().items(): + if not hasattr(config, k): + setattr(config, k, v) + + if not hasattr(config, "pad_token_id") or config.pad_token_id is None: + config.pad_token_id = tokenizer.pad_token_id + model = model_class.from_pretrained( hf_model_path, + config=config, dtype=torch_dtype, trust_remote_code=trust_remote_code, ) diff --git a/tests/unit/deepseek_ocr_layers_test.py b/tests/unit/deepseek_ocr_layers_test.py new file mode 100644 index 0000000000..46862b78b6 --- /dev/null +++ b/tests/unit/deepseek_ocr_layers_test.py @@ -0,0 +1,1418 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for DeepSeek-OCR-2 layers and multimodal plumbing.""" + +import os +import sys +import unittest +from typing import Optional, Tuple + +os.environ.setdefault("JAX_PLATFORMS", "cpu") +os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") + +import torch +from torch import nn +import torch.nn.functional as F +import numpy as np + +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +from flax import nnx + +from maxtext.models import deepseek_ocr +from maxtext.models import deepseek +from maxtext.multimodal import utils as mm_utils +from maxtext.multimodal.processor_deepseek_ocr import ( + DEEPSEEK_OCR_CROP_TOKENS, + DEEPSEEK_OCR_GLOBAL_TOKENS, + DEEPSEEK_OCR_IMAGE_TOKEN_ID, + DEEPSEEK_OCR_SEPARATOR_TOKENS, + DeepseekOCR2PreprocessorOutput, + add_extra_tokens_for_images_deepseek_ocr, + get_image_offsets_deepseek_ocr, +) +from maxtext.configs import pyconfig +from maxtext.common.common_types import DecoderBlockType, MODEL_MODE_TRAIN +from PIL import Image, ImageOps + +import transformers +from transformers.models.llama.configuration_llama import LlamaConfig +from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRotaryEmbedding +from transformers.models.qwen2.configuration_qwen2 import Qwen2Config as PTQwen2Config +from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer as PTQwen2DecoderLayer + +torch.set_grad_enabled(False) + +# ============================================================================== +# Helper functions for HF Image Preprocessing Reference +# ============================================================================== + + +def _normalize_hwc(image): + arr = np.array(image, dtype=np.float32) / 255.0 + return (arr - 0.5) / 0.5 + + +def _find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size): + best_ratio_diff = float("inf") + best_ratio = (1, 1) + area = width * height + for ratio in target_ratios: + target_aspect_ratio = ratio[0] / ratio[1] + ratio_diff = abs(aspect_ratio - target_aspect_ratio) + if ratio_diff < best_ratio_diff: + best_ratio_diff = ratio_diff + best_ratio = ratio + elif ratio_diff == best_ratio_diff: + if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: + best_ratio = ratio + return best_ratio + + +def _hf_dynamic_preprocess(image, min_num=2, max_num=6, image_size=768): + orig_width, orig_height = image.size + aspect_ratio = orig_width / orig_height + target_ratios = set( + (i, j) + for n in range(min_num, max_num + 1) + for i in range(1, n + 1) + for j in range(1, n + 1) + if i * j <= max_num and i * j >= min_num + ) + target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) + target_aspect_ratio = _find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size) + target_width = image_size * target_aspect_ratio[0] + target_height = image_size * target_aspect_ratio[1] + blocks = target_aspect_ratio[0] * target_aspect_ratio[1] + + resized_img = image.resize((target_width, target_height)) + processed_images = [] + for i in range(blocks): + box = ( + (i % (target_width // image_size)) * image_size, + (i // (target_width // image_size)) * image_size, + ((i % (target_width // image_size)) + 1) * image_size, + ((i // (target_width // image_size)) + 1) * image_size, + ) + processed_images.append(resized_img.crop(box)) + return processed_images, target_aspect_ratio + + +def _hf_reference_inputs(image, base_size=1024, image_size=768): + if image.size[0] <= image_size and image.size[1] <= image_size: + crops_raw = [] + crop_ratio = [1, 1] + else: + crops_raw, crop_ratio = _hf_dynamic_preprocess(image, min_num=2, max_num=6, image_size=image_size) + + global_view = ImageOps.pad(image, (base_size, base_size), color=(127, 127, 127)) + images_ori = _normalize_hwc(global_view) + images_crop = np.stack([_normalize_hwc(crop) for crop in crops_raw], axis=0) if crops_raw else np.zeros((0, image_size, image_size, 3), dtype=np.float32) + + return { + "images_ori": images_ori, + "images_crop": images_crop, + "crop_ratio": np.array(crop_ratio, dtype=np.int32), + } + +# ============================================================================== +# Helper functions for weight copying +# ============================================================================== + + +def to_jax(pt_tensor: torch.Tensor) -> jax.Array: + return jnp.asarray(pt_tensor.detach().cpu().numpy()) + + +def copy_linear_weights(torch_linear, jax_linear): + jax_linear.kernel.value = jnp.array(torch_linear.weight.detach().cpu().numpy().T) + if torch_linear.bias is not None and jax_linear.bias is not None: + jax_linear.bias.value = jnp.array(torch_linear.bias.detach().cpu().numpy()) + + +def copy_rmsnorm_weights(torch_norm, jax_norm): + if hasattr(torch_norm, "weight") and hasattr(jax_norm, "scale"): + jax_norm.scale.value = jnp.array(torch_norm.weight.detach().cpu().numpy()) + + +def copy_qwen2_attention_weights(torch_attn, jax_attn): + num_heads = jax_attn.num_query_heads + num_kv_heads = jax_attn.num_kv_heads + head_dim = jax_attn.head_dim + hidden_size = num_heads * head_dim + kv_dim = num_kv_heads * head_dim + output_dim = hidden_size + + q_weight = torch_attn.q_proj.weight.detach().cpu().numpy() + k_weight = torch_attn.k_proj.weight.detach().cpu().numpy() + v_weight = torch_attn.v_proj.weight.detach().cpu().numpy() + + q_bias = torch_attn.q_proj.bias.detach().cpu().numpy() if torch_attn.q_proj.bias is not None else np.zeros(hidden_size) + k_bias = torch_attn.k_proj.bias.detach().cpu().numpy() if torch_attn.k_proj.bias is not None else np.zeros(kv_dim) + v_bias = torch_attn.v_proj.bias.detach().cpu().numpy() if torch_attn.v_proj.bias is not None else np.zeros(kv_dim) + + jax_attn.query.kernel.value = jnp.array(q_weight.T.reshape(hidden_size, num_heads, head_dim)) + if jax_attn.query.bias is not None: + jax_attn.query.bias.value = jnp.array(q_bias.reshape(num_heads, head_dim)) + + jax_attn.key.kernel.value = jnp.array(k_weight.T.reshape(hidden_size, num_kv_heads, head_dim)) + if jax_attn.key.bias is not None: + jax_attn.key.bias.value = jnp.array(k_bias.reshape(num_kv_heads, head_dim)) + + jax_attn.value.kernel.value = jnp.array(v_weight.T.reshape(hidden_size, num_kv_heads, head_dim)) + if jax_attn.value.bias is not None: + jax_attn.value.bias.value = jnp.array(v_bias.reshape(num_kv_heads, head_dim)) + + out_weight = torch_attn.o_proj.weight.detach().cpu().numpy() + jax_attn.out.kernel.value = jnp.array(out_weight.T.reshape(num_heads, head_dim, output_dim)) + if torch_attn.o_proj.bias is not None and jax_attn.out.bias is not None: + jax_attn.out.bias.value = jnp.array(torch_attn.o_proj.bias.detach().cpu().numpy()) + + +def copy_llama_attention_weights(torch_attn, jax_attn): + num_heads = jax_attn.num_query_heads + num_kv_heads = jax_attn.num_kv_heads + head_dim = jax_attn.head_dim + hidden_size = num_heads * head_dim + + q_weight = torch_attn.q_proj.weight.detach().cpu().numpy() + k_weight = torch_attn.k_proj.weight.detach().cpu().numpy() + v_weight = torch_attn.v_proj.weight.detach().cpu().numpy() + out_weight = torch_attn.o_proj.weight.detach().cpu().numpy() + + q_scale = head_dim**-0.5 + jax_attn.query.kernel.value = jnp.asarray((q_weight.T * q_scale).reshape(hidden_size, num_heads, head_dim)) + jax_attn.key.kernel.value = jnp.asarray(k_weight.T.reshape(hidden_size, num_kv_heads, head_dim)) + jax_attn.value.kernel.value = jnp.asarray(v_weight.T.reshape(hidden_size, num_kv_heads, head_dim)) + jax_attn.out.kernel.value = jnp.asarray(out_weight.T.reshape(num_heads, head_dim, hidden_size)) + + if torch_attn.q_proj.bias is not None and jax_attn.query.bias is not None: + jax_attn.query.bias.value = jnp.asarray(torch_attn.q_proj.bias.detach().cpu().numpy().reshape(num_heads, head_dim)) + if torch_attn.k_proj.bias is not None and jax_attn.key.bias is not None: + jax_attn.key.bias.value = jnp.asarray(torch_attn.k_proj.bias.detach().cpu().numpy().reshape(num_kv_heads, head_dim)) + if torch_attn.v_proj.bias is not None and jax_attn.value.bias is not None: + jax_attn.value.bias.value = jnp.asarray(torch_attn.v_proj.bias.detach().cpu().numpy().reshape(num_kv_heads, head_dim)) + if torch_attn.o_proj.bias is not None and jax_attn.out.bias is not None: + jax_attn.out.bias.value = jnp.asarray(torch_attn.o_proj.bias.detach().cpu().numpy()) + + +def copy_llama_decoder_layer_weights(torch_layer, jax_layer): + copy_rmsnorm_weights(torch_layer.input_layernorm, jax_layer.pre_self_attention_layer_norm) + copy_rmsnorm_weights(torch_layer.post_attention_layernorm, jax_layer.post_self_attention_layer_norm) + copy_llama_attention_weights(torch_layer.self_attn, jax_layer.self_attention) + copy_linear_weights(torch_layer.mlp.gate_proj, jax_layer.mlp.wi_0) + copy_linear_weights(torch_layer.mlp.up_proj, jax_layer.mlp.wi_1) + copy_linear_weights(torch_layer.mlp.down_proj, jax_layer.mlp.wo) + + +def make_llama_decoder_layer(hidden_size, num_heads, num_kv_heads, intermediate_size, seq_len): + config = LlamaConfig( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_hidden_layers=1, + num_attention_heads=num_heads, + num_key_value_heads=num_kv_heads, + vocab_size=128, + max_position_embeddings=seq_len, + rms_norm_eps=1e-6, + attention_dropout=0.0, + hidden_act="silu", + attention_bias=False, + ) + config._attn_implementation = "eager" # pylint: disable=protected-access + return LlamaDecoderLayer(config, layer_idx=0).eval(), LlamaRotaryEmbedding(config) + + +# ============================================================================== +# PyTorch Reference Implementation (from deepencoderv2.py) +# ============================================================================== + + +def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: + max_rel_dist = int(2 * max(q_size, k_size) - 1) + if rel_pos.shape[0] != max_rel_dist: + dtype = rel_pos.dtype + rel_pos = rel_pos.to(torch.float32) + rel_pos_resized = F.interpolate( + rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), + size=max_rel_dist, + mode="linear", + ).to(dtype) + rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) + else: + rel_pos_resized = rel_pos + + q_coords = torch.arange(q_size, device=rel_pos.device)[:, None] * max(k_size / q_size, 1.0) + k_coords = torch.arange(k_size, device=rel_pos.device)[None, :] * max(q_size / k_size, 1.0) + relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) + + return rel_pos_resized[relative_coords.long()] + + +def add_decomposed_rel_pos( + q: torch.Tensor, + rel_pos_h: torch.Tensor, + rel_pos_w: torch.Tensor, + q_size: Tuple[int, int], + k_size: Tuple[int, int], +) -> Tuple[torch.Tensor, torch.Tensor]: + q_h, q_w = q_size + k_h, k_w = k_size + Rh = get_rel_pos(q_h, k_h, rel_pos_h) + Rw = get_rel_pos(q_w, k_w, rel_pos_w) + + B, _, dim = q.shape + r_q = q.reshape(B, q_h, q_w, dim) + rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) + rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) + rel_h = rel_h.unsqueeze(-1) + rel_w = rel_w.unsqueeze(-2) + rel_h = rel_h.reshape(B, q_h * q_w, k_h, 1) + rel_w = rel_w.reshape(B, q_h * q_w, 1, k_w) + + return rel_h, rel_w + + +class PTPatchEmbed(nn.Module): + + def __init__( + self, + kernel_size: Tuple[int, int] = (16, 16), + stride: Tuple[int, int] = (16, 16), + in_chans: int = 3, + embed_dim: int = 768, + ) -> None: + super().__init__() + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) + x = x.permute(0, 2, 3, 1) + return x + + +class PTMLPBlock(nn.Module): + + def __init__(self, embedding_dim: int, mlp_dim: int) -> None: + super().__init__() + self.lin1 = nn.Linear(embedding_dim, mlp_dim) + self.lin2 = nn.Linear(mlp_dim, embedding_dim) + self.act = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lin2(self.act(self.lin1(x))) + + +class PTAttention(nn.Module): + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = True, + use_rel_pos: bool = False, + input_size: Optional[Tuple[int, int]] = None, + ) -> None: + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.proj = nn.Linear(dim, dim) + + self.use_rel_pos = use_rel_pos + if self.use_rel_pos: + assert input_size is not None + self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) + self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, H, W, _ = x.shape + qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) + q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0) + + rel_h, rel_w = None, None + if self.use_rel_pos: + rel_h, rel_w = add_decomposed_rel_pos(q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)) + + q = q.view(B, self.num_heads, H * W, -1) + k = k.view(B, self.num_heads, H * W, -1) + v = v.view(B, self.num_heads, H * W, -1) + + if self.use_rel_pos: + rel_h = rel_h.view(B, self.num_heads, rel_h.size(1), rel_h.size(2), rel_h.size(3)) + rel_w = rel_w.view(B, self.num_heads, rel_w.size(1), rel_w.size(2), rel_w.size(3)) + attn_bias = (rel_h + rel_w).view(B, self.num_heads, rel_h.size(2), rel_h.size(3) * rel_w.size(4)) + x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias) + else: + x = torch.nn.functional.scaled_dot_product_attention(q, k, v) + + x = x.view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1) + x = self.proj(x) + return x + + +def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]: + B, H, W, C = x.shape + pad_h = (window_size - H % window_size) % window_size + pad_w = (window_size - W % window_size) % window_size + if pad_h > 0 or pad_w > 0: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + Hp, Wp = H + pad_h, W + pad_w + x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows, (Hp, Wp) + + +def window_unpartition( + windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int] +) -> torch.Tensor: + Hp, Wp = pad_hw + H, W = hw + B = windows.shape[0] // (Hp * Wp // window_size // window_size) + x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1) + if Hp > H or Wp > W: + x = x[:, :H, :W, :].contiguous() + return x + + +class PTBlock(nn.Module): + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + window_size: int = 0, + input_size: Optional[Tuple[int, int]] = None, + ) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = PTAttention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + use_rel_pos=True, + input_size=input_size if window_size == 0 else (window_size, window_size), + ) + self.norm2 = nn.LayerNorm(dim) + self.mlp = PTMLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio)) + self.window_size = window_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shortcut = x + x = self.norm1(x) + if self.window_size > 0: + H, W = x.shape[1], x.shape[2] + x, pad_hw = window_partition(x, self.window_size) + + x = self.attn(x) + if self.window_size > 0: + x = window_unpartition(x, self.window_size, pad_hw, (H, W)) + + x = shortcut + x + x = x + self.mlp(self.norm2(x)) + return x + + +class PTLayerNorm2d(nn.Module): + + def __init__(self, num_channels: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(num_channels)) + self.bias = nn.Parameter(torch.zeros(num_channels)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + + +class PTImageEncoderViT(nn.Module): + + def __init__( + self, + img_size: int = 1024, + patch_size: int = 16, + in_chans: int = 3, + embed_dim: int = 768, + depth: int = 12, + num_heads: int = 12, + mlp_ratio: float = 4.0, + out_chans: int = 256, + qkv_bias: bool = True, + window_size: int = 14, + global_attn_indexes: Tuple[int, ...] = (2, 5, 8, 11), + ) -> None: + super().__init__() + self.patch_embed = PTPatchEmbed( + kernel_size=(patch_size, patch_size), + stride=(patch_size, patch_size), + in_chans=in_chans, + embed_dim=embed_dim, + ) + self.pos_embed = nn.Parameter(torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim)) + self.blocks = nn.ModuleList() + for i in range(depth): + block = PTBlock( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + window_size=window_size if i not in global_attn_indexes else 0, + input_size=(img_size // patch_size, img_size // patch_size), + ) + self.blocks.append(block) + + self.neck = nn.Sequential( + nn.Conv2d(embed_dim, out_chans, kernel_size=1, bias=False), + PTLayerNorm2d(out_chans), + nn.Conv2d(out_chans, out_chans, kernel_size=3, padding=1, bias=False), + PTLayerNorm2d(out_chans), + ) + self.net_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1, bias=False) + self.net_3 = nn.Conv2d(512, 896, kernel_size=3, stride=2, padding=1, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.patch_embed(x) + tgt_size = x.shape[1] + src_size = self.pos_embed.shape[1] + if src_size != tgt_size: + pos_embed_resized = self.pos_embed.permute(0, 3, 1, 2) + pos_embed_resized = F.interpolate( + pos_embed_resized, size=(tgt_size, tgt_size), mode="bicubic", align_corners=False + ) + pos_embed_resized = pos_embed_resized.permute(0, 2, 3, 1) + else: + pos_embed_resized = self.pos_embed + x = x + pos_embed_resized + for blk in self.blocks: + x = blk(x) + x = self.neck(x.permute(0, 3, 1, 2)) + x2 = self.net_2(x) + x3 = self.net_3(x2) + return x3 + + +class PTCustomQwen2Decoder(nn.Module): + + def __init__( + self, + decoder_layer: int = 24, + max_position_embeddings: int = 131072, + hidden_dimension: int = 896, + num_attention_heads: int = 14, + num_key_value_heads: int = 2, + intermediate_size: int = 4864, + vocab_size: int = 151936, + attn_implementation: str = "sdpa", + rms_norm_eps: float = 1e-06, + rope_theta: float = 1000000.0, + attention_dropout: float = 0.0, + hidden_act: str = "silu", + ): + super().__init__() + Qwen2Model = getattr(transformers.models.qwen2.modeling_qwen2, "Qwen2Model") + Qwen2Config = getattr(transformers, "Qwen2Config") + + config = Qwen2Config( + hidden_size=hidden_dimension, + num_hidden_layers=decoder_layer, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + intermediate_size=intermediate_size, + max_position_embeddings=max_position_embeddings, + vocab_size=vocab_size, + rms_norm_eps=rms_norm_eps, + rope_theta=rope_theta, + attention_dropout=attention_dropout, + hidden_act=hidden_act, + _attn_implementation=attn_implementation, + ) + + self.model = self._create_custom_model(Qwen2Model, config) + del self.model.embed_tokens + + def _create_custom_model(self, Qwen2Model, config): + class CustomQwen2ModelInner(Qwen2Model): + + def forward( + self, + inputs_embeds=None, + attention_mask=None, + position_ids=None, + past_key_values=None, + token_type_ids=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + cache_position=None, + ): + self._current_token_type_ids = token_type_ids + if token_type_ids is not None and not isinstance(attention_mask, dict): + attention_mask = { + "full_attention": self._update_causal_mask( + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, + ) + } + return super().forward( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + def _update_causal_mask( + self, + attention_mask, + input_tensor, + cache_position, + past_key_values, + output_attentions, + ): + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + batch_size, sequence_length = input_tensor.shape[0], input_tensor.shape[1] + token_type_ids = self._current_token_type_ids + + causal_mask = self._create_custom_4d_mask( + sequence_length=sequence_length, + dtype=dtype, + device=device, + batch_size=batch_size, + token_type_ids=token_type_ids, + ) + + if attention_mask is not None and attention_mask.dim() == 2: + padding_mask = attention_mask[:, None, None, :].to(dtype=dtype) + padding_mask = (1.0 - padding_mask) * min_dtype + causal_mask = causal_mask + padding_mask + + return causal_mask + + def _create_custom_4d_mask( + self, + sequence_length, + dtype, + device, + batch_size, + token_type_ids, + ): + min_dtype = torch.finfo(dtype).min + masks = [] + for b in range(batch_size): + mask = torch.full((sequence_length, sequence_length), fill_value=min_dtype, dtype=dtype, device=device) + type_ids = token_type_ids[b] + image_positions = (type_ids == 0).nonzero(as_tuple=True)[0] + text_positions = (type_ids == 1).nonzero(as_tuple=True)[0] + + if len(image_positions) > 0: + mask[image_positions[:, None], image_positions] = 0.0 + + for i, text_pos in enumerate(text_positions): + if len(image_positions) > 0: + mask[text_pos, image_positions] = 0.0 + mask[text_pos, text_positions[: i + 1]] = 0.0 + + masks.append(mask) + + mask = torch.stack(masks, dim=0).unsqueeze(1) + return mask + + return CustomQwen2ModelInner(config) + + def forward(self, inputs_embeds, token_type_ids, attention_mask=None, **kwargs): + return self.model(inputs_embeds=inputs_embeds, token_type_ids=token_type_ids, attention_mask=attention_mask, **kwargs) + + +class PTQwen2Decoder2Encoder(nn.Module): + + def __init__( + self, + decoder_layer: int, + hidden_dimension: int, + num_attention_heads: int, + num_key_value_heads: int, + intermediate_size: int, + ): + super().__init__() + self.model = PTCustomQwen2Decoder( + decoder_layer=decoder_layer, + hidden_dimension=hidden_dimension, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + intermediate_size=intermediate_size, + attn_implementation="sdpa", + ) + self.query_768 = nn.Embedding(144, hidden_dimension) + self.query_1024 = nn.Embedding(256, hidden_dimension) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.flatten(1, 2) + bs, n_query, _ = x.shape + if n_query == 144: + param_img = self.query_768.weight + elif n_query == 256: + param_img = self.query_1024.weight + + batch_query_imgs = param_img.unsqueeze(0).expand(bs, -1, -1) + x_combined = torch.cat([x, batch_query_imgs], dim=1) + + token_type_ids = torch.cat( + [ + torch.zeros(bs, n_query, dtype=torch.long), + torch.ones(bs, n_query, dtype=torch.long), + ], + dim=1, + ) + + y = self.model(x_combined, token_type_ids)[0] + y = y[:, n_query:, :] + return y + + +# ============================================================================== +# Test Cases +# ============================================================================== + + +class DeepseekOCRLayersTest(unittest.TestCase): + + def setUp(self): + super().setUp() + self.batch_size = 1 + self.decoder_seq_len = 5 + self.decoder_hidden_size = 32 + self.decoder_num_heads = 4 + self.decoder_num_kv_heads = 2 + self.decoder_head_dim = self.decoder_hidden_size // self.decoder_num_heads + self.decoder_intermediate_size = 64 + + self.config = pyconfig.initialize( + [ + sys.argv[0], + "src/maxtext/configs/base.yml", + "run_name=test", + "skip_jax_distributed_system=True", + "attention=dot_product", + "dtype=float32", + "weight_dtype=float32", + "attention_bias=True", + "matmul_precision=highest", + "rope_max_timescale=1000000.0", + "normalization_layer_epsilon=1e-6", + "base_emb_dim=1280", + "base_num_query_heads=10", + "base_num_kv_heads=10", + "base_mlp_dim=6848", + "base_num_decoder_layers=12", + "vision_connector_emb_dim=896", + "vision_connector_num_query_heads=14", + "vision_connector_num_kv_heads=2", + "vision_connector_mlp_dim=4864", + "vision_connector_num_layers=24", + "decoder_block=deepseek", + "use_mla=false", + "attention_type=global", + "dropout_rate=0.0", + "enable_dropout=false", + "base_emb_dim=32", + "base_num_query_heads=4", + "base_num_kv_heads=2", + "base_mlp_dim=64", + "base_moe_mlp_dim=16", + "base_num_decoder_layers=2", + "first_num_dense_layers=1", + "head_dim=8", + "mlp_activations=['silu', 'linear']", + "rope_interleave=false", + "max_target_length=5", + "max_prefill_predict_length=5", + "per_device_batch_size=1", + "global_batch_size_to_load=1", + "global_batch_size_to_train_on=1", + "num_experts=4", + "num_experts_per_tok=2", + "shared_experts=1", + "routed_score_func=softmax", + "routed_scaling_factor=1.0", + "routed_bias=false", + "sparse_matmul=false", + "scan_layers=false", + "fused_mlp=false", + ] + ) + self.mesh = Mesh(np.array(jax.devices()[:1]), ("data",)) + + def _decoder_inputs(self): + torch.manual_seed(123) + x_pt = torch.randn(self.batch_size, self.decoder_seq_len, self.decoder_hidden_size) + x_jax = to_jax(x_pt) + positions = jnp.broadcast_to(jnp.arange(self.decoder_seq_len, dtype=jnp.int32), (self.batch_size, self.decoder_seq_len)) + segment_ids = jnp.ones((self.batch_size, self.decoder_seq_len), dtype=jnp.int32) + return x_pt, x_jax, positions, segment_ids + + def _jax_dense_decoder_layer(self): + return deepseek.DeepSeekDenseLayer( + config=self.config, + model_mode=MODEL_MODE_TRAIN, + mesh=self.mesh, + rngs=nnx.Rngs(0), + quant=None, + layer_idx=0, + ) + + def test_processor_expands_actual_visual_token_count(self): + pixel_mask = np.zeros( + (1, DEEPSEEK_OCR_GLOBAL_TOKENS + 6 * DEEPSEEK_OCR_CROP_TOKENS + DEEPSEEK_OCR_SEPARATOR_TOKENS), + dtype=np.bool_, + ) + pixel_mask[:, : 2 * DEEPSEEK_OCR_CROP_TOKENS] = True + pixel_mask[:, 6 * DEEPSEEK_OCR_CROP_TOKENS :] = True + processor_output = DeepseekOCR2PreprocessorOutput( + pixel_values=np.zeros((1, 7, 1024, 1024, 3), dtype=np.float32), + pixel_mask=pixel_mask, + num_images=1, + ) + + tokens = np.array([11, DEEPSEEK_OCR_IMAGE_TOKEN_ID, 12], dtype=np.int32) + expanded = add_extra_tokens_for_images_deepseek_ocr(tokens, processor_output) + + expected_visual_tokens = DEEPSEEK_OCR_GLOBAL_TOKENS + 2 * DEEPSEEK_OCR_CROP_TOKENS + DEEPSEEK_OCR_SEPARATOR_TOKENS + self.assertEqual(int((expanded == DEEPSEEK_OCR_IMAGE_TOKEN_ID).sum()), expected_visual_tokens) + self.assertEqual(expanded.shape[0], tokens.shape[0] + expected_visual_tokens - 1) + self.assertEqual(get_image_offsets_deepseek_ocr(processor_output), expected_visual_tokens - 1) + + def test_merge_mm_embeddings_accepts_deepseek_token_level_mask(self): + text_embeddings = jnp.zeros((1, 5, 2), dtype=jnp.float32) + multimodal_embeddings = jnp.arange(10, dtype=jnp.float32).reshape(1, 1, 5, 2) + text_mask = jnp.array([[False, True, True, True, False]]) + token_mask = jnp.array([[True, True, True, False, False]]) + + merged = mm_utils.merge_mm_embeddings(text_embeddings, multimodal_embeddings, text_mask, token_mask) + + np.testing.assert_array_equal(np.asarray(merged[0, 1:4]), np.asarray(multimodal_embeddings[0, 0, :3])) + np.testing.assert_array_equal(np.asarray(merged[0, 4]), np.zeros((2,), dtype=np.float32)) + + def test_deepseek_ocr_text_config_stays_non_mla(self): + self.assertEqual(self.config.decoder_block, DecoderBlockType.DEEPSEEK) + self.assertFalse(self.config.use_mla) + self.assertEqual(self.config.attention_type, "global") + self.assertEqual(self.config.num_query_heads, self.decoder_num_heads) + self.assertEqual(self.config.num_kv_heads, self.decoder_num_kv_heads) + self.assertEqual(self.config.head_dim, self.decoder_head_dim) + + def test_dense_decoder_layer_forward_shape_and_finite(self): + _, x_jax, positions, segment_ids = self._decoder_inputs() + layer = self._jax_dense_decoder_layer() + + out, kv_cache = layer( + x_jax, + decoder_segment_ids=segment_ids, + decoder_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + self.assertIsNone(kv_cache) + self.assertEqual(out.shape, (self.batch_size, self.decoder_seq_len, self.decoder_hidden_size)) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + + def test_dense_decoder_layer_is_causal(self): + _, x_jax, positions, segment_ids = self._decoder_inputs() + layer = self._jax_dense_decoder_layer() + + out_a, _ = layer( + x_jax, + decoder_segment_ids=segment_ids, + decoder_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + x_changed = x_jax.at[:, -1, :].add(25.0) + out_b, _ = layer( + x_changed, + decoder_segment_ids=segment_ids, + decoder_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + np.testing.assert_allclose(np.asarray(out_a[:, :-1, :]), np.asarray(out_b[:, :-1, :]), rtol=1e-5, atol=1e-5) + self.assertGreater(float(jnp.max(jnp.abs(out_a[:, -1, :] - out_b[:, -1, :]))), 1e-3) + + def test_moe_decoder_layer_forward_shape_and_finite(self): + _, x_jax, positions, segment_ids = self._decoder_inputs() + layer = deepseek.DeepSeekMoELayer( + config=self.config, + model_mode=MODEL_MODE_TRAIN, + mesh=self.mesh, + rngs=nnx.Rngs(1), + quant=None, + layer_idx=1, + ) + + out, kv_cache = layer( + x_jax, + decoder_segment_ids=segment_ids, + decoder_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + self.assertIsNone(kv_cache) + self.assertEqual(out.shape, (self.batch_size, self.decoder_seq_len, self.decoder_hidden_size)) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + + def test_one_layer_dense_forward_matches_hf_random_weights(self): + torch.manual_seed(7) + hf_layer, hf_rope = make_llama_decoder_layer( + hidden_size=self.decoder_hidden_size, + num_heads=self.decoder_num_heads, + num_kv_heads=self.decoder_num_kv_heads, + intermediate_size=self.decoder_intermediate_size, + seq_len=self.decoder_seq_len, + ) + jax_layer = self._jax_dense_decoder_layer() + copy_llama_decoder_layer_weights(hf_layer, jax_layer) + + x_pt, x_jax, positions, segment_ids = self._decoder_inputs() + position_ids_pt = torch.zeros((self.batch_size, self.decoder_seq_len), dtype=torch.long) + positions = jnp.zeros_like(positions) + causal_mask = torch.full( + (self.batch_size, 1, self.decoder_seq_len, self.decoder_seq_len), + torch.finfo(torch.float32).min, + dtype=torch.float32, + ) + causal_mask = torch.triu(causal_mask, diagonal=1) + + with torch.no_grad(): + hf_out = hf_layer( + x_pt, + attention_mask=causal_mask, + position_ids=position_ids_pt, + position_embeddings=hf_rope(x_pt, position_ids_pt), + use_cache=False, + ) + + jax_out, _ = jax_layer( + x_jax, + decoder_segment_ids=segment_ids, + decoder_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + if isinstance(hf_out, tuple): + hf_out = hf_out[0] + np.testing.assert_allclose(np.asarray(to_jax(hf_out)), np.asarray(jax_out), rtol=1e-4, atol=1e-4) + + def test_sam_attention(self): + dim = 768 + num_heads = 12 + input_size = (64, 64) + batch_size = 2 + + # PyTorch + pt_layer = PTAttention(dim=dim, num_heads=num_heads, qkv_bias=True, use_rel_pos=True, input_size=input_size) + pt_layer.eval() + + # JAX + rngs = nnx.Rngs(0) + jax_layer = deepseek_ocr.SAMAttention( + dim=dim, + num_heads=num_heads, + qkv_bias=True, + use_rel_pos=True, + input_size=input_size, + rngs=rngs, + ) + + # Copy weights + jax_layer.qkv.kernel.value = to_jax(pt_layer.qkv.weight.T) + jax_layer.qkv.bias.value = to_jax(pt_layer.qkv.bias) + jax_layer.proj.kernel.value = to_jax(pt_layer.proj.weight.T) + jax_layer.proj.bias.value = to_jax(pt_layer.proj.bias) + jax_layer.rel_pos_h.value = to_jax(pt_layer.rel_pos_h) + jax_layer.rel_pos_w.value = to_jax(pt_layer.rel_pos_w) + + # Input + x_pt = torch.randn(batch_size, input_size[0], input_size[1], dim) + x_jax = to_jax(x_pt) + + # Forward + with torch.no_grad(): + out_pt = pt_layer(x_pt) + out_jax = jax_layer(x_jax) + + # Compare + np.testing.assert_allclose(to_jax(out_pt), out_jax, rtol=3e-4, atol=3e-4) + + def test_sam_block(self): + dim = 768 + num_heads = 12 + input_size = (64, 64) + batch_size = 2 + window_size = 14 + + # PyTorch + pt_layer = PTBlock( + dim=dim, num_heads=num_heads, mlp_ratio=4.0, qkv_bias=True, window_size=window_size, input_size=input_size + ) + pt_layer.eval() + + # JAX + rngs = nnx.Rngs(0) + jax_layer = deepseek_ocr.SAMBlock( + dim=dim, + num_heads=num_heads, + mlp_ratio=4.0, + qkv_bias=True, + use_rel_pos=True, + window_size=window_size, + input_size=input_size, + rngs=rngs, + ) + + # Copy weights + jax_layer.norm1.scale.value = to_jax(pt_layer.norm1.weight) + jax_layer.norm1.bias.value = to_jax(pt_layer.norm1.bias) + jax_layer.norm2.scale.value = to_jax(pt_layer.norm2.weight) + jax_layer.norm2.bias.value = to_jax(pt_layer.norm2.bias) + + jax_layer.attn.qkv.kernel.value = to_jax(pt_layer.attn.qkv.weight.T) + jax_layer.attn.qkv.bias.value = to_jax(pt_layer.attn.qkv.bias) + jax_layer.attn.proj.kernel.value = to_jax(pt_layer.attn.proj.weight.T) + jax_layer.attn.proj.bias.value = to_jax(pt_layer.attn.proj.bias) + jax_layer.attn.rel_pos_h.value = to_jax(pt_layer.attn.rel_pos_h) + jax_layer.attn.rel_pos_w.value = to_jax(pt_layer.attn.rel_pos_w) + + jax_layer.lin1.kernel.value = to_jax(pt_layer.mlp.lin1.weight.T) + jax_layer.lin1.bias.value = to_jax(pt_layer.mlp.lin1.bias) + jax_layer.lin2.kernel.value = to_jax(pt_layer.mlp.lin2.weight.T) + jax_layer.lin2.bias.value = to_jax(pt_layer.mlp.lin2.bias) + + # Input + x_pt = torch.randn(batch_size, input_size[0], input_size[1], dim) + x_jax = to_jax(x_pt) + + # Forward + with torch.no_grad(): + out_pt = pt_layer(x_pt) + out_jax = jax_layer(x_jax) + + # Compare + np.testing.assert_allclose(to_jax(out_pt), out_jax, rtol=4e-3, atol=4e-3) + + def test_sam_vit_b(self): + batch_size = 2 + img_size = 1024 + + # PyTorch + pt_layer = PTImageEncoderViT(img_size=img_size) + pt_layer.eval() + + # JAX + rngs = nnx.Rngs(0) + jax_layer = deepseek_ocr.SAMViTB(config=self.config, mesh=self.mesh, rngs=rngs) + + # Copy weights + jax_layer.patch_embed.kernel.value = to_jax(pt_layer.patch_embed.proj.weight.permute(2, 3, 1, 0)) + jax_layer.patch_embed.bias.value = to_jax(pt_layer.patch_embed.proj.bias) + jax_layer.pos_embed.value = to_jax(pt_layer.pos_embed) + + for i in range(12): + pt_blk = pt_layer.blocks[i] + jax_blk = getattr(jax_layer, f"block_{i}") + + jax_blk.norm1.scale.value = to_jax(pt_blk.norm1.weight) + jax_blk.norm1.bias.value = to_jax(pt_blk.norm1.bias) + jax_blk.norm2.scale.value = to_jax(pt_blk.norm2.weight) + jax_blk.norm2.bias.value = to_jax(pt_blk.norm2.bias) + + jax_blk.attn.qkv.kernel.value = to_jax(pt_blk.attn.qkv.weight.T) + jax_blk.attn.qkv.bias.value = to_jax(pt_blk.attn.qkv.bias) + jax_blk.attn.proj.kernel.value = to_jax(pt_blk.attn.proj.weight.T) + jax_blk.attn.proj.bias.value = to_jax(pt_blk.attn.proj.bias) + jax_blk.attn.rel_pos_h.value = to_jax(pt_blk.attn.rel_pos_h) + jax_blk.attn.rel_pos_w.value = to_jax(pt_blk.attn.rel_pos_w) + + jax_blk.lin1.kernel.value = to_jax(pt_blk.mlp.lin1.weight.T) + jax_blk.lin1.bias.value = to_jax(pt_blk.mlp.lin1.bias) + jax_blk.lin2.kernel.value = to_jax(pt_blk.mlp.lin2.weight.T) + jax_blk.lin2.bias.value = to_jax(pt_blk.mlp.lin2.bias) + + jax_layer.neck_conv1.kernel.value = to_jax(pt_layer.neck[0].weight.permute(2, 3, 1, 0)) + jax_layer.neck_ln1.scale.value = to_jax(pt_layer.neck[1].weight) + jax_layer.neck_ln1.bias.value = to_jax(pt_layer.neck[1].bias) + jax_layer.neck_conv2.kernel.value = to_jax(pt_layer.neck[2].weight.permute(2, 3, 1, 0)) + jax_layer.neck_ln2.scale.value = to_jax(pt_layer.neck[3].weight) + jax_layer.neck_ln2.bias.value = to_jax(pt_layer.neck[3].bias) + + jax_layer.net_2.kernel.value = to_jax(pt_layer.net_2.weight.permute(2, 3, 1, 0)) + jax_layer.net_3.kernel.value = to_jax(pt_layer.net_3.weight.permute(2, 3, 1, 0)) + + # Input (PyTorch expects NCHW, JAX expects NHWC) + x_pt = torch.randn(batch_size, 3, img_size, img_size) + x_jax = to_jax(x_pt.permute(0, 2, 3, 1)) + + # Forward + with torch.no_grad(): + out_pt = pt_layer(x_pt) + out_jax = jax_layer(x_jax) + + # Compare + out_pt_nhwc = out_pt.permute(0, 2, 3, 1) + np.testing.assert_allclose(to_jax(out_pt_nhwc), out_jax, rtol=1.2e-2, atol=1.2e-2) + + def test_qwen2_encoder_layer(self): + # Config + pt_config = PTQwen2Config( + hidden_size=896, + num_attention_heads=14, + num_key_value_heads=2, + intermediate_size=4864, + rms_norm_eps=1e-6, + ) + + # PyTorch + pt_layer = PTQwen2DecoderLayer(pt_config, layer_idx=0) + pt_layer.eval() + + # JAX + from maxtext.configs.pyconfig import HyperParameters + + pydantic_config = self.config._pydantic_config + update_dict = { + "decoder_block": DecoderBlockType.QWEN2, + "emb_dim": 896, + "num_query_heads": 14, + "num_kv_heads": 2, + "mlp_dim": 4864, + "head_dim": 896 // 14, + "max_target_length": 288, + "max_prefill_predict_length": 288, + } + if hasattr(pydantic_config, "model_copy"): + new_pydantic_config = pydantic_config.model_copy(update=update_dict) + else: + new_pydantic_config = pydantic_config.copy(update=update_dict) + connector_config = HyperParameters(new_pydantic_config) + rngs = nnx.Rngs(0) + jax_layer = deepseek_ocr.Qwen2EncoderLayer( + config=connector_config, + mesh=self.mesh, + quant=None, + rngs=rngs, + ) + + # Copy weights + copy_rmsnorm_weights(pt_layer.input_layernorm, jax_layer.pre_self_attention_layer_norm) + copy_rmsnorm_weights(pt_layer.post_attention_layernorm, jax_layer.post_self_attention_layer_norm) + copy_qwen2_attention_weights(pt_layer.self_attn, jax_layer.self_attention) + copy_linear_weights(pt_layer.mlp.gate_proj, jax_layer.mlp.wi_0) + copy_linear_weights(pt_layer.mlp.up_proj, jax_layer.mlp.wi_1) + copy_linear_weights(pt_layer.mlp.down_proj, jax_layer.mlp.wo) + + # Input + batch_size = 1 + seq_len = 288 + x_pt = torch.randn(batch_size, seq_len, 896) + x_jax = to_jax(x_pt) + + # Mask & Positions + token_type_ids = torch.cat( + [ + torch.zeros(batch_size, 144, dtype=torch.long), + torch.ones(batch_size, 144, dtype=torch.long), + ], + dim=1, + ) + + decoder = PTCustomQwen2Decoder( + decoder_layer=1, + hidden_dimension=896, + num_attention_heads=14, + num_key_value_heads=2, + intermediate_size=4864, + ) + decoder.model.layers[0].load_state_dict(pt_layer.state_dict()) + decoder.eval() + + with torch.no_grad(): + decoder.model._current_token_type_ids = token_type_ids + position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) + mask_pt = decoder.model._update_causal_mask(None, x_pt, None, None, None) + position_embeddings = decoder.model.rotary_emb(x_pt, position_ids) + + # Run PyTorch Layer directly + out_pt = pt_layer( + hidden_states=x_pt, + attention_mask=mask_pt, + position_ids=position_ids, + position_embeddings=position_embeddings, + )[0] + + # JAX Run + mask_image = jnp.ones((batch_size, 144), dtype=jnp.bool_) + mask_query = jnp.zeros((batch_size, 144), dtype=jnp.bool_) + bidirectional_mask = jnp.concatenate([mask_image, mask_query], axis=1) + decoder_positions = jnp.arange(seq_len)[None, :] + + out_jax = jax_layer( + x_jax, + bidirectional_mask=bidirectional_mask, + decoder_positions=decoder_positions, + deterministic=True, + ) + + if out_pt.ndim != out_jax.ndim: + out_pt_compared = out_pt.unsqueeze(0) + else: + out_pt_compared = out_pt + np.testing.assert_allclose(to_jax(out_pt_compared), out_jax, rtol=1.5e-3, atol=1.5e-3) + + def test_qwen2_decoder2encoder(self): + batch_size = 1 + + # PyTorch + pt_model = PTQwen2Decoder2Encoder( + decoder_layer=2, # Use 2 layers for faster test + hidden_dimension=896, + num_attention_heads=14, + num_key_value_heads=2, + intermediate_size=4864, + ) + pt_model.eval() + + # JAX + from maxtext.configs.pyconfig import HyperParameters + + pydantic_config = self.config._pydantic_config + update_dict = { + "decoder_block": DecoderBlockType.QWEN2, + "vision_connector_num_layers": 2, + "max_target_length": 512, + "max_prefill_predict_length": 512, + } + if hasattr(pydantic_config, "model_copy"): + new_pydantic_config = pydantic_config.model_copy(update=update_dict) + else: + new_pydantic_config = pydantic_config.copy(update=update_dict) + test_config = HyperParameters(new_pydantic_config) + rngs = nnx.Rngs(0) + jax_model = deepseek_ocr.Qwen2Decoder2Encoder( + config=test_config, + mesh=self.mesh, + quant=None, + rngs=rngs, + ) + + # Copy weights + jax_model.query_768.embedding.value = to_jax(pt_model.query_768.weight) + jax_model.query_1024.embedding.value = to_jax(pt_model.query_1024.weight) + copy_rmsnorm_weights(pt_model.model.model.norm, jax_model.norm) + + for i in range(2): + pt_layer = pt_model.model.model.layers[i] + jax_layer = getattr(jax_model, f"layer_{i}") + + copy_rmsnorm_weights(pt_layer.input_layernorm, jax_layer.pre_self_attention_layer_norm) + copy_rmsnorm_weights(pt_layer.post_attention_layernorm, jax_layer.post_self_attention_layer_norm) + copy_qwen2_attention_weights(pt_layer.self_attn, jax_layer.self_attention) + copy_linear_weights(pt_layer.mlp.gate_proj, jax_layer.mlp.wi_0) + copy_linear_weights(pt_layer.mlp.up_proj, jax_layer.mlp.wi_1) + copy_linear_weights(pt_layer.mlp.down_proj, jax_layer.mlp.wo) + + for H, W in [(12, 12)]: + x_pt = torch.randn(batch_size, H, W, 896) + x_jax = to_jax(x_pt) + + with torch.no_grad(): + out_pt = pt_model(x_pt) + out_jax = jax_model(x_jax) + np.testing.assert_allclose(to_jax(out_pt), out_jax, rtol=1.5e-3, atol=1.5e-3) + + def test_image_preprocessing(self): + # Create a dummy image with a specific size that triggers cropping + # e.g., 1600 x 1200 + np.random.seed(42) + random_array = np.random.randint(0, 256, (1200, 1600, 3), dtype=np.uint8) + image = Image.fromarray(random_array) + + # HF Reference + hf_res = _hf_reference_inputs(image) + + # MaxText + from maxtext.multimodal.processor_deepseek_ocr import preprocess_mm_data_deepseek_ocr + mt_res = preprocess_mm_data_deepseek_ocr(image) + + # Compare aspect ratios + np.testing.assert_array_equal(hf_res["crop_ratio"], mt_res.aspect_ratios[0]) + + # Compare global view + mt_global = mt_res.pixel_values[0, 0] + np.testing.assert_allclose(hf_res["images_ori"], mt_global, atol=1e-5) + + # Compare crops + num_crops = hf_res["images_crop"].shape[0] + for i in range(num_crops): + mt_crop = mt_res.pixel_values[0, i + 1, :768, :768, :] + np.testing.assert_allclose(hf_res["images_crop"][i], mt_crop, atol=1e-5) + + def test_deepseek_ocr_vision_encoder(self): + # JAX Config override for fast test + from maxtext.configs.pyconfig import HyperParameters + pydantic_config = self.config._pydantic_config + update_dict = { + "decoder_block": DecoderBlockType.QWEN2, + "vision_connector_num_layers": 2, + "max_target_length": 512, + "max_prefill_predict_length": 512, + } + if hasattr(pydantic_config, "model_copy"): + new_pydantic_config = pydantic_config.model_copy(update=update_dict) + else: + new_pydantic_config = pydantic_config.copy(update=update_dict) + test_config = HyperParameters(new_pydantic_config) + + rngs = nnx.Rngs(0) + jax_model = deepseek_ocr.DeepseekOCR2VisionEncoder(config=test_config, mesh=self.mesh, rngs=rngs) + + # PT Models + pt_sam = PTImageEncoderViT(img_size=1024) + pt_qwen = PTQwen2Decoder2Encoder( + decoder_layer=2, + hidden_dimension=896, + num_attention_heads=14, + num_key_value_heads=2, + intermediate_size=4864, + ) + pt_sam.eval() + pt_qwen.eval() + + # Copy weights + # SAM + jax_model.sam_model.patch_embed.kernel.value = to_jax(pt_sam.patch_embed.proj.weight.permute(2, 3, 1, 0)) + jax_model.sam_model.patch_embed.bias.value = to_jax(pt_sam.patch_embed.proj.bias) + jax_model.sam_model.pos_embed.value = to_jax(pt_sam.pos_embed) + for i in range(12): + pt_blk = pt_sam.blocks[i] + jax_blk = getattr(jax_model.sam_model, f"block_{i}") + jax_blk.norm1.scale.value = to_jax(pt_blk.norm1.weight) + jax_blk.norm1.bias.value = to_jax(pt_blk.norm1.bias) + jax_blk.norm2.scale.value = to_jax(pt_blk.norm2.weight) + jax_blk.norm2.bias.value = to_jax(pt_blk.norm2.bias) + jax_blk.attn.qkv.kernel.value = to_jax(pt_blk.attn.qkv.weight.T) + jax_blk.attn.qkv.bias.value = to_jax(pt_blk.attn.qkv.bias) + jax_blk.attn.proj.kernel.value = to_jax(pt_blk.attn.proj.weight.T) + jax_blk.attn.proj.bias.value = to_jax(pt_blk.attn.proj.bias) + jax_blk.attn.rel_pos_h.value = to_jax(pt_blk.attn.rel_pos_h) + jax_blk.attn.rel_pos_w.value = to_jax(pt_blk.attn.rel_pos_w) + jax_blk.lin1.kernel.value = to_jax(pt_blk.mlp.lin1.weight.T) + jax_blk.lin1.bias.value = to_jax(pt_blk.mlp.lin1.bias) + jax_blk.lin2.kernel.value = to_jax(pt_blk.mlp.lin2.weight.T) + jax_blk.lin2.bias.value = to_jax(pt_blk.mlp.lin2.bias) + + jax_model.sam_model.neck_conv1.kernel.value = to_jax(pt_sam.neck[0].weight.permute(2, 3, 1, 0)) + jax_model.sam_model.neck_ln1.scale.value = to_jax(pt_sam.neck[1].weight) + jax_model.sam_model.neck_ln1.bias.value = to_jax(pt_sam.neck[1].bias) + jax_model.sam_model.neck_conv2.kernel.value = to_jax(pt_sam.neck[2].weight.permute(2, 3, 1, 0)) + jax_model.sam_model.neck_ln2.scale.value = to_jax(pt_sam.neck[3].weight) + jax_model.sam_model.neck_ln2.bias.value = to_jax(pt_sam.neck[3].bias) + jax_model.sam_model.net_2.kernel.value = to_jax(pt_sam.net_2.weight.permute(2, 3, 1, 0)) + jax_model.sam_model.net_3.kernel.value = to_jax(pt_sam.net_3.weight.permute(2, 3, 1, 0)) + + # Qwen2 + jax_model.qwen2_model.query_768.embedding.value = to_jax(pt_qwen.query_768.weight) + jax_model.qwen2_model.query_1024.embedding.value = to_jax(pt_qwen.query_1024.weight) + copy_rmsnorm_weights(pt_qwen.model.model.norm, jax_model.qwen2_model.norm) + for i in range(2): + pt_layer = pt_qwen.model.model.layers[i] + jax_layer = getattr(jax_model.qwen2_model, f"layer_{i}") + copy_rmsnorm_weights(pt_layer.input_layernorm, jax_layer.pre_self_attention_layer_norm) + copy_rmsnorm_weights(pt_layer.post_attention_layernorm, jax_layer.post_self_attention_layer_norm) + copy_qwen2_attention_weights(pt_layer.self_attn, jax_layer.self_attention) + copy_linear_weights(pt_layer.mlp.gate_proj, jax_layer.mlp.wi_0) + copy_linear_weights(pt_layer.mlp.up_proj, jax_layer.mlp.wi_1) + copy_linear_weights(pt_layer.mlp.down_proj, jax_layer.mlp.wo) + + # Input: 1 image, 6 crops + 1 global (total 7 views) + # JAX expects [B, N, 7, H, W, C] + x_jax = jax.random.normal(jax.random.PRNGKey(42), (1, 1, 7, 1024, 1024, 3)) + x_pt_global = torch.from_numpy(np.array(x_jax[0, 0, 0].transpose(2, 0, 1))).unsqueeze(0) + x_pt_crops = torch.from_numpy(np.array(x_jax[0, 0, 1:, :768, :768, :].transpose(0, 3, 1, 2))) + + # Run PT + with torch.no_grad(): + sam_global = pt_sam(x_pt_global) + qwen_global = pt_qwen(sam_global.permute(0, 2, 3, 1)) + + sam_crops = pt_sam(x_pt_crops) + qwen_crops = pt_qwen(sam_crops.permute(0, 2, 3, 1)) + + pt_out = torch.cat([qwen_crops.reshape(864, 896), qwen_global.reshape(256, 896)], dim=0).unsqueeze(0).unsqueeze(0) + + # Run JAX + jax_out = jax_model(x_jax, deterministic=True) + + # Compare (uses tolerance of 1e-2 to account for cumulative numerical error across SAM and Qwen2) + np.testing.assert_allclose(to_jax(pt_out), jax_out, rtol=1e-2, atol=1e-2) + + def test_text_image_embedding_merge_realistic(self): + text_embeddings = jnp.arange(1024 * 128, dtype=jnp.float32).reshape(1, 1024, 128) + multimodal_embeddings = (jnp.arange(1121 * 128, dtype=jnp.float32) + 1000000.0).reshape(1, 1, 1121, 128) + + mask = np.zeros((1, 1024), dtype=np.bool_) + mask[0, 100:933] = True + mask = jnp.array(mask) + + pixel_mask = np.zeros((1, 1121), dtype=np.bool_) + pixel_mask[0, :576] = True + pixel_mask[0, 864:1121] = True + pixel_mask = jnp.array(pixel_mask) + + merged = mm_utils.merge_mm_embeddings( + text_embeddings=text_embeddings, + multimodal_embeddings=multimodal_embeddings, + mask=mask, + token_masks=pixel_mask, + ) + + active_mm = multimodal_embeddings[0, 0][pixel_mask[0]] + np.testing.assert_allclose(merged[0, 100:933], active_mm, atol=1e-5) + np.testing.assert_allclose(merged[0, :100], text_embeddings[0, :100], atol=1e-5) + np.testing.assert_allclose(merged[0, 933:], text_embeddings[0, 933:], atol=1e-5) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/param_mapping_test.py b/tests/unit/param_mapping_test.py index 8a18c64b9d..435e8f95f2 100644 --- a/tests/unit/param_mapping_test.py +++ b/tests/unit/param_mapping_test.py @@ -23,6 +23,13 @@ class ParamMappingTest(unittest.TestCase): + def _apply_hooks(self, weight, target_shape, hooks): + if not isinstance(hooks, list): + hooks = [hooks] + for hook in hooks: + weight = hook(weight, target_shape) + return weight + def test_gemma3_mapping_unscanned(self): config = { "text_config": {"num_hidden_layers": 2, "hidden_size": 256}, @@ -125,6 +132,66 @@ def test_deepseek_mapping_scanned(self): mapping = param_mapping.DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True) self.assertIn("params-decoder-dense_layers-self_attention-query-kernel", mapping) + def test_deepseek_non_mla_query_hook_scales_dense_layer(self): + config = { + "num_hidden_layers": 2, + "first_k_dense_replace": 1, + "n_routed_experts": 2, + "use_mla": False, + "head_dim": 3, + } + maxtext_config = mock.Mock() + hooks = param_mapping.DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN( + config, maxtext_config, scan_layers=False, saving_to_hf=False + ) + hook = hooks["params-decoder-dense_layer_0-self_attention-query-kernel"] + + hf_q = np.arange(36, dtype=np.float32).reshape(6, 6) + target_shape = (6, 2, 3) + output = self._apply_hooks(hf_q, target_shape, hook) + expected = hf_q.T.reshape(target_shape) / np.sqrt(config["head_dim"]) + np.testing.assert_allclose(output, expected) + + def test_deepseek_non_mla_query_hook_scales_moe_layer(self): + config = { + "num_hidden_layers": 2, + "first_k_dense_replace": 1, + "n_routed_experts": 2, + "use_mla": False, + "head_dim": 3, + } + maxtext_config = mock.Mock() + hooks = param_mapping.DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN( + config, maxtext_config, scan_layers=False, saving_to_hf=False + ) + hook = hooks["params-decoder-layers_1-self_attention-query-kernel"] + + hf_q = np.arange(36, dtype=np.float32).reshape(6, 6) + target_shape = (6, 2, 3) + output = self._apply_hooks(hf_q, target_shape, hook) + expected = hf_q.T.reshape(target_shape) / np.sqrt(config["head_dim"]) + np.testing.assert_allclose(output, expected) + + def test_deepseek_non_mla_key_hook_only_reshapes(self): + config = { + "num_hidden_layers": 2, + "first_k_dense_replace": 1, + "n_routed_experts": 2, + "use_mla": False, + "head_dim": 3, + } + maxtext_config = mock.Mock() + hooks = param_mapping.DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN( + config, maxtext_config, scan_layers=False, saving_to_hf=False + ) + hook = hooks["params-decoder-dense_layer_0-self_attention-key-kernel"] + + hf_k = np.arange(36, dtype=np.float32).reshape(6, 6) + target_shape = (6, 2, 3) + output = self._apply_hooks(hf_k, target_shape, hook) + expected = hf_k.T.reshape(target_shape) + np.testing.assert_allclose(output, expected) + def test_gpt_oss_mapping(self): config = { "num_hidden_layers": 2,