Merge pull request #3966 from AI-Hypercomputer:darisoy-fix-pathways-conversion

Google-ML-Automation · Google-ML-Automation · commit 7249840d7e9a · 2026-05-21T20:42:33.000-07:00
PiperOrigin-RevId: 919409009
diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py
@@ -790,6 +790,8 @@ def from_pretrained(
             "skip_jax_distributed_system=True",
             "--lazy_load_tensors=True",
             f"--simulated_cpu_devices_count={simulated_cpu_devices_count}",
+            f"checkpoint_storage_use_ocdbt={config.checkpoint_storage_use_ocdbt}",
+            f"checkpoint_storage_use_zarr3={config.checkpoint_storage_use_zarr3}",
         ]
 
         try:

Original file line number	Diff line number	Diff line change
`@@ -790,6 +790,8 @@ def from_pretrained(`
`790`	`790`	`"skip_jax_distributed_system=True",`
`791`	`791`	`"--lazy_load_tensors=True",`
`792`	`792`	`f"--simulated_cpu_devices_count={simulated_cpu_devices_count}",`
	`793`	`+ f"checkpoint_storage_use_ocdbt={config.checkpoint_storage_use_ocdbt}",`
	`794`	`+ f"checkpoint_storage_use_zarr3={config.checkpoint_storage_use_zarr3}",`
`793`	`795`	`]`
`794`	`796`
`795`	`797`	`try:`