memory optimization

Author: khatwanimohit

src/MaxText/configs/base.yml

@@ -67,6 +67,8 @@ checkpoint_storage_use_zarr3: True
 # larger models requires higher concurrent GB for I/O
 # default concurrent gb for PytreeCheckpointHandler is 96GB
 checkpoint_storage_concurrent_gb: 96
+# lazyily load tensors for checkpoint conversion
+lazy_load_tensors: False
 
 # Bool flag for enabling Orbax v1.
 enable_orbax_v1: False