Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,21 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/4f87c6488b21c3c03780afba253c8af0 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/f43da177788619ad3cd686b1345608ef filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/ce563dd53d0cc8894641b5996dc5453f filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/156c1d2fa9763c506b9a075a5729b7de filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/8d2385834c5546d282b60ec4c29d4378 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/ca081ceb253837d87f9bc28d6dfe2b0a filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/44d6c5046ef1352f82da17c3daffc66d filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/ae742711b7e4d0ae7a8eaa239393ec58 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/88b7b61e8b4c3f22ec014bd4fcd5cbe5 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/78efb519de44bcebc6eb6fd345ff0af0 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/df6d4f33d5f964745802e5eedfe1425e filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/124c8d55117974805f4a148ebd670601 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/82c6ee1c6166b74a75b93badf224480e filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/97411cb3209e6a07d86332228942a45c filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/e61206338bedb2c7eecd25dc76511b00 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/2f806db10cde1c0f439f288f34ee8b95 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/0c600442de678f199ab5e8d0152c8bbc filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/a4f076e56e9ec995bf74b8f5d19e4101 filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,55 @@
+Gemma Terms of Use
+
+Last modified: April 1, 2024
+
+Thank you for using Gemma! These terms of use ("Terms") are entered into by and between Google LLC ("Google") and you and govern your use of Gemma. This is an agreement between you and Google. If you are accepting these Terms and using Gemma on behalf of another person or entity, you represent and warrant that you have full authority to bind that person or entity to these Terms.
+
+1. THE AGREEMENT
+
+1.1 These Terms apply only to Gemma and any use or access you obtain to Gemma. For the purpose of these Terms, "Gemma" means the Gemma model and software and tools made available by Google, and any software, documentation, or other materials provided by Google relating to Gemma. Gemma may be: (a) a standalone product, or (b) a component of other software products. Gemma is not a Google Service or a Google Product as described in Google’s Terms of Service (https://policies.google.com/terms). These Terms do not apply to any Third Party Products (defined in Section 2.3) or any other Google products or services.
+
+1.2 “You” means you as an individual or entity, or your authorized representative. If you or your organization have a separate, written agreement with Google governing your use of Gemma, that agreement will supersede these Terms.
+
+2. LICENSE AND USE
+
+2.1 Subject to your compliance with these Terms, Google grants you a perpetual, non-exclusive, worldwide, non-transferable, non-sublicensable license to use, reproduce, distribute, perform, and display Gemma, to prepare Derivatives of Gemma, and otherwise use Gemma in your Derivative Models.
+
+2.2 You may publicly distribute or make available to third parties, without restriction or charge, Gemma, Derivative Models, and Derivative Works (collectively, the “Artifacts”), provided you: (a) comply with these Terms; (b) provide a copy of these Terms; (c) provide attribution to Google as the source of Gemma in a manner that is reasonable and customary within the industry or context; (d) include a prominent notice with any Derivative Works stating that you have modified Gemma; and (e) include the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms". 
+
+2.3 You may use Third Party Products with Gemma. "Third Party Products" means third-party software, data, or other resources. Third Party Products may be subject to their own terms and Google is not responsible for Third Party Products.
+
+3. RESTRICTIONS AND PROHIBITED USE POLICY
+
+3.1 You will comply with the Gemma Prohibited Use Policy (provided at ai.google.dev/gemma/terms). Google may update the Gemma Prohibited Use Policy from time to time by publishing a new version at ai.google.dev/gemma/terms.
+
+3.2 You will not use Gemma or Derivatives, or allow or facilitate its use, for any of the following:
+
+(a) High Risk Activities: any use that could lead to death, personal injury, or environmental damage, such as in: (i) emergency services; (ii) military, law enforcement, or surveillance; (iii) management or operation of critical infrastructure; or (iv) safety applications.
+
+(b) Development or Use of Weapons. Prohibited uses include, but are not limited to, research, development, production, testing, maintenance or use of: (i) weapons; (ii) technologies that cause physical injury or death; or (iii) agents intended to cause harm to humans, animals, or property.
+
+(c) Malicious, Abusive, or Unlawful Activities. Prohibited uses include, but are not limited to, use of Gemma or Derivatives to: (i) develop malware, including spyware; (ii) engage in phishing, fraud, or other deceptive activities; (iii) engage in cyber abuse, including denial of service or hacking; (iv) exploit vulnerabilities; (v) generate or disseminate spam; or (vi) facilitate or encourage violations of law.
+
+(d) Automated Decision-Making. Prohibited uses include, but are not limited to, use of Gemma or Derivatives to: (i) make decisions on behalf of decision makers within government; (ii) make decisions that determine an individual’s eligibility or selection within education, employment, housing, insurance, or credit; or (iii) make decisions, recommendations, or rankings that determine a person’s eligibility or access within healthcare or legal services.
+
+(e) Harmful or Scams. Prohibited uses include, but are not limited to, use of Gemma or Derivatives to: (i) facilitate, encourage, or further scams; (ii) generate or disseminate content that incites violence or hatred; (iii) assist in exploitation or abuse of minors; or (iv) promote or facilitate discrimination.
+
+4. TERM AND TERMINATION
+
+4.1 These Terms are effective upon your first use, access, or download of Gemma and will continue until terminated.
+
+4.2 Google may suspend or terminate your use of Gemma at any time if it believes you have breached these Terms or to comply with applicable law. If Google suspends or terminates your use, you must stop using Gemma.
+
+4.3 Upon termination, you must cease all use of Gemma and delete Gemma and all Derivatives.
+
+5. DISCLAIMERS
+
+5.1 GEMMA IS PROVIDED "AS IS" WITHOUT WARRANTIES OF ANY KIND. GOOGLE DISCLAIMS ALL WARRANTIES AND CONDITIONS, EXPRESS OR IMPLIED, INCLUDING ANY WARRANTIES OR CONDITIONS OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, OR NON-INFRINGEMENT.
+
+5.2 YOU ARE RESPONSIBLE FOR YOUR USE OF GEMMA AND DERIVATIVES.
+
+5.3 IN NO EVENT WILL GOOGLE OR ITS AFFILIATES BE LIABLE FOR ANY INDIRECT, INCIDENTAL, SPECIAL, CONSEQUENTIAL, OR EXEMPLARY DAMAGES OR LOST PROFITS ARISING FROM OR IN CONNECTION WITH THESE TERMS OR USE OF GEMMA.
+
+If you have any questions regarding these Terms, you may contact Google at gemma@google.com.
+
+Please read these Terms carefully and keep a copy for your reference.

NOTICE

@@ -0,0 +1 @@
+Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms

README.md

@@ -0,0 +1,52 @@
+---
+license: gemma
+base_model: google/gemma-3-27b-pt
+tags:
+- jax
+- orbax
+- tensorstore
+- gemma
+- conversion
+---
+
+# Gemma 3 27B PT – Orbax/TensorStore (OCDBT)
+
+This repository contains the **Orbax/TensorStore (OCDBT)** checkpoint converted from
+`google/gemma-3-27b-pt` Hugging Face safetensors. The conversion stacks transformer
+layer weights along the depth axis for efficient JAX/Orbax loading.
+
+## What’s included
+
+- Orbax/TensorStore checkpoint files (OCDBT)
+- `convert.py` (the conversion script used)
+- `LICENSE` and `NOTICE` per Gemma Terms of Use
+
+## Conversion details
+
+- Source: `google/gemma-3-27b-pt`
+- Format change only: **no weight changes** beyond layout/format
+- Layer weights are stacked under `layers_stacked/...`
+
+## Loading example (JAX/Orbax)
+
+```python
+from huggingface_hub import snapshot_download
+import orbax.checkpoint as ocp
+
+path = snapshot_download("dffarr/gemma-3-27b-pt-orbax")
+ckpt = ocp.StandardCheckpointer().restore(path)
+
+# Example access
+embed_tokens = ckpt["language_model/model/embed_tokens"]
+q_proj = ckpt["layers_stacked/language_model/model/self_attn/q_proj/weight"]
+```
+
+## License & Use
+
+This repository is provided under the **Gemma Terms of Use**. Please read `LICENSE`
+and comply with the **Gemma Prohibited Use Policy**:
+https://ai.google.dev/gemma/terms
+
+## Conversion script
+
+The conversion script used to generate this checkpoint is included as `convert.py`.

_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1770325910424047211, "commit_timestamp_nsecs": 1770326068343296057, "custom_metadata": {}}

