XPK Start: Thu Apr 23 18:40:06 UTC 2026 2026-04-23 18:40:10.595098: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered WARNING: All log messages before absl::InitializeLog() is called are written to STDERR E0000 00:00:1776969610.607373 10 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered E0000 00:00:1776969610.611114 10 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered W0000 00:00:1776969610.622624 10 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776969610.622645 10 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776969610.622647 10 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776969610.622649 10 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. 2026-04-23 18:40:29.715998: E external/local_xla/xla/stream_executor/cuda/cuda_platform.cc:51] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303) I0423 18:40:30.238543 137979238270784 max_utils.py:273] Attempting to initialize the jax distributed system... INFO:2026-04-23 18:40:39,277:jax._src.distributed:140: Starting JAX distributed service on [::]:8482 I0423 18:40:39.277990 137979238270784 distributed.py:140] Starting JAX distributed service on [::]:8482 INFO:2026-04-23 18:40:39,280:jax._src.distributed:157: Connecting to JAX distributed service on mt-15-ocdbt-false-sv-wfyjm-slice-job-0-0.mt-15-ocdbt-false-sv-wfyjm:8482 I0423 18:40:39.280234 137979238270784 distributed.py:157] Connecting to JAX distributed service on mt-15-ocdbt-false-sv-wfyjm-slice-job-0-0.mt-15-ocdbt-false-sv-wfyjm:8482 I0423 18:40:40.494333 137979238270784 max_utils.py:284] Jax distributed system initialized! I0423 18:40:47.328984 137979238270784 max_utils.py:800] System Information: Jax Version: 0.8.1 I0423 18:40:47.329089 137979238270784 max_utils.py:801] System Information: Jaxlib Version: 0.8.1 I0423 18:40:47.329130 137979238270784 max_utils.py:802] System Information: Jax Backend: PJRT C API TFRT TPU v6 lite Built on Nov 12 2025 14:16:36 (1762985796) cl/831091709 I0423 18:40:47.329165 137979238270784 train_utils.py:377] WARNING: Sequence packing is essentially ignored for synthetic data. Please use a real dataset to use sequence packing. I0423 18:40:47.947432 137979238270784 maxtext_utils.py:1631] Num_devices: 32, shape (1, 1, 1, 32, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0423 18:40:48.061638 137979238270784 checkpointing.py:688] Setting up checkpoint logger... I0423 18:40:48.061756 137979238270784 checkpointing.py:234] Creating checkpoint manager with ocdbt=False and zarr3=True I0423 18:40:48.061807 137979238270784 pytree_checkpoint_handler.py:589] save_device_host_concurrent_bytes=None I0423 18:40:48.062038 137979238270784 base_pytree_checkpoint_handler.py:415] Created BasePyTreeCheckpointHandler: use_ocdbt=False, use_zarr3=True, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7d7d4f174d10>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0423 18:40:50.964496 137979238270784 checkpointing.py:266] Enabling policy for fixed interval checkpointing. I0423 18:40:50.964795 137979238270784 checkpoint_manager.py:709] [process=3][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=('items',), item_handlers={'items': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7d686865d5e0>}, handler_registry=None I0423 18:40:50.965043 137979238270784 composite_checkpoint_handler.py:237] Deferred registration for item: "items". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7d686865d5e0>` for item "items" and save args `<class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>` to `_handler_registry`. I0423 18:40:50.965095 137979238270784 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7d69b2b4a510>` for item "metrics" and save args `<class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>` to `_handler_registry`. I0423 18:40:50.965133 137979238270784 composite_checkpoint_handler.py:505] Initialized registry DefaultCheckpointHandlerRegistry({('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7d686865d5e0>, ('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7d686865d5e0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7d69b2b4a510>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7d69b2b4a510>}). I0423 18:40:50.965461 137979238270784 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.33 I0423 18:40:50.965541 137979238270784 async_checkpointer.py:177] [process=3][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7d68687e9260> timeout: 600 secs and primary_host=0 for async checkpoint writes I0423 18:40:52.441988 137979238270784 checkpoint_manager.py:1818] Found 0 checkpoint steps in gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints I0423 18:40:52.454207 137979238270784 checkpoint_manager.py:929] [process=3][thread=MainThread] CheckpointManager created, primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=None, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=False, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_background_delete=False, read_only=False, enable_async_checkpointing=True, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=FixedIntervalPolicy(interval=5), preservation_policy=LatestN(n=None), prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False, lightweight_initialize=False), root_directory=gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7d69b3ae5ac0> I0423 18:40:52.454327 137979238270784 checkpointing.py:302] Checkpoint manager created! I0423 18:40:53.386147 137979238270784 checkpointing.py:578] checkpoint manager exists so trying to load this run's existing checkpoint I0423 18:40:53.386268 137979238270784 checkpointing.py:676] No existing checkpoints found, not restoring checkpoint. fsdp: 32 I0423 18:40:54.800799 137979238270784 maxtext_utils.py:1749] decoder/decoder_norm/scale/value Shape: float32[2048] Physical: (None,) I0423 18:40:54.800919 137979238270784 maxtext_utils.py:1749] decoder/layers/mlp/wi_0/kernel/value Shape: float32[2048,16,7168] Physical: ('fsdp', None, None) I0423 18:40:54.800990 137979238270784 maxtext_utils.py:1749] decoder/layers/mlp/wi_1/kernel/value Shape: float32[2048,16,7168] Physical: ('fsdp', None, None) I0423 18:40:54.801037 137979238270784 maxtext_utils.py:1749] decoder/layers/mlp/wo/kernel/value Shape: float32[7168,16,2048] Physical: (None, None, 'fsdp') I0423 18:40:54.801076 137979238270784 maxtext_utils.py:1749] decoder/layers/post_self_attention_layer_norm/scale/value Shape: float32[2048,16] Physical: (None, None) I0423 18:40:54.801112 137979238270784 maxtext_utils.py:1749] decoder/layers/pre_self_attention_layer_norm/scale/value Shape: float32[2048,16] Physical: (None, None) I0423 18:40:54.801150 137979238270784 maxtext_utils.py:1749] decoder/layers/self_attention/key/kernel/value Shape: float32[2048,16,16,128] Physical: ('fsdp', None, None, None) I0423 18:40:54.801188 137979238270784 maxtext_utils.py:1749] decoder/layers/self_attention/out/kernel/value Shape: float32[16,16,128,2048] Physical: (None, None, None, 'fsdp') I0423 18:40:54.801223 137979238270784 maxtext_utils.py:1749] decoder/layers/self_attention/query/kernel/value Shape: float32[2048,16,16,128] Physical: ('fsdp', None, None, None) I0423 18:40:54.801258 137979238270784 maxtext_utils.py:1749] decoder/layers/self_attention/value/kernel/value Shape: float32[2048,16,16,128] Physical: ('fsdp', None, None, None) I0423 18:40:54.801293 137979238270784 maxtext_utils.py:1749] decoder/logits_dense/kernel/value Shape: float32[2048,32000] Physical: ('fsdp', None) I0423 18:40:54.801327 137979238270784 maxtext_utils.py:1749] token_embedder/embedding/value Shape: float32[32000,2048] Physical: (None, 'fsdp') I0423 18:40:55.196304 137979238270784 nnx_decoders.py:616] nnx_decoders/carry Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0423 18:40:55.196452 137979238270784 nnx_decoders.py:616] nnx_decoders/carry Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0423 18:40:55.201808 137979238270784 nnx_decoders.py:616] Unknown Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0423 18:40:55.201902 137979238270784 nnx_decoders.py:616] Unknown Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0423 18:40:55.350368 137979238270784 attentions.py:1084] attentions/inputs_q Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0423 18:40:55.350501 137979238270784 attentions.py:1084] attentions/inputs_q Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0423 18:40:55.498318 137979238270784 attentions.py:1085] attentions/inputs_kv Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0423 18:40:55.498452 137979238270784 attentions.py:1085] attentions/inputs_kv Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0423 18:40:55.655109 137979238270784 attentions.py:1150] attentions/query Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0423 18:40:55.655245 137979238270784 attentions.py:1150] attentions/query Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0423 18:40:55.797845 137979238270784 attentions.py:1151] attentions/key Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0423 18:40:55.797999 137979238270784 attentions.py:1151] attentions/key Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0423 18:40:55.937627 137979238270784 attentions.py:1152] attentions/value Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0423 18:40:55.937764 137979238270784 attentions.py:1152] attentions/value Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0423 18:40:56.088733 137979238270784 attentions.py:1193] attentions/out Logical: bfloat16[32,2048,16,128].................................... ('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv'). I0423 18:40:56.088872 137979238270784 attentions.py:1193] attentions/out Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0423 18:40:56.111522 137979238270784 linears.py:541] linears/x Logical: bfloat16[32,2048,7168]...................................... ('activation_batch', 'activation_length', 'activation_mlp'). I0423 18:40:56.111617 137979238270784 linears.py:541] linears/x Physical: bfloat16[32,2048,7168]...................................... ('fsdp', None, None). I0423 18:41:07.010857 137979238270784 max_utils.py:791] Total memory size: 1.5 GB, Output size: 0.4 GB, Temp size: 1.1 GB, Argument size: 0.4 GB, Host temp size: 0.0 GB. I0423 18:41:07.013206 137979238270784 metric_logger.py:289] number parameters: 1.104 billion I0423 18:41:18.046080 137979238270784 checkpointing.py:794] Waiting for step 0 to finish before checkpoint... I0423 18:41:18.197125 137979238270784 checkpointing.py:798] Waited 0.15103459358215332 seconds for step 0 to finish before starting checkpointing. I0423 18:41:18.199347 137979238270784 checkpoint_manager.py:2013] [process=3][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0423 18:41:18.200945 137979238270784 checkpoint_manager.py:1518] [process=3] Saving checkpoint at step 0 I0423 18:41:18.203731 137979238270784 async_checkpointer.py:452] [process=3] Started async saving checkpoint to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/0. I0423 18:41:18.810507 137979238270784 signaling_client.py:364] Using JaxDistributedSignalingClient I0423 18:41:18.811624 137979238270784 jax_array_handlers.py:358] Scheduling D2H of 69 prioritized jax.Array. I0423 18:41:18.811683 137979238270784 replica_slices.py:424] Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=False I0423 18:41:19.056055 137979238270784 base_pytree_checkpoint_handler.py:153] [process=3][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.246680s I0423 18:41:19.056228 137979238270784 base_pytree_checkpoint_handler.py:129] [process=3] /jax/checkpoint/write/blocking_gbytes_per_sec: 2.817 GiB/s (total gbytes: 1.5 GiB) (time elapsed: 0.5476527214050293 s) (per-host) I0423 18:41:19.056282 137979238270784 base_pytree_checkpoint_handler.py:737] [process=3][thread=MainThread] Initiated Pytree async_save. Time taken: 0.547718s (batch_requests_ready=0.294477s, total_serialization_initiated=0.253168s, others=0.000073s) I0423 18:41:19.056380 137979238270784 composite_checkpoint_handler.py:715] [process=3][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.551870s (all_items=0.000018s, per_item={'items': '0.00001836'}, temp_paths=0.551852) I0423 18:41:19.057325 137851933026048 async_checkpointer.py:79] [process=3][thread=async_save] Background save thread started. I0423 18:41:19.057490 137979238270784 async_checkpointer.py:561] Finished blocking save. Time taken: 0.856477s. Continuing background save to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/0. I0423 18:41:19.065973 137979238270784 checkpoint_manager.py:1566] [process=3][thread=MainThread][step=0] Starting CheckpointManager Save Finalize thread=save_finalize I0423 18:41:19.066246 137850867676928 async_checkpointer.py:265] [process=3][thread=save_finalize] Waiting for background save thread=async_save. I0423 18:41:19.066377 137979238270784 standard_logger.py:34] {'step': 0, 'event_type': 'save', 'directory': 'gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776969678.199327, 'wait_for_prev_duration_secs': 6.628036499023438e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776969678.2009838, 'checkpointer_blocking_duration_secs': 0.8566124439239502, 'get_old_steps_start_time': 1776969679.0576174, 'get_old_steps_duration_secs': 2.4557113647460938e-05, 'checkpoint_manager_blocking_start_time': 1776969678.1975121, 'checkpoint_manager_blocking_duration_secs': 0.8688294887542725} I0423 18:41:19.066488 137979238270784 checkpointing.py:409] Started an asynchronous checkpoint save for step 0 I0423 18:41:19.066558 137979238270784 max_utils.py:750] Memstats: After params initialized: I0423 18:41:19.066607 137979238270784 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_10(process=3,(2,2,0,0)) I0423 18:41:19.066642 137979238270784 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_11(process=3,(3,2,0,0)) I0423 18:41:19.066669 137979238270784 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_14(process=3,(2,3,0,0)) I0423 18:41:19.066695 137979238270784 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_15(process=3,(3,3,0,0)) I0423 18:41:19.333010 137979238270784 metric_logger.py:185] completed step: 0, seconds: 11.032, TFLOP/s/device: 1.232, Tokens/s/device: 185.645, total_weights: 65536, loss: 10.874 I0423 18:41:19.511440 137979238270784 metric_logger.py:185] completed step: 1, seconds: 1.285, TFLOP/s/device: 10.572, Tokens/s/device: 1593.593, total_weights: 65536, loss: 10.874 I0423 18:41:19.939476 137979238270784 metric_logger.py:185] completed step: 2, seconds: 0.033, TFLOP/s/device: 413.259, Tokens/s/device: 62290.894, total_weights: 65536, loss: 10.560 I0423 18:41:20.086971 137979238270784 metric_logger.py:185] completed step: 3, seconds: 0.419, TFLOP/s/device: 32.432, Tokens/s/device: 4888.493, total_weights: 65536, loss: 9.979 I0423 18:41:20.092824 137979238270784 checkpointing.py:794] Waiting for step 5 to finish before checkpoint... I0423 18:41:20.382068 137979238270784 checkpointing.py:798] Waited 0.28921937942504883 seconds for step 5 to finish before starting checkpointing. I0423 18:41:20.384639 137979238270784 checkpoint_manager.py:2024] [process=3][thread=MainThread][step=0][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0423 18:41:21.975830 2688 google_auth_provider.cc:181] Running on GCE, using service account 562977990677-compute@developer.gserviceaccount.com I0423 18:41:23.852852 137851941418752 array_metadata_store.py:203] [process=3][thread=array_type_handler] Wrote 69 array_metadata.ArrayMetadata to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/0/items/array_metadatas/process_3 I0423 18:41:29.654651 137851933026048 base_pytree_checkpoint_handler.py:129] [process=3] /jax/checkpoint/write/gbytes_per_sec: 141.718 MiB/s (total gbytes: 1.5 GiB) (time elapsed: 11.14602279663086 s) (per-host) I0423 18:41:29.654770 137851933026048 async_checkpointer.py:90] [process=3][thread=async_save] 3 Handler Commit operations completed. Time taken: 10.597331s. I0423 18:41:37.640173 137851933026048 async_checkpointer.py:144] [process=3][thread=async_save] Background save thread done. Time taken: 18.582719s. I0423 18:41:37.640546 137850867676928 async_checkpointer.py:273] [process=3][thread=save_finalize] Done with waiting for background save thread=async_save. I0423 18:41:37.640690 137850867676928 async_checkpointer.py:283] [process=3][thread=save_finalize] No errors found in background save thread=async_save. I0423 18:41:37.640758 137850867676928 checkpoint_manager.py:2133] [process=3][thread=save_finalize][step=0] CheckpointManager Save Finalize is syncing with other hosts... I0423 18:41:37.642474 137850867676928 checkpoint_manager.py:2142] [process=3][thread=save_finalize][step=0] CheckpointManager Save Finalize is done on all hosts. I0423 18:41:37.642668 137979238270784 checkpoint_manager.py:2036] [process=3][thread=MainThread][step=0][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=0. W0423 18:41:37.642812 137979238270784 checkpoint_manager.py:1458] Waiting for previous save to complete took 17.258177 seconds. If this number is high, consider checkpointing less frequently. I0423 18:41:37.644967 137979238270784 checkpoint_manager.py:1518] [process=3] Saving checkpoint at step 5 I0423 18:41:37.649626 137979238270784 async_checkpointer.py:452] [process=3] Started async saving checkpoint to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/5. I0423 18:41:38.553965 137979238270784 jax_array_handlers.py:358] Scheduling D2H of 69 prioritized jax.Array. I0423 18:41:38.554131 137979238270784 replica_slices.py:424] Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=False I0423 18:41:38.595332 137979238270784 base_pytree_checkpoint_handler.py:153] [process=3][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.043465s I0423 18:41:38.595507 137979238270784 base_pytree_checkpoint_handler.py:129] [process=3] /jax/checkpoint/write/blocking_gbytes_per_sec: 5.031 GiB/s (total gbytes: 1.5 GiB) (time elapsed: 0.30660367012023926 s) (per-host) I0423 18:41:38.595564 137979238270784 base_pytree_checkpoint_handler.py:737] [process=3][thread=MainThread] Initiated Pytree async_save. Time taken: 0.306668s (batch_requests_ready=0.258387s, total_serialization_initiated=0.048210s, others=0.000071s) I0423 18:41:38.595664 137979238270784 composite_checkpoint_handler.py:715] [process=3][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.311523s (all_items=0.000015s, per_item={'items': '0.00001526'}, temp_paths=0.311508) I0423 18:41:38.596557 137845465396992 async_checkpointer.py:79] [process=3][thread=async_save] Background save thread started. I0423 18:41:38.596715 137979238270784 async_checkpointer.py:561] Finished blocking save. Time taken: 0.951677s. Continuing background save to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/5. I0423 18:41:38.602705 137979238270784 checkpoint_manager.py:1566] [process=3][thread=MainThread][step=5] Starting CheckpointManager Save Finalize thread=save_finalize I0423 18:41:38.602993 137850867676928 async_checkpointer.py:265] [process=3][thread=save_finalize] Waiting for background save thread=async_save. I0423 18:41:38.603155 137979238270784 standard_logger.py:34] {'step': 5, 'event_type': 'save', 'directory': 'gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776969680.3846107, 'wait_for_prev_duration_secs': 17.258176565170288, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776969697.645006, 'checkpointer_blocking_duration_secs': 0.9518222808837891, 'get_old_steps_start_time': 1776969698.5968478, 'get_old_steps_duration_secs': 2.574920654296875e-05, 'checkpoint_manager_blocking_start_time': 1776969680.3824055, 'checkpoint_manager_blocking_duration_secs': 18.22071623802185} I0423 18:41:38.603264 137979238270784 checkpointing.py:409] Started an asynchronous checkpoint save for step 5 I0423 18:41:38.604015 137979238270784 metric_logger.py:185] completed step: 4, seconds: 0.163, TFLOP/s/device: 83.516, Tokens/s/device: 12588.513, total_weights: 65536, loss: 9.457 I0423 18:41:38.611497 137979238270784 metric_logger.py:185] completed step: 5, seconds: 0.147, TFLOP/s/device: 92.421, Tokens/s/device: 13930.646, total_weights: 65536, loss: 8.955 I0423 18:41:58.305206 137979238270784 metric_logger.py:185] completed step: 6, seconds: 18.517, TFLOP/s/device: 0.734, Tokens/s/device: 110.602, total_weights: 65536, loss: 8.465 I0423 18:41:58.452564 137979238270784 metric_logger.py:185] completed step: 7, seconds: 19.549, TFLOP/s/device: 0.695, Tokens/s/device: 104.765, total_weights: 65536, loss: 7.998 I0423 18:41:58.599986 137979238270784 metric_logger.py:185] completed step: 8, seconds: 0.153, TFLOP/s/device: 88.536, Tokens/s/device: 13345.062, total_weights: 65536, loss: 7.567 I0423 18:41:58.748373 137979238270784 checkpointing.py:794] Waiting for step 9 to finish before checkpoint... I0423 18:41:58.749042 137979238270784 checkpointing.py:798] Waited 0.0006804466247558594 seconds for step 9 to finish before starting checkpointing. I0423 18:41:58.751313 137979238270784 checkpoint_manager.py:2024] [process=3][thread=MainThread][step=5][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0423 18:42:03.059075 137851941418752 array_metadata_store.py:203] [process=3][thread=array_type_handler] Wrote 69 array_metadata.ArrayMetadata to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/5/items/array_metadatas/process_3 I0423 18:42:09.265310 137845465396992 base_pytree_checkpoint_handler.py:129] [process=3] /jax/checkpoint/write/gbytes_per_sec: 50.994 MiB/s (total gbytes: 1.5 GiB) (time elapsed: 30.976366996765137 s) (per-host) I0423 18:42:09.265435 137845465396992 async_checkpointer.py:90] [process=3][thread=async_save] 3 Handler Commit operations completed. Time taken: 30.668767s. I0423 18:42:15.778262 137845465396992 async_checkpointer.py:144] [process=3][thread=async_save] Background save thread done. Time taken: 37.181579s. I0423 18:42:15.778507 137850867676928 async_checkpointer.py:273] [process=3][thread=save_finalize] Done with waiting for background save thread=async_save. I0423 18:42:15.778575 137850867676928 async_checkpointer.py:283] [process=3][thread=save_finalize] No errors found in background save thread=async_save. I0423 18:42:15.778624 137850867676928 checkpoint_manager.py:2133] [process=3][thread=save_finalize][step=5] CheckpointManager Save Finalize is syncing with other hosts... I0423 18:42:15.780150 137850867676928 checkpoint_manager.py:2142] [process=3][thread=save_finalize][step=5] CheckpointManager Save Finalize is done on all hosts. I0423 18:42:15.780263 137979238270784 checkpoint_manager.py:2036] [process=3][thread=MainThread][step=5][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=5. W0423 18:42:15.780324 137979238270784 checkpoint_manager.py:1458] Waiting for previous save to complete took 17.029026 seconds. If this number is high, consider checkpointing less frequently. I0423 18:42:15.782181 137979238270784 checkpoint_manager.py:1518] [process=3] Saving checkpoint at step 9 I0423 18:42:15.785757 137979238270784 async_checkpointer.py:452] [process=3] Started async saving checkpoint to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/9. I0423 18:42:16.772353 137979238270784 jax_array_handlers.py:358] Scheduling D2H of 69 prioritized jax.Array. I0423 18:42:16.772463 137979238270784 replica_slices.py:424] Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=False I0423 18:42:16.816678 137979238270784 base_pytree_checkpoint_handler.py:153] [process=3][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.046336s I0423 18:42:16.816864 137979238270784 base_pytree_checkpoint_handler.py:129] [process=3] /jax/checkpoint/write/blocking_gbytes_per_sec: 4.765 GiB/s (total gbytes: 1.5 GiB) (time elapsed: 0.3236968517303467 s) (per-host) I0423 18:42:16.816940 137979238270784 base_pytree_checkpoint_handler.py:737] [process=3][thread=MainThread] Initiated Pytree async_save. Time taken: 0.323769s (batch_requests_ready=0.272892s, total_serialization_initiated=0.050797s, others=0.000080s) I0423 18:42:16.817075 137979238270784 composite_checkpoint_handler.py:715] [process=3][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.327891s (all_items=0.000010s, per_item={'items': '0.00000954'}, temp_paths=0.327881) I0423 18:42:16.817979 137846555924224 async_checkpointer.py:79] [process=3][thread=async_save] Background save thread started. I0423 18:42:16.818135 137979238270784 async_checkpointer.py:561] Finished blocking save. Time taken: 1.035887s. Continuing background save to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/9. I0423 18:42:17.240811 137979238270784 checkpoint_manager.py:1566] [process=3][thread=MainThread][step=9] Starting CheckpointManager Save Finalize thread=save_finalize I0423 18:42:17.241203 137850867676928 async_checkpointer.py:265] [process=3][thread=save_finalize] Waiting for background save thread=async_save. I0423 18:42:17.241372 137979238270784 standard_logger.py:34] {'step': 9, 'event_type': 'save', 'directory': 'gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776969718.7512841, 'wait_for_prev_duration_secs': 17.02902579307556, 'time_between_consecutive_saves_sec': 21.108762979507446, 'checkpointer_blocking_start_time': 1776969735.7822187, 'checkpointer_blocking_duration_secs': 1.0360262393951416, 'get_old_steps_start_time': 1776969736.818257, 'get_old_steps_duration_secs': 2.0503997802734375e-05, 'checkpoint_manager_blocking_start_time': 1776969718.749261, 'checkpoint_manager_blocking_duration_secs': 18.492078065872192} I0423 18:42:17.241503 137979238270784 checkpointing.py:409] Started an asynchronous checkpoint save for step 9 I0423 18:42:17.241547 137979238270784 checkpoint_manager.py:2024] [process=3][thread=MainThread][step=9][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0423 18:42:21.876970 137851941418752 array_metadata_store.py:203] [process=3][thread=array_type_handler] Wrote 69 array_metadata.ArrayMetadata to gs://lance-maxtext/nnx_ckpt_xpk_feat_nnx_set_defaults_true_20260423_155242/nnx_xpk_feat_nnx_set_defaults_true_20260423_155242_15_ocdbt_false/checkpoints/9/items/array_metadatas/process_3 I0423 18:42:27.424042 137846555924224 base_pytree_checkpoint_handler.py:129] [process=3] /jax/checkpoint/write/gbytes_per_sec: 144.508 MiB/s (total gbytes: 1.5 GiB) (time elapsed: 10.930845499038696 s) (per-host) I0423 18:42:27.424144 137846555924224 async_checkpointer.py:90] [process=3][thread=async_save] 3 Handler Commit operations completed. Time taken: 10.606054s. I0423 18:42:32.131460 137846555924224 async_checkpointer.py:144] [process=3][thread=async_save] Background save thread done. Time taken: 15.313355s. I0423 18:42:32.131783 137850867676928 async_checkpointer.py:273] [process=3][thread=save_finalize] Done with waiting for background save thread=async_save. I0423 18:42:32.131906 137850867676928 async_checkpointer.py:283] [process=3][thread=save_finalize] No errors found in background save thread=async_save. I0423 18:42:32.131965 137850867676928 checkpoint_manager.py:2133] [process=3][thread=save_finalize][step=9] CheckpointManager Save Finalize is syncing with other hosts... I0423 18:42:32.133584 137850867676928 checkpoint_manager.py:2142] [process=3][thread=save_finalize][step=9] CheckpointManager Save Finalize is done on all hosts. I0423 18:42:32.133768 137979238270784 checkpoint_manager.py:2036] [process=3][thread=MainThread][step=9][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=9. I0423 18:42:32.133960 137979238270784 checkpoint_manager.py:2013] [process=3][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0423 18:42:32.134704 137979238270784 metric_logger.py:185] completed step: 9, seconds: 0.147, TFLOP/s/device: 92.482, Tokens/s/device: 13939.938, total_weights: 65536, loss: 7.176 Per train step: Total TFLOPs: 13.59 split as 93.93% learnable weight flops and 6.07% attention flops XPK End: Thu Apr 23 18:42:42 UTC 2026 EXIT_CODE=0