Add roundtrip ckpt conversion tests in CI by hengtaoguo · Pull Request #3909 · AI-Hypercomputer/maxtext

hengtaoguo · 2026-05-14T18:54:45Z

Description

Make import torch lazy import in to_maxtext, so we can run torch-free conversion in CI.
Add a 2-layered qwen3-30b-a3b ckpt-conversion-roundtrip cpu-only test in CI.
-- We lazy-load safetensors so only 2-out-of-16 chunks were actually loaded for to_maxtext conversion.
-- Within the tmpdir, we convert maxtext Orbax ckpt back to HF for to_huggingface conversion.
-- All converted MaxText/HF ckpts will be auto removed from tmpdir.
-- Success criteria: (1) Subprocess finishes without errors (2) Result ckpts exist in the tmpdir. We didn't enforce logits check in the CI because that will actually require torch dependency. But I've tested forward_pass_logit_checker offline.
Fix b/513636729, introduced in cl/914979909.

Tests

# to maxtext
python -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml model_name=qwen3-30b-a3b base_output_directory=/dev/shm/mini base_num_decoder_layers=2 override_model_config=True scan_layers=false hf_access_token=\<token\> hardware=cpu skip_jax_distributed_system=True checkpoint_storage_use_ocdbt=False checkpoint_storage_use_zarr3=False --eager_load_method=safetensors --save_dtype=bfloat16 --lazy_load_tensors=True

# to huggingface
python -m maxtext.checkpoint_conversion.to_huggingface model_name=qwen3-30b-a3b load_parameters_path=/dev/shm/mini/0/items base_output_directory=/dev/shm/mini/hf_safetensor/qwen3-30b-a3b base_num_decoder_layers=2 override_model_config=True scan_layers=false weight_dtype=bfloat16 hardware=cpu skip_jax_distributed_system=True hf_access_token=\<token\> --override_model_architecture=True

# logits check (roundtrip above)
python -m tests.utils.forward_pass_logit_checker maxtext/configs/base.yml run_name=ht_test model_name=qwen3-30b-a3b tokenizer_path=Qwen/Qwen3-30B-A3B-Thinking-2507 load_parameters_path=/dev/shm/mini/0/items base_num_decoder_layers=2 override_model_config=True scan_layers=false max_prefill_predict_length=4 ici_tensor_parallelism=4 hf_access_token=\<token\> --run_hf_model=True --hf_model_path=Qwen/Qwen3-30B-A3B-Thinking-2507 --hf_model_path=/dev/shm/mini/hf_safetensor/qwen3-30b-a3b --max_kl_div=0.015

2-layerd mini checkpoint with matching gibberish logits: Logs

Tests active in CI: https://screenshot.googleplex.com/4ZhRPiqJZLxTXk5

For this single test on my v5p-8 VM, it took ~2.5 minutes and the peak memory was ~12G.

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.

codecov · 2026-05-14T18:59:28Z

Codecov Report

❌ Patch coverage is 0% with 12 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/checkpoint_conversion/to_maxtext.py	0.00%	12 Missing ⚠️

📢 Thoughts on this report? Let us know!

shuningjin

Looks good! Some minor comments.

shuningjin · 2026-05-15T21:51:46Z

+          "weight_dtype=bfloat16",
+          "hardware=cpu",
+          "skip_jax_distributed_system=True",
+          "--override_model_architecture=True",


unrelated to this pr, for follow-up: I wonder if --override_model_architecture=True can be directly replaced with override_model_config=True

IIUC override_model_config is for overriding MaxText models, while override_model_architecture is for overriding/dumping HF config.json. I also agree we should consolidate this flag in a follow up PR.

github-actions · 2026-05-16T20:54:52Z

🤖 CI Failure Investigation Report (Automated)

I have analyzed the recent test failures that occurred after this PR was merged.

Findings:
PR #3909 intended to make the checkpoint conversion torch-free, but several torch. references remain in to_maxtext.py, causing failures in environments without torch (like the tpu-integration CI job).

Representative Failure:

Failing Test: tests/integration/checkpoint_conversion_test.py::Qwen3CheckpointConversionTest::test_qwen3_30b_a3b_roundtrip_conversion
Likely Error: AttributeError: 'NoneType' object has no attribute 'bfloat16'

Recommended Fix:
Guard all torch. attribute accesses in src/maxtext/checkpoint_conversion/to_maxtext.py and ensure safetensors uses framework="np" when torch is not available.

See .gemini/findings.md for full details.

github-actions · 2026-05-17T17:03:45Z

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline and identified the following:

🔍 What Failed

Job/Matrix: tpu-integration (and other torch-free environments)
Failing Test: tests/integration/checkpoint_conversion_test.py::Qwen3CheckpointConversionTest::test_qwen3_30b_a3b_roundtrip_conversion
Error: AttributeError: 'NoneType' object has no attribute 'bfloat16'

🪵 Error Details & Stack Trace

# Representative traceback from to_maxtext.py (inferred from code analysis and automated triage)
  File "src/maxtext/checkpoint_conversion/to_maxtext.py", line 945, in _eager_getter
    if v.dtype == torch.bfloat16:
AttributeError: 'NoneType' object has no attribute 'bfloat16'

💡 Root Cause Analysis & Context

Confidence: High

The failure was caused by unguarded accesses to torch. attributes in src/maxtext/checkpoint_conversion/to_maxtext.py. While PR #3909 introduced lazy loading of torch to support torch-free environments, it missed several locations where torch.bfloat16 or torch.float32 were accessed. In environments where torch is not installed (such as some TPU integration jobs), torch is initialized to None, leading to the AttributeError.

Additionally, some calls to load_hf_dict_from_safetensors were defaulting to framework="pt", which also requires torch.

🛠️ Recommended Fix

I have applied a fix to src/maxtext/checkpoint_conversion/to_maxtext.py that:

Guards all torch. attribute accesses with if torch and ....
Handles both torch tensors and numpy arrays in _eager_getter.
Automatically switches safetensors framework to "np" when torch is not available.

# Example of the fix applied:
if torch and hasattr(v, "dtype") and v.dtype == torch.bfloat16:
    v = v.to(torch.float32)
if hasattr(v, "numpy"):
    v = v.numpy()
return v.astype(ml_dtypes.bfloat16)

github-actions · 2026-05-18T17:33:40Z

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline and identified the following:

🔍 What Failed

Job/Matrix: tpu-integration (and other environments where torch is not installed)
Failing Test: tests/integration/checkpoint_conversion_test.py::Qwen3CheckpointConversionTest::test_qwen3_30b_a3b_roundtrip_conversion
Error: AttributeError: 'NoneType' object has no attribute 'bfloat16'

🪵 Error Details & Stack Trace

# The failure occurs in environments where torch is not available.
# torch is imported as None via _lazy_import("torch").
# At line 945 (and other locations), the code attempts to access torch.bfloat16.

  File "src/maxtext/checkpoint_conversion/to_maxtext.py", line 945, in _eager_getter
    if v.dtype == torch.bfloat16:
AttributeError: 'NoneType' object has no attribute 'bfloat16'

💡 Root Cause Analysis & Context

Confidence: High

The root cause is the lack of guards around torch attribute accesses in src/maxtext/checkpoint_conversion/to_maxtext.py. While PR #3909 aimed to make the checkpoint conversion process torch-free, it introduced several direct references to torch.bfloat16 and torch.float32. In CI environments where torch is not installed (common for JAX/TPU-focused jobs), _lazy_import("torch") returns None, causing AttributeError when these attributes are accessed.

Furthermore, load_hf_dict_from_safetensors was defaulting to the pt (PyTorch) framework, which also requires torch.

🛠️ Recommended Fix

I have applied a fix to src/maxtext/checkpoint_conversion/to_maxtext.py in the workspace that:

Guards all torch attribute accesses with if torch and ....
Dynamically switches the safetensors framework between pt and np based on torch availability.
Added hasattr(v, "numpy") and hasattr(v, "detach") checks to handle both PyTorch tensors and NumPy arrays gracefully.

# Example of the fix applied:
if torch and hasattr(v, "dtype") and v.dtype == torch.bfloat16:
    v = v.to(torch.float32)
if hasattr(v, "numpy"):
    v = v.numpy()
return np.array(v).astype(ml_dtypes.bfloat16)

hengtaoguo force-pushed the hengtaoguo-ckpt-ci branch 8 times, most recently from 592aedf to e6fcf3a Compare May 15, 2026 20:46

hengtaoguo changed the title ~~Try ckpt conversion tests in CI~~ Add roundtrip ckpt conversion tests in CI May 15, 2026

hengtaoguo marked this pull request as ready for review May 15, 2026 20:50

hengtaoguo requested review from A9isha, NicoGrande, NuojCheng, RissyRan, SurbhiJainUSC, abhinavclemson, aireenmei, bvandermoon, gagika, gobbleturk, jesselu-google, jiangjy1982, khatwanimohit, parambole, richjames0, shralex, shuningjin, suexu1025 and vipannalla as code owners May 15, 2026 20:50

hengtaoguo requested review from darisoy, dipannita08 and igorts-git as code owners May 15, 2026 20:50

SurbhiJainUSC approved these changes May 15, 2026

View reviewed changes

shuningjin approved these changes May 15, 2026

View reviewed changes

github-actions Bot added pull ready labels May 15, 2026

Add ckpt conversion test in CI

4f2de47

hengtaoguo force-pushed the hengtaoguo-ckpt-ci branch from 4683db9 to 4f2de47 Compare May 15, 2026 22:20

doc

6977365

copybara-service Bot merged commit b5e5330 into main May 16, 2026
28 of 29 checks passed

copybara-service Bot deleted the hengtaoguo-ckpt-ci branch May 16, 2026 01:03

github-actions Bot mentioned this pull request May 18, 2026

Failed build: MaxText Package Tests #3765

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add roundtrip ckpt conversion tests in CI#3909

Add roundtrip ckpt conversion tests in CI#3909
copybara-service[bot] merged 2 commits into
mainfrom
hengtaoguo-ckpt-ci

hengtaoguo commented May 14, 2026 •

edited

Loading

Uh oh!

codecov Bot commented May 14, 2026 •

edited

Loading

Uh oh!

shuningjin left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

shuningjin May 15, 2026

Uh oh!

hengtaoguo May 15, 2026

Uh oh!

Uh oh!

github-actions Bot commented May 16, 2026

Uh oh!

github-actions Bot commented May 17, 2026

Uh oh!

github-actions Bot commented May 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

hengtaoguo commented May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

codecov Bot commented May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

shuningjin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

shuningjin May 15, 2026

Choose a reason for hiding this comment

Uh oh!

hengtaoguo May 15, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot commented May 16, 2026

🤖 CI Failure Investigation Report (Automated)

Uh oh!

github-actions Bot commented May 17, 2026

🤖 CI Failure Investigation Report

🔍 What Failed

🪵 Error Details & Stack Trace

💡 Root Cause Analysis & Context

🛠️ Recommended Fix

Uh oh!

github-actions Bot commented May 18, 2026

🤖 CI Failure Investigation Report

🔍 What Failed

🪵 Error Details & Stack Trace

💡 Root Cause Analysis & Context

🛠️ Recommended Fix

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hengtaoguo commented May 14, 2026 •

edited

Loading

codecov Bot commented May 14, 2026 •

edited

Loading