Skip to content

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

Merged
gulsumgudukbay merged 1 commit into
rocm-mainfrom
fix/dynamic-loader-decoupled-gcs
Jul 16, 2026
Merged

[DECOUPLED-MODE] Checkpointing: import GCS via decoupling stub in dynamic loader#91
gulsumgudukbay merged 1 commit into
rocm-mainfrom
fix/dynamic-loader-decoupled-gcs

Conversation

@gulsumgudukbay

Copy link
Copy Markdown
Collaborator

Description

Route the dynamic HuggingFace checkpoint loader's Google Cloud Storage access through the existing decoupling helper instead of importing it directly.

Previously, src/maxtext/checkpoint_conversion/utils/load_dynamic.py did a top-level from google.cloud import storage. In "decoupled" environments (DECOUPLE_GCLOUD=TRUE), where optional Google Cloud dependencies are intentionally not installed, this unconditional import raised ImportError: cannot import name 'storage' from 'google.cloud' at module load.
Because load_dynamic is pulled in transitively by maxtext.common.checkpointing (and thus by maxtext/__init__), the error
aborted the entire pytest collection at tests/conftest.py import time, so the full decoupled test suite failed before running a single test.

The rest of the codebase already accesses GCS through maxtext.common.gcloud_stub.gcs_storage(), which returns the real
google.cloud.storage module when available and a no-op stub otherwise. This PR makes load_dynamic.py follow that same pattern:

  • Replace the direct from google.cloud import storage with from maxtext.common.gcloud_stub import gcs_storage.
  • Bind storage = gcs_storage() at module scope, so the existing storage.Client() / list_blobs(...) call sites are unchanged.

This is a minimal, low-risk change: in non-decoupled runs behavior is identical (real google.cloud.storage, now with transfer_manager attached by the helper); in decoupled runs the module imports cleanly and GCS calls are stubbed.
No functional GCS behavior changes for real workloads.

Tests

Verified in the ROCm TheRock 7.14 MaxText container with the decoupled venv (DECOUPLE_GCLOUD=TRUE, PYTHONPATH=src):

  • Direct import sanity: python -c "import maxtext.checkpoint_conversion.utils.load_dynamic as m; print(m.storage._IS_STUB)"
    imports OK, prints True (stub) in decoupled mode.
  • Decoupled pytest collection (previously crashed) now succeeds: pytest -m 'not cpu_only and not tpu_only and not post_training and decoupled' tests --ignore=tests/post_training 2664 collected / 481 deselected / 2 skipped / 2183 selected, no collection error; tests execute normally.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@gulsumgudukbay
gulsumgudukbay merged commit 3bf4312 into rocm-main Jul 16, 2026
4 checks passed
@gulsumgudukbay
gulsumgudukbay deleted the fix/dynamic-loader-decoupled-gcs branch July 16, 2026 17:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant