Skip to content

[DECOUPLED-MODE] Checkpointing: import GCS via decoupling stub in dynamic loader #91

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
gulsumgudukbay merged 1 commit into from
Jul 16, 2026
Merged

[DECOUPLED-MODE] Checkpointing: import GCS via decoupling stub in dynamic loader #91

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

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion src/maxtext/checkpoint_conversion/utils/load_dynamic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,19 +73,23 @@

from flax import nnx
import flax.traverse_util
from google.cloud import storage
import huggingface_hub
import jax
from maxtext.checkpoint_conversion.utils import hf_model_configs
from maxtext.checkpoint_conversion.utils import param_mapping
from maxtext.checkpoint_conversion.utils import tensor_handling
from maxtext.common.gcloud_stub import gcs_storage
from maxtext.utils import gcs_utils
from maxtext.utils import globals as maxtext_globals
from maxtext.utils import max_logging
from orbax.checkpoint import v1 as ocp_v1
from orbax.checkpoint._src.arrays import sharding as sharding_utils


# Route GCS through the decoupling helper so this module imports cleanly in
# decoupled environments where google-cloud-storage is intentionally absent.
storage = gcs_storage()

HF_MODEL_CONFIGS = hf_model_configs.HF_MODEL_CONFIGS
get_hf_loading_function = tensor_handling.get_hf_loading_function

Expand Down
Loading