_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('language_model/model/embed_tokens',)": {"key_metadata": [{"key": "language_model/model/embed_tokens", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('language_model/model/norm',)": {"key_metadata": [{"key": "language_model/model/norm", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/input_layernorm/weight',)": {"key_metadata": [{"key": "layers_stacked/input_layernorm/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/language_model/model/self_attn/k_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/language_model/model/self_attn/k_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/language_model/model/self_attn/q_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/language_model/model/self_attn/q_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/language_model/model/self_attn/v_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/language_model/model/self_attn/v_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/layer_norm1/bias',)": {"key_metadata": [{"key": "layers_stacked/layer_norm1/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/layer_norm1/weight',)": {"key_metadata": [{"key": "layers_stacked/layer_norm1/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/layer_norm2/bias',)": {"key_metadata": [{"key": "layers_stacked/layer_norm2/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/layer_norm2/weight',)": {"key_metadata": [{"key": "layers_stacked/layer_norm2/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/mlp/down_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/mlp/down_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/mlp/fc1/bias',)": {"key_metadata": [{"key": "layers_stacked/mlp/fc1/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/mlp/fc1/weight',)": {"key_metadata": [{"key": "layers_stacked/mlp/fc1/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/mlp/fc2/bias',)": {"key_metadata": [{"key": "layers_stacked/mlp/fc2/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/mlp/fc2/weight',)": {"key_metadata": [{"key": "layers_stacked/mlp/fc2/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/mlp/gate_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/mlp/gate_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/mlp/up_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/mlp/up_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/post_attention_layernorm/weight',)": {"key_metadata": [{"key": "layers_stacked/post_attention_layernorm/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/post_feedforward_layernorm/weight',)": {"key_metadata": [{"key": "layers_stacked/post_feedforward_layernorm/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/pre_feedforward_layernorm/weight',)": {"key_metadata": [{"key": "layers_stacked/pre_feedforward_layernorm/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/self_attn/k_norm/weight',)": {"key_metadata": [{"key": "layers_stacked/self_attn/k_norm/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/self_attn/k_proj/bias',)": {"key_metadata": [{"key": "layers_stacked/self_attn/k_proj/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/self_attn/o_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/self_attn/o_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/self_attn/out_proj/bias',)": {"key_metadata": [{"key": "layers_stacked/self_attn/out_proj/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/self_attn/out_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/self_attn/out_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/self_attn/q_norm/weight',)": {"key_metadata": [{"key": "layers_stacked/self_attn/q_norm/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/self_attn/q_proj/bias',)": {"key_metadata": [{"key": "layers_stacked/self_attn/q_proj/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/self_attn/v_proj/bias',)": {"key_metadata": [{"key": "layers_stacked/self_attn/v_proj/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/vision_tower/vision_model/encoder/self_attn/k_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/vision_tower/vision_model/encoder/self_attn/k_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/vision_tower/vision_model/encoder/self_attn/q_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/vision_tower/vision_model/encoder/self_attn/q_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('layers_stacked/vision_tower/vision_model/encoder/self_attn/v_proj/weight',)": {"key_metadata": [{"key": "layers_stacked/vision_tower/vision_model/encoder/self_attn/v_proj/weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('multi_modal_projector/mm_input_projection_weight',)": {"key_metadata": [{"key": "multi_modal_projector/mm_input_projection_weight", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('multi_modal_projector/mm_soft_emb_norm',)": {"key_metadata": [{"key": "multi_modal_projector/mm_soft_emb_norm", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('vision_tower/vision_model/embeddings/patch_embedding',)": {"key_metadata": [{"key": "vision_tower/vision_model/embeddings/patch_embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('vision_tower/vision_model/embeddings/patch_embedding/bias',)": {"key_metadata": [{"key": "vision_tower/vision_model/embeddings/patch_embedding/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('vision_tower/vision_model/embeddings/position_embedding',)": {"key_metadata": [{"key": "vision_tower/vision_model/embeddings/position_embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('vision_tower/vision_model/post_layernorm',)": {"key_metadata": [{"key": "vision_tower/vision_model/post_layernorm", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('vision_tower/vision_model/post_layernorm/bias',)": {"key_metadata": [{"key": "vision_tower/vision_model/post_layernorm/bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

commit_success.txt

@@ -0,0 +1 @@
+Checkpoint commit was successful to gs://gemma-3-27b-pt-orbax-b76114af/gemma-3-27b-pt-orbax

convert.py

@@ -0,0 +1,469 @@
+#!/usr/bin/env python3
+"""Convert Hugging Face safetensors to Orbax (TensorStore) with smart stacking.
+
+This program streams a Hugging Face safetensors snapshot, discovers transformer
+layer parameters, stacks them by suffix across depth, and writes a topology-
+agnostic Orbax checkpoint to GCS. It enforces CPU execution, processes one stack
+at a time to limit peak RAM, and keeps non-layer parameters unstacked with clean
+names for downstream JAX/Orbax training loops.
+"""
+
+from __future__ import annotations
+
+import gc
+import fnmatch
+import json
+import logging
+import os
+import shutil
+import sys
+import tempfile
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Tuple
+
+# Force CPU before importing JAX.
+os.environ["JAX_PLATFORM_NAME"] = "cpu"
+
+import numpy as np
+import typer
+import ml_dtypes  # register bfloat16 dtype for numpy  # noqa: F401
+from huggingface_hub import hf_hub_download, list_repo_files, snapshot_download
+from safetensors import safe_open
+from tqdm import tqdm
+
+import orbax.checkpoint as ocp
+
+
+APP_NAME = "hf-safetensors-to-orbax"
+ALLOWED_LAYER_PREFIXES = {"layers", "layer", "h", "blocks", "block"}
+PREFIXES_TO_STRIP = ("model.", "transformer.", "module.")
+WEIGHT_SUFFIX = ".weight"
+
+
+def normalize_output_path(path: str, local: bool) -> str:
+    """Validate and normalize the output location.
+
+    By default the converter writes to GCS and enforces a `gs://` prefix. When
+    the local flag is set, the path is treated as a local filesystem destination
+    and `gs://` paths are rejected to avoid accidental uploads.
+    """
+    if local:
+        if path.startswith("gs://"):
+            raise typer.BadParameter("--local requires a filesystem path")
+        return os.path.abspath(path)
+    if not path.startswith("gs://") or len(path) <= 5:
+        raise typer.BadParameter(
+            "--output must be a valid gs:// path unless --local is set"
+        )
+    return path.rstrip("/")
+
+
+def configure_logging() -> None:
+    """Configure structured logging for CLI execution.
+
+    The goal is to surface high-level progress milestones and critical context
+    (like stack shapes and layer ranges) without drowning the user in per-tensor
+    noise. The format is timestamped for long-running conversions.
+    """
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s | %(levelname)s | %(message)s",
+    )
+
+
+def configure_hf_transfer(byte_progress: bool) -> None:
+    """Enable or disable byte-level download progress via hf_transfer.
+
+    Hugging Face's default snapshot progress bar is per-file, which makes large
+    single-file models look stuck at 0/1 for a long time. Enabling hf_transfer
+    switches to a byte-level stream so progress reflects actual download volume.
+    """
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" if byte_progress else "0"
+
+
+def infer_snapshot_dir(file_path: str) -> str:
+    """Infer the snapshot directory from a cached Hugging Face file path.
+
+    Cached files live under `.../snapshots/<commit>/...`. This helper finds that
+    segment and returns the snapshot directory so subsequent code can operate on
+    a stable root without assuming a particular cache location.
+    """
+    parts = Path(file_path).parts
+    for idx, part in enumerate(parts):
+        if part == "snapshots" and idx + 1 < len(parts):
+            return str(Path(*parts[: idx + 2]))
+    raise ValueError(f"Unable to infer snapshot dir from path: {file_path}")
+
+
+def download_snapshot(repo_id: str, token: Optional[str], byte_progress: bool) -> str:
+    """Download the safetensors snapshot from Hugging Face.
+
+    The fast path uses `snapshot_download`, which only exposes file-count
+    progress. When byte-level progress is enabled, this function lists matching
+    files and downloads each with `hf_hub_download` to get per-byte visibility.
+    """
+    allow_patterns = ["*.safetensors", "*.safetensors.index.json"]
+    logging.info("Downloading snapshot for %s", repo_id)
+
+    if not byte_progress:
+        return snapshot_download(
+            repo_id=repo_id,
+            token=token,
+            allow_patterns=allow_patterns,
+        )
+
+    repo_files = list_repo_files(repo_id=repo_id, token=token)
+    filtered = [
+        name
+        for name in repo_files
+        if any(fnmatch.fnmatch(name, pat) for pat in allow_patterns)
+    ]
+    if not filtered:
+        raise ValueError("No matching safetensors files found in repository")
+
+    snapshot_dir: Optional[str] = None
+    for name in tqdm(sorted(filtered), desc="Downloading files"):
+        path = hf_hub_download(repo_id=repo_id, filename=name, token=token)
+        inferred = infer_snapshot_dir(path)
+        if snapshot_dir is None:
+            snapshot_dir = inferred
+        elif snapshot_dir != inferred:
+            raise ValueError("Downloaded files belong to different snapshot dirs")
+
+    if snapshot_dir is None:
+        raise ValueError("Unable to determine snapshot directory")
+    return snapshot_dir
+
+
+def find_index_json(snapshot_dir: str) -> Optional[str]:
+    """Locate a safetensors index JSON, if present.
+
+    Sharded HF checkpoints publish a single index JSON that maps tensor names
+    to shard files. We treat multiple index files as an error because it is not
+    a standard layout and likely indicates a mixed or malformed snapshot.
+    """
+    matches = list(Path(snapshot_dir).rglob("*.safetensors.index.json"))
+    if not matches:
+        return None
+    if len(matches) > 1:
+        raise ValueError("Multiple safetensors index JSON files found")
+    return str(matches[0])
+
+
+def build_key_to_file(snapshot_dir: str) -> Dict[str, str]:
+    """Build a mapping from tensor key to safetensors file path.
+
+    Prefer the official shard index when available to avoid redundant scans.
+    If no index exists (common for single-file models), scan the safetensors
+    files directly. Duplicate keys are treated as fatal to prevent corruption.
+    """
+    index_path = find_index_json(snapshot_dir)
+    if index_path:
+        with open(index_path, "r", encoding="utf-8") as handle:
+            data = json.load(handle)
+        weight_map = data.get("weight_map")
+        if not weight_map:
+            raise ValueError(f"Index JSON missing weight_map: {index_path}")
+        mapping: Dict[str, str] = {}
+        for key, rel_path in weight_map.items():
+            abs_path = os.path.join(snapshot_dir, rel_path)
+            if key in mapping:
+                raise ValueError(f"Duplicate key in index: {key}")
+            mapping[key] = abs_path
+        return mapping
+
+    safetensor_files = list(Path(snapshot_dir).rglob("*.safetensors"))
+    if not safetensor_files:
+        raise ValueError("No .safetensors files found in snapshot")
+
+    mapping = {}
+    for file_path in safetensor_files:
+        with safe_open(str(file_path), framework="numpy") as handle:
+            for key in handle.keys():
+                if key in mapping:
+                    raise ValueError(f"Duplicate key across files: {key}")
+                mapping[key] = str(file_path)
+    return mapping
+
+
+def parse_layer_key(key: str) -> Optional[Tuple[str, int, str]]:
+    """Parse a parameter key into (prefix, layer_index, suffix) if it looks stackable.
+
+    The heuristic looks for an integer segment that is preceded by a known layer
+    container name (layers, h, blocks, ...). The suffix is everything after the
+    index, which becomes the grouping key for stacking. The prefix is everything
+    before the layer container; it disambiguates multi-tower models that reuse
+    the same suffixes across different component stacks.
+    """
+    parts = key.split(".")
+    for idx in range(1, len(parts) - 1):
+        if parts[idx].isdigit() and parts[idx - 1] in ALLOWED_LAYER_PREFIXES:
+            prefix = ".".join(parts[: idx - 1]).strip(".")
+            suffix = ".".join(parts[idx + 1 :]).strip(".")
+            if suffix:
+                return prefix, int(parts[idx]), suffix
+    return None
+
+
+def group_layer_keys(
+    keys: Iterable[str],
+) -> Tuple[Dict[Tuple[str, str], Dict[int, str]], List[str], Dict[str, set[str]]]:
+    """Group layer parameters by (prefix, suffix) and separate non-layer parameters.
+
+    This pass identifies all stackable keys, buckets them by (prefix, suffix),
+    and filters the remaining parameters into a non-layer list. Only groups with
+    two or more distinct layer indices qualify for stacking to reduce false positives.
+    It also tracks which prefixes appear for a given suffix to drive naming.
+    """
+    suffix_groups: Dict[Tuple[str, str], Dict[int, str]] = {}
+    suffix_prefixes: Dict[str, set[str]] = {}
+
+    for key in keys:
+        parsed = parse_layer_key(key)
+        if not parsed:
+            continue
+        prefix, layer_idx, suffix = parsed
+        group = suffix_groups.setdefault((prefix, suffix), {})
+        if layer_idx in group:
+            raise ValueError(
+                f"Duplicate layer index {layer_idx} for prefix '{prefix}' and suffix '{suffix}'"
+            )
+        group[layer_idx] = key
+        suffix_prefixes.setdefault(suffix, set()).add(prefix)
+
+    stackable_suffixes = {
+        key: group for key, group in suffix_groups.items() if len(group) >= 2
+    }
+
+    stackable_keys = {
+        key for group in stackable_suffixes.values() for key in group.values()
+    }
+    non_layer_keys = [key for key in keys if key not in stackable_keys]
+
+    return stackable_suffixes, non_layer_keys, suffix_prefixes
+
+
+def load_tensor(file_path: str, key: str) -> np.ndarray:
+    """Load a single tensor from a safetensors file.
+
+    Safetensors reads are isolated to the specific key to keep IO and memory
+    tight. Callers handle validation and stacking; this only returns the array.
+    """
+    with safe_open(file_path, framework="numpy") as handle:
+        return handle.get_tensor(key)
+
+
+def validate_contiguous(indices: List[int], suffix: str) -> None:
+    """Assert that a suffix group has a contiguous layer index range.
+
+    Stacking assumes a dense layer axis. Missing indices likely mean a keying
+    bug or an unexpected model layout, so we fail fast with a helpful message.
+    """
+    expected = list(range(indices[0], indices[-1] + 1))
+    if indices != expected:
+        raise ValueError(
+            f"Missing layer indices for suffix '{suffix}': got {indices}, expected {expected}"
+        )
+
+
+def make_stacked_name(prefix: str, suffix: str, suffix_prefixes: Dict[str, set[str]]) -> str:
+    """Translate a (prefix, suffix) pair into the output stacked parameter name.
+
+    The suffix path is preserved while converting dots to slashes, and it is
+    rooted under `layers_stacked/`. If the same suffix appears under multiple
+    prefixes, the prefix path is included to avoid collisions and make the
+    output explicit for multi-tower models.
+    """
+    suffix_path = suffix.replace(".", "/")
+    if len(suffix_prefixes.get(suffix, set())) > 1:
+        prefix_path = prefix.replace(".", "/").strip("/")
+        if prefix_path:
+            return f"layers_stacked/{prefix_path}/{suffix_path}"
+    return f"layers_stacked/{suffix_path}"
+
+
+def clean_global_name(key: str) -> str:
+    """Normalize a non-layer parameter key into a clean output name.
+
+    This strips common framework prefixes and a trailing `.weight`, then
+    converts dot segments into path separators for a stable, readable tree.
+    """
+    cleaned = key
+    for prefix in PREFIXES_TO_STRIP:
+        if cleaned.startswith(prefix):
+            cleaned = cleaned[len(prefix) :]
+            break
+    if cleaned.endswith(WEIGHT_SUFFIX):
+        cleaned = cleaned[: -len(WEIGHT_SUFFIX)]
+    cleaned = cleaned.strip(".")
+    if not cleaned:
+        raise ValueError(f"Invalid empty name after cleaning: {key}")
+    return cleaned.replace(".", "/")
+
+
+def save_array_to_temp(temp_dir: str, name: str, array: np.ndarray) -> np.ndarray:
+    """Persist an array to disk and reopen it as a memory-mapped view.
+
+    Building the full output tree in RAM can be expensive for large models. By
+    staging arrays to disk and reopening with mmap, we limit peak memory while
+    still handing Orbax a stable array-like object for checkpoint writing. We
+    store raw bytes instead of `.npy` so non-standard dtypes like bfloat16 are
+    preserved when memory-mapped.
+    """
+    file_path = os.path.join(temp_dir, f"{name}.bin")
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+    if not array.flags["C_CONTIGUOUS"]:
+        array = np.ascontiguousarray(array)
+    array.tofile(file_path)
+    return np.memmap(
+        file_path,
+        mode="r",
+        dtype=array.dtype,
+        shape=array.shape,
+    )
+
+
+def convert(
+    hf_repo: str,
+    hf_token: Optional[str],
+    output_path: str,
+    stacking_config: str,
+    byte_progress: bool,
+    local: bool,
+) -> None:
+    """Run the full conversion pipeline from HF snapshot to Orbax checkpoint.
+
+    The flow is: download snapshot, build key-to-file map, identify stackable
+    groups, stack one suffix group at a time with shape validation, write staged
+    arrays into an output tree, then save via Orbax. This function enforces all
+    invariants and ensures temporary storage is cleaned up on failure.
+    """
+    if stacking_config != "auto":
+        raise typer.BadParameter("--stacking-config only supports 'auto'")
+
+    output_path = normalize_output_path(output_path, local)
+    configure_hf_transfer(byte_progress)
+    snapshot_dir = download_snapshot(hf_repo, hf_token, byte_progress)
+    key_to_file = build_key_to_file(snapshot_dir)
+
+    keys = sorted(key_to_file.keys())
+    stackable, non_layer_keys, suffix_prefixes = group_layer_keys(keys)
+
+    logging.info("Found %d stackable groups", len(stackable))
+    logging.info("Found %d non-layer params", len(non_layer_keys))
+
+    output_tree: Dict[str, np.ndarray] = {}
+    used_names = set()
+
+    temp_dir = tempfile.mkdtemp(prefix=f"{APP_NAME}-")
+    logging.info("Using temp dir: %s", temp_dir)
+
+    try:
+        for prefix, suffix in tqdm(sorted(stackable.keys()), desc="Stacking groups"):
+            group = stackable[(prefix, suffix)]
+            indices = sorted(group.keys())
+            validate_contiguous(indices, suffix)
+
+            first_key = group[indices[0]]
+            first_arr = load_tensor(key_to_file[first_key], first_key)
+            stacked = np.empty((len(indices),) + first_arr.shape, dtype=first_arr.dtype)
+            stacked[0] = first_arr
+            base_shape = first_arr.shape
+            logging.info(
+                "Stacking %s | prefix %s | layers %d-%d | shape %s | dtype %s",
+                suffix,
+                prefix or "<root>",
+                indices[0],
+                indices[-1],
+                base_shape,
+                stacked.dtype,
+            )
+            del first_arr
+
+            for pos, layer_idx in enumerate(indices[1:], start=1):
+                key = group[layer_idx]
+                arr = load_tensor(key_to_file[key], key)
+                if arr.shape != base_shape:
+                    raise ValueError(
+                        f"Shape mismatch for suffix '{suffix}': {arr.shape} vs {base_shape}"
+                    )
+                stacked[pos] = arr
+
+            stacked_name = make_stacked_name(prefix, suffix, suffix_prefixes)
+            if stacked_name in used_names:
+                raise ValueError(f"Output name collision: {stacked_name}")
+            used_names.add(stacked_name)
+
+            output_tree[stacked_name] = save_array_to_temp(
+                temp_dir, stacked_name, stacked
+            )
+            del stacked
+            gc.collect()
+
+        for key in tqdm(non_layer_keys, desc="Saving non-layer params"):
+            global_name = clean_global_name(key)
+            if global_name in used_names:
+                raise ValueError(f"Output name collision: {global_name}")
+            used_names.add(global_name)
+
+            arr = load_tensor(key_to_file[key], key)
+            output_tree[global_name] = save_array_to_temp(temp_dir, global_name, arr)
+            del arr
+
+        logging.info("Saving Orbax checkpoint to %s", output_path)
+        checkpointer = ocp.StandardCheckpointer()
+        checkpointer.save(output_path, output_tree)
+        checkpointer.wait_until_finished()
+        checkpointer.close()
+        logging.info("Save complete")
+    finally:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+
+
+app = typer.Typer(add_completion=False)
+
+
+@app.command()
+def main(
+    hf_repo: str = typer.Option(..., help="Hugging Face repo ID"),
+    hf_token: Optional[str] = typer.Option(None, help="Hugging Face auth token"),
+    output: str = typer.Option(
+        ...,
+        "--output",
+        help="Output path (gs://... unless --local)",
+    ),
+    stacking_config: str = typer.Option("auto", help="Stacking strategy (only 'auto')"),
+    byte_progress: bool = typer.Option(
+        True,
+        "--byte-progress/--no-byte-progress",
+        help="Use byte-level progress via hf_transfer",
+    ),
+    local: bool = typer.Option(
+        False,
+        "--local/--no-local",
+        help="Write to a local filesystem path instead of GCS",
+    ),
+) -> None:
+    """CLI entrypoint that wires arguments to the conversion routine.
+
+    This keeps the CLI surface thin while centralizing the conversion logic in
+    `convert`, which makes it easier to test or reuse programmatically.
+    """
+
+    configure_logging()
+    convert(
+        hf_repo=hf_repo,
+        hf_token=hf_token,
+        output_path=output,
+        stacking_config=stacking_config,
+        byte_progress=byte_progress,
+        local=local,
+    )
+
+
+if __name__ == "__main__":
+    try:
+        app()
+    except Exception as exc:  # pragma: no cover - CLI failure handling
+        logging.error("Conversion failed: %s", exc)
+        sys.exit(1)

ocdbt.process_0/d/0c600442de678f199ab5e8d0152c8bbc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c76d615abe2ee81f238c013b0c08bcba9a670b533ed4db511cf9db9f46ea401
+size 2754031616

ocdbt.process_0/d/124c8d55117974805f4a148ebd670601

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0517cabc606bdc13d666180e4c6b8c87f79561733cc17e9ffcc40eef4cc2f6a1
+size 2780028928

ocdbt.process_0/d/156c1d2fa9763c506b9a075a5729b7de

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e30ef7664e63a6d0f57cd3504ff41b6a4ef6069b885beeb55c3e4a4bd89e666
+size 2825613312

ocdbt.process_0/d/2f806db10cde1c0f439f288f34ee8b95

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b01357e5cc845093e19d68445e120595d5f601655f5eb9be19bb4bd59724332
+size 2905985024

ocdbt.process_0/d/44d6c5046ef1352f82da17c3daffc66d

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4e105c721ad7c6acc26f4649a392a83573424fb9aa2cd0626c35b98afeaf571
+size 2779865088

ocdbt.process_0/d/4f87c6488b21c3c03780afba253c8af0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd0ce378f0c0abbc96571437f4abbedc881082cfb4ab99e1a217d4728f69a9ea
+size 2906173440

ocdbt.process_0/d/78efb519de44bcebc6eb6fd345ff0af0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:456831e79cd197fcf3a54bb82d96cc29d5d12c2cd1fa985e2839479ecd5c5d55
+size 2754076672

ocdbt.process_0/d/82c6ee1c6166b74a75b93badf224480e

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f59acf6703a68ed463ffdeb6167cc96c5c1a2c929230624e74c8789b916559e
+size 19894272

ocdbt.process_0/d/88b7b61e8b4c3f22ec014bd4fcd5cbe5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0338a2136ca337df63f08ebfb302dc5cede3d86ba4575d38bd7c9589ae849c07
+size 477347

ocdbt.process_0/d/8d2385834c5546d282b60ec4c29d4378

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a0d8676748e0392c1cb5f555e19e3f31a7cfd3177e79fe934a3a7594ab84842
+size 2779918336

ocdbt.process_0/d/97411cb3209e6a07d86332228942a45c

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d210bbe1657bfc07fc0fa3d3d113c9c8cf4a3cf772a51daffdc00b7fe069335
+size 2303647744

ocdbt.process_0/d/a4f076e56e9ec995bf74b8f5d19e4101

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b277181db8e3ad9504f4ab24e6039b3d0a553237d7c1b9f887f676c53b550c8
+size 2906382336

ocdbt.process_0/d/ae742711b7e4d0ae7a8eaa239393ec58

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a3620e6e429f02c3ed717690a1baed22654e9fae8a89cfc02323a035590006f
+size 2779967488

ocdbt.process_0/d/ca081ceb253837d87f9bc28d6dfe2b0a

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34504dc0900d79fff51a81e34dc12b56d1e1b4daaee31e70b23e6c3e26bd8d84
+size 2906353664

ocdbt.process_0/d/ce563dd53d0cc8894641b5996dc5453f

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:811d04411b96a8c78452d8f8348f03a9aa8bab73aca80953fdfdda74ac990d82
+size 3030417408

ocdbt.process_0/d/df6d4f33d5f964745802e5eedfe1425e

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e545021423358467819c42e0dfa71bf7ab7f0f5bd117dac4884a3a092a6921fa
+size 2303664128

ocdbt.process_0/d/e61206338bedb2c7eecd25dc76511b00

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b89841d5d259a3d222eda23ee072c1daf9382418fdc3f52eb1d581c87e83ae75
+size 926654464

ocdbt.process_0/d/f43da177788619ad3cd686b1345608ef

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8a56b172509bc843430abfcf32526a185ee59d312f1ded36ee2c317fb740399
+size 2754002944