[WIP] deepseek_ocr_2 exp

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

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,

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,13 +1713,26 @@ 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"]
 
   mapping = {
       "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,

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}")