main (574ad3fb9) vs feat/Migrate-Decoder-And-Tests-to-NNX (574ad3fb9)| Metric | main 574ad3fb9 | feat/Migrate-Decoder-And-Tests-to-NNX 574ad3fb9 | Diff (feat/Migrate-Decoder-And-Tests-to-NNX − main) |
|---|---|---|---|
| Parameters | 1.104 billion | 1.104 billion | — |
| Final loss | 8.1790 | 8.1790 | 0 |
| TFLOP/s | 92.234 | 92.121 | -0.113 |
| Tok/s | 13902.5 | 13885.4 | -17.061 |
| Avg s/step | 3.071 | 2.426 | -0.645 |
| Memory % | 1.38 | 1.38 | 0 |
| JAX | 0.9.2 | 0.8.1 | — |
Diff = branch value − main value. Green = branch improved. Red = branch regressed.
XPK Start: Sat Apr 18 18:01:13 UTC 2026 PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used. 2026-04-18 18:01:38.448324: 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) I0418 18:01:38.627807 136016632297280 max_utils.py:273] Attempting to initialize the jax distributed system... I0418 18:01:47.667797 136016632297280 distributed.py:149] Starting JAX distributed service on [::]:8482 I0418 18:01:47.670194 136016632297280 distributed.py:172] Connecting to JAX distributed service on mt-01-base-2sag8-slice-job-0-0.mt-01-base-2sag8:8482 I0418 18:01:49.137647 136016632297280 max_utils.py:284] Jax distributed system initialized! I0418 18:01:56.197314 136016632297280 max_utils.py:800] System Information: Jax Version: 0.9.2 I0418 18:01:56.197413 136016632297280 max_utils.py:801] System Information: Jaxlib Version: 0.9.2 I0418 18:01:56.197469 136016632297280 max_utils.py:802] System Information: Jax Backend: PJRT C API TFRT TPU v6 lite Built on Mar 4 2026 11:32:08 (1772652728) cl/878335365 I0418 18:01:56.197514 136016632297280 train_utils.py:348] WARNING: Sequence packing is essentially ignored for synthetic data. Please use a real dataset to use sequence packing. I0418 18:01:56.911277 136016632297280 maxtext_utils.py:1551] Num_devices: 32, shape (1, 1, 1, 32, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0418 18:01:56.911563 136016632297280 checkpointing.py:677] Setting up checkpoint logger... I0418 18:01:56.911619 136016632297280 checkpointing.py:233] Creating checkpoint manager with ocdbt=True and zarr3=True I0418 18:01:56.911664 136016632297280 pytree_checkpoint_handler.py:592] save_device_host_concurrent_bytes=None I0418 18:01:56.912001 136016632297280 base_pytree_checkpoint_handler.py:441] Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=True, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7bb436918260>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0418 18:02:01.482331 136016632297280 checkpointing.py:265] Enabling policy for fixed interval checkpointing. I0418 18:02:01.482582 136016632297280 checkpoint_manager.py:708] [process=5][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=('items',), item_handlers={'items': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7b9f743869c0>}, handler_registry=None I0418 18:02:01.482824 136016632297280 composite_checkpoint_handler.py:237] Deferred registration for item: "items". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7b9f743869c0>` 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`. I0418 18:02:01.482872 136016632297280 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7b9f74388290>` 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`. I0418 18:02:01.482908 136016632297280 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 0x7b9f743869c0>, ('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7b9f743869c0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7b9f74388290>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7b9f74388290>}). I0418 18:02:01.483229 136016632297280 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.34 I0418 18:02:01.483299 136016632297280 async_checkpointer.py:192] [process=5][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7b9f546afba0> timeout: 1200 secs and primary_host=0 for async checkpoint writes I0418 18:02:02.509403 136016632297280 checkpoint_manager.py:1812] Found 0 checkpoint steps in gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints I0418 18:02:03.389142 136016632297280 checkpoint_manager.py:929] [process=5][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=10), 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/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7bb315bdffb0> I0418 18:02:03.389307 136016632297280 checkpointing.py:301] Checkpoint manager created! I0418 18:02:04.318331 136016632297280 nnx_wrappers.py:437] Unknown Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0418 18:02:04.318439 136016632297280 nnx_wrappers.py:437] Unknown Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0418 18:02:04.697327 136016632297280 attentions.py:1088] attentions/inputs_q Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0418 18:02:04.697414 136016632297280 attentions.py:1088] attentions/inputs_q Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0418 18:02:04.713699 136016632297280 attentions.py:1089] attentions/inputs_kv Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0418 18:02:04.713760 136016632297280 attentions.py:1089] attentions/inputs_kv Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0418 18:02:04.737531 136016632297280 attentions.py:1154] attentions/query Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0418 18:02:04.737595 136016632297280 attentions.py:1154] attentions/query Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0418 18:02:04.753716 136016632297280 attentions.py:1155] attentions/key Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0418 18:02:04.753776 136016632297280 attentions.py:1155] attentions/key Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0418 18:02:04.769909 136016632297280 attentions.py:1156] attentions/value Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0418 18:02:04.769973 136016632297280 attentions.py:1156] attentions/value Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0418 18:02:04.794766 136016632297280 attentions.py:1197] attentions/out Logical: bfloat16[32,2048,16,128].................................... ('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv'). I0418 18:02:04.794831 136016632297280 attentions.py:1197] attentions/out Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0418 18:02:04.815628 136016632297280 linears.py:525] linears/x Logical: bfloat16[32,2048,7168]...................................... ('activation_batch', 'activation_length', 'activation_mlp'). I0418 18:02:04.815692 136016632297280 linears.py:525] linears/x Physical: bfloat16[32,2048,7168]...................................... ('fsdp', None, None). I0418 18:02:05.031522 136016632297280 checkpointing.py:577] checkpoint manager exists so trying to load this run's existing checkpoint I0418 18:02:05.031626 136016632297280 checkpointing.py:665] No existing checkpoints found, not restoring checkpoint. fsdp: 32 I0418 18:02:06.457116 136016632297280 maxtext_utils.py:1654] params/params/decoder/decoder_norm/scale Shape: float32[2048] Logical: P('norm',) Physical: (None,) I0418 18:02:06.457235 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/mlp/wi_0/kernel Shape: float32[2048,16,7168] Logical: P('embed', 'layers', 'mlp') Physical: ('fsdp', None, None) I0418 18:02:06.457286 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/mlp/wi_1/kernel Shape: float32[2048,16,7168] Logical: P('embed', 'layers', 'mlp') Physical: ('fsdp', None, None) I0418 18:02:06.457343 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/mlp/wo/kernel Shape: float32[7168,16,2048] Logical: P('mlp', 'layers', 'embed') Physical: (None, None, 'fsdp') I0418 18:02:06.457395 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/post_self_attention_layer_norm/scale Shape: float32[2048,16] Logical: P('norm', 'layers') Physical: (None, None) I0418 18:02:06.457432 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/pre_self_attention_layer_norm/scale Shape: float32[2048,16] Logical: P('norm', 'layers') Physical: (None, None) I0418 18:02:06.457497 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/self_attention/key/kernel Shape: float32[2048,16,16,128] Logical: P('embed', 'layers', 'kv_heads', 'kv_head_dim') Physical: ('fsdp', None, None, None) I0418 18:02:06.457551 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/self_attention/out/kernel Shape: float32[16,16,128,2048] Logical: P('heads', 'layers', 'kv', 'embed') Physical: (None, None, None, 'fsdp') I0418 18:02:06.457590 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/self_attention/query/kernel Shape: float32[2048,16,16,128] Logical: P('embed', 'layers', 'q_heads', 'kv') Physical: ('fsdp', None, None, None) I0418 18:02:06.457626 136016632297280 maxtext_utils.py:1654] params/params/decoder/layers/self_attention/value/kernel Shape: float32[2048,16,16,128] Logical: P('embed', 'layers', 'kv_heads', 'kv_head_dim') Physical: ('fsdp', None, None, None) I0418 18:02:06.457673 136016632297280 maxtext_utils.py:1654] params/params/decoder/logits_dense/kernel Shape: float32[2048,32000] Logical: P('embed_vocab', 'vocab') Physical: ('fsdp', None) I0418 18:02:06.457719 136016632297280 maxtext_utils.py:1654] params/params/token_embedder/embedding Shape: float32[32000,2048] Logical: P('vocab', 'embed_vocab') Physical: (None, 'fsdp') I0418 18:02:06.945318 136016632297280 train.py:155] train/xent Logical: float32[32,2048]............................................ ('activation_embed_and_logits_batch', 'activation_length'). I0418 18:02:06.945410 136016632297280 train.py:155] train/xent Physical: float32[32,2048]............................................ ('fsdp', None). I0418 18:02:06.960833 136016632297280 train.py:162] train/z_loss Logical: float32[32,2048]............................................ ('activation_embed_and_logits_batch', 'activation_length'). I0418 18:02:06.960896 136016632297280 train.py:162] train/z_loss Physical: float32[32,2048]............................................ ('fsdp', None). I0418 18:02:17.925596 136016632297280 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. I0418 18:02:17.926374 136016632297280 metric_logger.py:301] number parameters: 1.104 billion I0418 18:02:28.987720 136016632297280 checkpointing.py:772] Waiting for step 0 to finish before checkpoint... I0418 18:02:29.319164 136016632297280 checkpointing.py:776] Waited 0.3314321041107178 seconds for step 0 to finish before starting checkpointing. I0418 18:02:29.321613 136016632297280 checkpoint_manager.py:2009] [process=5][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0418 18:02:29.323310 136016632297280 checkpoint_manager.py:1512] [process=5] Saving checkpoint at step 0 I0418 18:02:29.324758 136016632297280 event_tracking.py:70] [process=5] [async] Started save checkpoint @ gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints/0. I0418 18:02:30.509867 136016632297280 signaling_client.py:364] Using JaxDistributedSignalingClient I0418 18:02:30.510792 136016632297280 jax_array_handlers.py:360] Scheduling D2H of 39 prioritized jax.Array. I0418 18:02:30.510847 136016632297280 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 I0418 18:02:30.788037 136016632297280 base_pytree_checkpoint_handler.py:154] [process=5][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.278241s I0418 18:02:30.788214 136016632297280 base_pytree_checkpoint_handler.py:130] [process=5] /jax/orbax/write/blocking_gbytes_per_sec: 5.429 GiB/s (total gbytes: 1.5 GiB) (time elapsed: 0.28411173820495605 s) (per-host) I0418 18:02:30.788267 136016632297280 base_pytree_checkpoint_handler.py:768] [process=5][thread=MainThread] Initiated Pytree async_save. Time taken: 0.284174s (batch_requests_ready=0.002627s, total_serialization_initiated=0.281476s, others=0.000071s) I0418 18:02:30.788356 136016632297280 composite_checkpoint_handler.py:715] [process=5][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.288152s (all_items=0.000017s, per_item={'items': '0.00001693'}, temp_paths=0.288136) I0418 18:02:30.789228 136016632297280 event_tracking.py:125] [process=5] [async] Finished blocking save in 1.47 seconds. Continuing save @ gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints/0. I0418 18:02:30.789592 135888973575936 async_checkpointer.py:76] [process=5][thread=async_save] Background save thread started. Deadline for this save operation is 2026-04-18 18:22:30.789554 I0418 18:02:30.796347 136016632297280 checkpoint_manager.py:1560] [process=5][thread=MainThread][step=0] Starting CheckpointManager Save Finalize thread=save_finalize I0418 18:02:30.796626 135887344273152 async_checkpointer.py:280] [process=5][thread=save_finalize] Waiting for background save thread=async_save. I0418 18:02:30.796753 136016632297280 standard_logger.py:34] {'step': 0, 'event_type': 'save', 'directory': 'gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776535349.3215866, 'wait_for_prev_duration_secs': 8.726119995117188e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776535349.3233628, 'checkpointer_blocking_duration_secs': 1.466383457183838, 'get_old_steps_start_time': 1776535350.7897708, 'get_old_steps_duration_secs': 3.0517578125e-05, 'checkpoint_manager_blocking_start_time': 1776535349.319707, 'checkpoint_manager_blocking_duration_secs': 1.4770071506500244} I0418 18:02:30.796863 136016632297280 checkpointing.py:408] Started an asynchronous checkpoint save for step 0 I0418 18:02:30.796917 136016632297280 max_utils.py:750] Memstats: After params initialized: I0418 18:02:30.796967 136016632297280 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_18(process=5,(2,4,0,0)) I0418 18:02:30.797000 136016632297280 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_19(process=5,(3,4,0,0)) I0418 18:02:30.797027 136016632297280 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_22(process=5,(2,5,0,0)) I0418 18:02:30.797050 136016632297280 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_23(process=5,(3,5,0,0)) I0418 18:02:31.105650 136016632297280 metric_logger.py:196] completed step: 0, seconds: 11.061, TFLOP/s/device: 1.228, Tokens/s/device: 185.152, total_weights: 65536, loss: 10.874, lm_loss: 10.874, perplexity: 52779.875 I0418 18:02:31.296370 136016632297280 metric_logger.py:196] completed step: 1, seconds: 2.116, TFLOP/s/device: 6.420, Tokens/s/device: 967.643, total_weights: 65536, loss: 10.874, lm_loss: 10.874, perplexity: 52779.875 I0418 18:02:31.692446 136016632297280 metric_logger.py:196] completed step: 2, seconds: 0.044, TFLOP/s/device: 312.262, Tokens/s/device: 47067.476, total_weights: 65536, loss: 10.262, lm_loss: 10.262, perplexity: 28637.129 I0418 18:02:31.839858 136016632297280 metric_logger.py:196] completed step: 3, seconds: 0.397, TFLOP/s/device: 34.233, Tokens/s/device: 5159.990, total_weights: 65536, loss: 9.734, lm_loss: 9.734, perplexity: 16889.203 I0418 18:02:32.135124 136016632297280 metric_logger.py:196] completed step: 4, seconds: 0.152, TFLOP/s/device: 89.128, Tokens/s/device: 13434.353, total_weights: 65536, loss: 9.277, lm_loss: 9.277, perplexity: 10694.614 I0418 18:02:32.140931 136016632297280 metric_logger.py:196] completed step: 5, seconds: 0.147, TFLOP/s/device: 92.207, Tokens/s/device: 13898.409, total_weights: 65536, loss: 8.892, lm_loss: 8.892, perplexity: 7270.717 I0418 18:02:36.769223 2828 google_auth_provider.cc:181] Running on GCE, using service account 562977990677-compute@developer.gserviceaccount.com I0418 18:02:40.319857 135887393478400 array_metadata_store.py:203] [process=5][thread=array_type_handler] Wrote 39 array_metadata.ArrayMetadata to gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints/0/items/array_metadatas/process_5 I0418 18:02:56.479434 136016632297280 metric_logger.py:196] completed step: 6, seconds: 0.296, TFLOP/s/device: 45.968, Tokens/s/device: 6928.844, total_weights: 65536, loss: 8.592, lm_loss: 8.592, perplexity: 5390.759 I0418 18:02:56.626775 136016632297280 metric_logger.py:196] completed step: 7, seconds: 24.192, TFLOP/s/device: 0.562, Tokens/s/device: 84.658, total_weights: 65536, loss: 8.384, lm_loss: 8.384, perplexity: 4376.905 I0418 18:02:56.774048 136016632297280 metric_logger.py:196] completed step: 8, seconds: 0.152, TFLOP/s/device: 89.207, Tokens/s/device: 13446.261, total_weights: 65536, loss: 8.255, lm_loss: 8.255, perplexity: 3846.406 I0418 18:02:56.920902 136016632297280 checkpointing.py:772] Waiting for step 9 to finish before checkpoint... I0418 18:02:56.921597 136016632297280 checkpointing.py:776] Waited 0.0007135868072509766 seconds for step 9 to finish before starting checkpointing. I0418 18:02:56.923789 136016632297280 checkpoint_manager.py:2020] [process=5][thread=MainThread][step=0][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0418 18:03:09.217653 135888973575936 base_pytree_checkpoint_handler.py:130] [process=5] /jax/orbax/write/gbytes_per_sec: 40.802 MiB/s (total gbytes: 1.5 GiB) (time elapsed: 38.713520526885986 s) (per-host) I0418 18:03:09.217759 135888973575936 async_checkpointer.py:90] [process=5][thread=async_save] 3 Handler Commit operations completed. Time taken: 38.428054s. I0418 18:03:21.081743 135888973575936 async_checkpointer.py:160] [process=5][thread=async_save] Background save thread done. Time taken: 50.292021s. I0418 18:03:21.082024 135887344273152 async_checkpointer.py:288] [process=5][thread=save_finalize] Done with waiting for background save thread=async_save. I0418 18:03:21.082161 135887344273152 async_checkpointer.py:298] [process=5][thread=save_finalize] No errors found in background save thread=async_save. I0418 18:03:21.082236 135887344273152 checkpoint_manager.py:2137] [process=5][thread=save_finalize][step=0] CheckpointManager Save Finalize is syncing with other hosts... I0418 18:03:21.084198 135887344273152 checkpoint_manager.py:2146] [process=5][thread=save_finalize][step=0] CheckpointManager Save Finalize is done on all hosts. I0418 18:03:21.084340 136016632297280 checkpoint_manager.py:2032] [process=5][thread=MainThread][step=0][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=0. W0418 18:03:21.084431 136016632297280 checkpoint_manager.py:1452] Waiting for previous save to complete took 24.160645 seconds. If this number is high, consider checkpointing less frequently. I0418 18:03:21.086312 136016632297280 checkpoint_manager.py:1512] [process=5] Saving checkpoint at step 9 I0418 18:03:21.088404 136016632297280 event_tracking.py:70] [process=5] [async] Started save checkpoint @ gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints/9. I0418 18:03:21.805513 136016632297280 jax_array_handlers.py:360] Scheduling D2H of 39 prioritized jax.Array. I0418 18:03:21.805606 136016632297280 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 I0418 18:03:21.851402 136016632297280 base_pytree_checkpoint_handler.py:154] [process=5][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.047068s I0418 18:03:21.851588 136016632297280 base_pytree_checkpoint_handler.py:130] [process=5] /jax/orbax/write/blocking_gbytes_per_sec: 30.451 GiB/s (total gbytes: 1.5 GiB) (time elapsed: 0.05065727233886719 s) (per-host) I0418 18:03:21.851641 136016632297280 base_pytree_checkpoint_handler.py:768] [process=5][thread=MainThread] Initiated Pytree async_save. Time taken: 0.050721s (batch_requests_ready=0.001908s, total_serialization_initiated=0.048741s, others=0.000072s) I0418 18:03:21.851732 136016632297280 composite_checkpoint_handler.py:715] [process=5][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.054832s (all_items=0.000016s, per_item={'items': '0.00001574'}, temp_paths=0.054816) I0418 18:03:21.852479 136016632297280 event_tracking.py:125] [process=5] [async] Finished blocking save in 0.77 seconds. Continuing save @ gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints/9. I0418 18:03:21.852841 135887344273152 async_checkpointer.py:76] [process=5][thread=async_save] Background save thread started. Deadline for this save operation is 2026-04-18 18:23:21.852802 I0418 18:03:21.854827 136016632297280 checkpoint_manager.py:1560] [process=5][thread=MainThread][step=9] Starting CheckpointManager Save Finalize thread=save_finalize I0418 18:03:21.855110 135887393478400 async_checkpointer.py:280] [process=5][thread=save_finalize] Waiting for background save thread=async_save. I0418 18:03:21.855244 136016632297280 standard_logger.py:34] {'step': 9, 'event_type': 'save', 'directory': 'gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776535376.9237592, 'wait_for_prev_duration_secs': 24.160645008087158, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776535401.086351, 'checkpointer_blocking_duration_secs': 0.766655445098877, 'get_old_steps_start_time': 1776535401.8530304, 'get_old_steps_duration_secs': 3.147125244140625e-05, 'checkpoint_manager_blocking_start_time': 1776535376.9218338, 'checkpoint_manager_blocking_duration_secs': 24.933377981185913} I0418 18:03:21.855412 136016632297280 checkpointing.py:408] Started an asynchronous checkpoint save for step 9 I0418 18:03:21.855471 136016632297280 checkpoint_manager.py:2020] [process=5][thread=MainThread][step=9][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0418 18:03:29.938111 135865812625152 array_metadata_store.py:203] [process=5][thread=array_type_handler] Wrote 39 array_metadata.ArrayMetadata to gs://lance-maxtext/linen_ckpt_xpk_main_20260418_180002/linen_xpk_main_20260418_180002_01_base/checkpoints/9/items/array_metadatas/process_5 I0418 18:04:05.668946 135887344273152 base_pytree_checkpoint_handler.py:130] [process=5] /jax/orbax/write/gbytes_per_sec: 36.008 MiB/s (total gbytes: 1.5 GiB) (time elapsed: 43.86797738075256 s) (per-host) I0418 18:04:05.669079 135887344273152 async_checkpointer.py:90] [process=5][thread=async_save] 3 Handler Commit operations completed. Time taken: 43.816113s. I0418 18:04:14.249446 135887344273152 async_checkpointer.py:160] [process=5][thread=async_save] Background save thread done. Time taken: 52.396466s. I0418 18:04:14.249764 135887393478400 async_checkpointer.py:288] [process=5][thread=save_finalize] Done with waiting for background save thread=async_save. I0418 18:04:14.249889 135887393478400 async_checkpointer.py:298] [process=5][thread=save_finalize] No errors found in background save thread=async_save. I0418 18:04:14.249945 135887393478400 checkpoint_manager.py:2137] [process=5][thread=save_finalize][step=9] CheckpointManager Save Finalize is syncing with other hosts... I0418 18:04:14.252395 135887393478400 checkpoint_manager.py:2146] [process=5][thread=save_finalize][step=9] CheckpointManager Save Finalize is done on all hosts. I0418 18:04:14.252585 136016632297280 checkpoint_manager.py:2032] [process=5][thread=MainThread][step=9][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=9. I0418 18:04:14.252735 136016632297280 checkpoint_manager.py:2009] [process=5][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0418 18:04:14.253656 136016632297280 metric_logger.py:196] completed step: 9, seconds: 0.147, TFLOP/s/device: 92.234, Tokens/s/device: 13902.466, total_weights: 65536, loss: 8.179, lm_loss: 8.179, perplexity: 3564.635 Per train step: Total TFLOPs: 13.59 split as 93.93% learnable weight flops and 6.07% attention flops XPK End: Sat Apr 18 18:04:23 UTC 2026 EXIT_CODE=0
574ad3fb9
· feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615
· full log
XPK Start: Sun Apr 19 06:37:21 UTC 2026 2026-04-19 06:37:25.391802: 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:1776580645.404689 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:1776580645.408495 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:1776580645.419672 10 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776580645.419692 10 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776580645.419695 10 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776580645.419697 10 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. 2026-04-19 06:37:45.146532: 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) I0419 06:37:45.660595 134452881033024 max_utils.py:273] Attempting to initialize the jax distributed system... INFO:2026-04-19 06:37:54,702:jax._src.distributed:140: Starting JAX distributed service on [::]:8482 I0419 06:37:54.702025 134452881033024 distributed.py:140] Starting JAX distributed service on [::]:8482 INFO:2026-04-19 06:37:54,704:jax._src.distributed:157: Connecting to JAX distributed service on mt-01-base-usj7h-slice-job-0-0.mt-01-base-usj7h:8482 I0419 06:37:54.704313 134452881033024 distributed.py:157] Connecting to JAX distributed service on mt-01-base-usj7h-slice-job-0-0.mt-01-base-usj7h:8482 I0419 06:37:58.503653 134452881033024 max_utils.py:284] Jax distributed system initialized! I0419 06:38:05.646684 134452881033024 max_utils.py:800] System Information: Jax Version: 0.8.1 I0419 06:38:05.646789 134452881033024 max_utils.py:801] System Information: Jaxlib Version: 0.8.1 I0419 06:38:05.646857 134452881033024 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 I0419 06:38:05.646896 134452881033024 train_utils.py:347] WARNING: Sequence packing is essentially ignored for synthetic data. Please use a real dataset to use sequence packing. I0419 06:38:06.263688 134452881033024 maxtext_utils.py:1517] Num_devices: 32, shape (1, 1, 1, 32, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0419 06:38:06.263991 134452881033024 checkpointing.py:677] Setting up checkpoint logger... I0419 06:38:06.264050 134452881033024 checkpointing.py:233] Creating checkpoint manager with ocdbt=True and zarr3=True I0419 06:38:06.264093 134452881033024 pytree_checkpoint_handler.py:589] save_device_host_concurrent_bytes=None I0419 06:38:06.264453 134452881033024 base_pytree_checkpoint_handler.py:415] Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=True, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7a483e575070>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0419 06:38:09.934162 134452881033024 checkpointing.py:265] Enabling policy for fixed interval checkpointing. I0419 06:38:09.934514 134452881033024 checkpoint_manager.py:709] [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=('items',), item_handlers={'items': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a46faf71b80>}, handler_registry=None I0419 06:38:09.934752 134452881033024 composite_checkpoint_handler.py:237] Deferred registration for item: "items". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a46faf71b80>` 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`. I0419 06:38:09.934801 134452881033024 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a33a830af00>` 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`. I0419 06:38:09.934851 134452881033024 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 0x7a46faf71b80>, ('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a46faf71b80>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a33a830af00>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a33a830af00>}). I0419 06:38:09.935346 134452881033024 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.33 I0419 06:38:09.935430 134452881033024 async_checkpointer.py:177] [process=0][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7a3390456ac0> timeout: 600 secs and primary_host=0 for async checkpoint writes I0419 06:38:11.966485 134452881033024 checkpoint_manager.py:1818] Found 0 checkpoint steps in gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints I0419 06:38:12.438722 134452881033024 checkpoint_manager.py:929] [process=0][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=10), 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/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7a33a830aa20> I0419 06:38:12.438916 134452881033024 checkpointing.py:301] Checkpoint manager created! I0419 06:38:13.320749 134452881033024 nnx_wrappers.py:437] Unknown Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0419 06:38:13.320869 134452881033024 nnx_wrappers.py:437] Unknown Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0419 06:38:13.690023 134452881033024 attentions.py:1088] attentions/inputs_q Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0419 06:38:13.690126 134452881033024 attentions.py:1088] attentions/inputs_q Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0419 06:38:13.705794 134452881033024 attentions.py:1089] attentions/inputs_kv Logical: bfloat16[32,2048,2048]...................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0419 06:38:13.705868 134452881033024 attentions.py:1089] attentions/inputs_kv Physical: bfloat16[32,2048,2048]...................................... ('fsdp', None, None). I0419 06:38:13.728729 134452881033024 attentions.py:1154] attentions/query Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0419 06:38:13.728800 134452881033024 attentions.py:1154] attentions/query Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0419 06:38:13.744619 134452881033024 attentions.py:1155] attentions/key Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0419 06:38:13.744679 134452881033024 attentions.py:1155] attentions/key Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0419 06:38:13.760378 134452881033024 attentions.py:1156] attentions/value Logical: bfloat16[32,2048,16,128].................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0419 06:38:13.760443 134452881033024 attentions.py:1156] attentions/value Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0419 06:38:13.785166 134452881033024 attentions.py:1197] attentions/out Logical: bfloat16[32,2048,16,128].................................... ('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv'). I0419 06:38:13.785243 134452881033024 attentions.py:1197] attentions/out Physical: bfloat16[32,2048,16,128].................................... ('fsdp', None, None, None). I0419 06:38:13.804666 134452881033024 linears.py:525] linears/x Logical: bfloat16[32,2048,7168]...................................... ('activation_batch', 'activation_length', 'activation_mlp'). I0419 06:38:13.804734 134452881033024 linears.py:525] linears/x Physical: bfloat16[32,2048,7168]...................................... ('fsdp', None, None). I0419 06:38:14.006085 134452881033024 checkpointing.py:577] checkpoint manager exists so trying to load this run's existing checkpoint I0419 06:38:14.006200 134452881033024 checkpointing.py:665] No existing checkpoints found, not restoring checkpoint. fsdp: 32 I0419 06:38:15.427833 134452881033024 maxtext_utils.py:1620] params/params/decoder/decoder_norm/scale Shape: float32[2048] Logical: PartitionSpec('norm',) Physical: (None,) I0419 06:38:15.427967 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/mlp/wi_0/kernel Shape: float32[2048,16,7168] Logical: PartitionSpec('embed', 'layers', 'mlp') Physical: ('fsdp', None, None) I0419 06:38:15.428022 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/mlp/wi_1/kernel Shape: float32[2048,16,7168] Logical: PartitionSpec('embed', 'layers', 'mlp') Physical: ('fsdp', None, None) I0419 06:38:15.428081 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/mlp/wo/kernel Shape: float32[7168,16,2048] Logical: PartitionSpec('mlp', 'layers', 'embed') Physical: (None, None, 'fsdp') I0419 06:38:15.428151 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/post_self_attention_layer_norm/scale Shape: float32[2048,16] Logical: PartitionSpec('norm', 'layers') Physical: (None, None) I0419 06:38:15.428193 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/pre_self_attention_layer_norm/scale Shape: float32[2048,16] Logical: PartitionSpec('norm', 'layers') Physical: (None, None) I0419 06:38:15.428247 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/self_attention/key/kernel Shape: float32[2048,16,16,128] Logical: PartitionSpec('embed', 'layers', 'kv_heads', 'kv_head_dim') Physical: ('fsdp', None, None, None) I0419 06:38:15.428300 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/self_attention/out/kernel Shape: float32[16,16,128,2048] Logical: PartitionSpec('heads', 'layers', 'kv', 'embed') Physical: (None, None, None, 'fsdp') I0419 06:38:15.428338 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/self_attention/query/kernel Shape: float32[2048,16,16,128] Logical: PartitionSpec('embed', 'layers', 'q_heads', 'kv') Physical: ('fsdp', None, None, None) I0419 06:38:15.428373 134452881033024 maxtext_utils.py:1620] params/params/decoder/layers/self_attention/value/kernel Shape: float32[2048,16,16,128] Logical: PartitionSpec('embed', 'layers', 'kv_heads', 'kv_head_dim') Physical: ('fsdp', None, None, None) I0419 06:38:15.428419 134452881033024 maxtext_utils.py:1620] params/params/decoder/logits_dense/kernel Shape: float32[2048,32000] Logical: PartitionSpec('embed', 'vocab') Physical: ('fsdp', None) I0419 06:38:15.428473 134452881033024 maxtext_utils.py:1620] params/params/token_embedder/embedding Shape: float32[32000,2048] Logical: PartitionSpec('vocab', 'embed') Physical: (None, 'fsdp') I0419 06:38:15.927445 134452881033024 train.py:155] train/xent Logical: float32[32,2048]............................................ ('activation_embed_and_logits_batch', 'activation_length'). I0419 06:38:15.927539 134452881033024 train.py:155] train/xent Physical: float32[32,2048]............................................ ('fsdp', None). I0419 06:38:15.942640 134452881033024 train.py:162] train/z_loss Logical: float32[32,2048]............................................ ('activation_embed_and_logits_batch', 'activation_length'). I0419 06:38:15.942701 134452881033024 train.py:162] train/z_loss Physical: float32[32,2048]............................................ ('fsdp', None). I0419 06:38:26.780308 134452881033024 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. I0419 06:38:26.786350 134452881033024 metric_logger.py:289] number parameters: 1.104 billion I0419 06:38:38.212936 134452881033024 checkpointing.py:772] Waiting for step 0 to finish before checkpoint... I0419 06:38:38.369526 134452881033024 checkpointing.py:776] Waited 0.1565718650817871 seconds for step 0 to finish before starting checkpointing. I0419 06:38:38.372123 134452881033024 checkpoint_manager.py:2013] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0419 06:38:38.373933 134452881033024 checkpoint_manager.py:1518] [process=0] Saving checkpoint at step 0 I0419 06:38:38.376710 134452881033024 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0. I0419 06:38:39.121659 134452881033024 signaling_client.py:364] Using JaxDistributedSignalingClient I0419 06:38:39.380861 134327182554880 atomicity.py:137] Creating tmp directory gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0 I0419 06:38:39.421802 134452881033024 jax_array_handlers.py:358] Scheduling D2H of 39 prioritized jax.Array. I0419 06:38:39.421911 134452881033024 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 I0419 06:38:39.664671 134452881033024 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.243950s I0419 06:38:39.665607 134452881033024 base_pytree_checkpoint_handler.py:129] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 2.875 GiB/s (total gbytes: 1.5 GiB) (time elapsed: 0.5365936756134033 s) (per-host) I0419 06:38:39.665683 134452881033024 base_pytree_checkpoint_handler.py:737] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.536688s (batch_requests_ready=0.288012s, total_serialization_initiated=0.247833s, others=0.000842s) I0419 06:38:39.665763 134452881033024 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.540687s (all_items=0.000019s, per_item={'items': '0.00001907'}, temp_paths=0.540668) I0419 06:38:39.666698 134327148984064 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0419 06:38:39.666867 134452881033024 async_checkpointer.py:561] Finished blocking save. Time taken: 1.292858s. Continuing background save to gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0. I0419 06:38:40.075743 134452881033024 checkpoint_manager.py:1566] [process=0][thread=MainThread][step=0] Starting CheckpointManager Save Finalize thread=save_finalize I0419 06:38:40.076134 134325563549440 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0419 06:38:40.076298 134452881033024 standard_logger.py:34] {'step': 0, 'event_type': 'save', 'directory': 'gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776580718.3721037, 'wait_for_prev_duration_secs': 7.033348083496094e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776580718.3739777, 'checkpointer_blocking_duration_secs': 1.2929940223693848, 'get_old_steps_start_time': 1776580719.666991, 'get_old_steps_duration_secs': 2.5033950805664062e-05, 'checkpoint_manager_blocking_start_time': 1776580718.3700128, 'checkpoint_manager_blocking_duration_secs': 1.7062435150146484} I0419 06:38:40.076410 134452881033024 checkpointing.py:408] Started an asynchronous checkpoint save for step 0 I0419 06:38:40.076462 134452881033024 max_utils.py:750] Memstats: After params initialized: I0419 06:38:40.076521 134452881033024 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_0(process=0,(0,0,0,0)) I0419 06:38:40.076555 134452881033024 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_1(process=0,(1,0,0,0)) I0419 06:38:40.076582 134452881033024 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_4(process=0,(0,1,0,0)) I0419 06:38:40.076606 134452881033024 max_utils.py:756] Using (GB) 0.43 / 31.25 (1.376000%) on TPU_5(process=0,(1,1,0,0)) I0419 06:38:40.369475 134452881033024 metric_logger.py:185] completed step: 0, seconds: 11.072, TFLOP/s/device: 1.227, Tokens/s/device: 184.976, total_weights: 65536, loss: 10.874 I0419 06:38:40.372324 134452881033024 metric_logger.py:269] To see full metrics 'tensorboard --logdir=gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/tensorboard/' I0419 06:38:40.772618 134452881033024 metric_logger.py:185] completed step: 1, seconds: 2.155, TFLOP/s/device: 6.305, Tokens/s/device: 950.373, total_weights: 65536, loss: 10.874 I0419 06:38:40.918942 134452881033024 metric_logger.py:185] completed step: 2, seconds: 0.404, TFLOP/s/device: 33.639, Tokens/s/device: 5070.449, total_weights: 65536, loss: 10.263 I0419 06:38:41.066591 134452881033024 metric_logger.py:185] completed step: 3, seconds: 0.008, TFLOP/s/device: 1647.324, Tokens/s/device: 248302.619, total_weights: 65536, loss: 9.734 I0419 06:38:59.740970 134452881033024 metric_logger.py:185] completed step: 4, seconds: 0.147, TFLOP/s/device: 92.197, Tokens/s/device: 13896.994, total_weights: 65536, loss: 9.278 I0419 06:38:59.749565 134452881033024 metric_logger.py:185] completed step: 5, seconds: 0.148, TFLOP/s/device: 91.971, Tokens/s/device: 13862.941, total_weights: 65536, loss: 8.892 I0419 06:38:59.895677 134452881033024 metric_logger.py:185] completed step: 6, seconds: 18.674, TFLOP/s/device: 0.728, Tokens/s/device: 109.673, total_weights: 65536, loss: 8.592 I0419 06:39:00.043463 134452881033024 metric_logger.py:185] completed step: 7, seconds: 0.006, TFLOP/s/device: 2182.321, Tokens/s/device: 328943.142, total_weights: 65536, loss: 8.384 I0419 06:39:00.190885 134452881033024 metric_logger.py:185] completed step: 8, seconds: 0.148, TFLOP/s/device: 92.097, Tokens/s/device: 13881.828, total_weights: 65536, loss: 8.255 I0419 06:39:00.337901 134452881033024 checkpointing.py:772] Waiting for step 9 to finish before checkpoint... I0419 06:39:00.338600 134452881033024 checkpointing.py:776] Waited 0.000720977783203125 seconds for step 9 to finish before starting checkpointing. I0419 06:39:00.340999 134452881033024 checkpoint_manager.py:2024] [process=0][thread=MainThread][step=0][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0419 06:39:00.677837 134327165769472 atomicity.py:137] Creating tmp directory gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0/items I0419 06:39:01.237961 134326602491648 checkpoint.py:188] Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776580740421672973, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1776580740421672973, "commit_timestamp_nsecs": null, "custom_metadata": {}} to gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0/_CHECKPOINT_METADATA I0419 06:39:02.230891 2689 google_auth_provider.cc:181] Running on GCE, using service account 562977990677-compute@developer.gserviceaccount.com I0419 06:39:04.791944 134329866909440 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 39 array_metadata.ArrayMetadata to gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0/items/array_metadatas/process_0 I0419 06:39:39.728554 134329858516736 base_pytree_checkpoint_handler.py:1232] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 37.498904s (commit=36.500008s, array_metadata_write=0.998896s) I0419 06:39:39.729637 134327148984064 base_pytree_checkpoint_handler.py:129] [process=0] /jax/checkpoint/write/gbytes_per_sec: 26.066 MiB/s (total gbytes: 1.5 GiB) (time elapsed: 60.600592851638794 s) (per-host) I0419 06:39:39.729761 134327148984064 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 60.062955s. I0419 06:39:40.812691 134327148984064 checkpoint.py:228] Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776580740421672973, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0/_CHECKPOINT_METADATA I0419 06:39:41.547217 134326602491648 checkpoint.py:247] Updated Metadata={'item_handlers': {'items': 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776580740421672973, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0/_CHECKPOINT_METADATA I0419 06:39:41.911551 134327148984064 array_metadata_store.py:411] [process=0][thread=async_save] Validated ArrayMetadata from all 8 hosts. Time taken: 0.000249s. I0419 06:39:42.209529 134327148984064 ocdbt_utils.py:49] Param validation support for Zarr3 will be added later (b/362328389). I0419 06:39:43.134888 134327148984064 base_pytree_checkpoint_handler.py:1356] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 2.189299s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0/items I0419 06:39:43.136977 134327148984064 atomicity.py:608] Finalizing gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0/items I0419 06:39:43.582844 134327148984064 atomicity.py:608] Finalizing gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0 I0419 06:39:44.869845 134327148984064 atomicity.py:790] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0`. I0419 06:39:44.870913 134327148984064 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 66.496908s. directory=gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/0 I0419 06:39:44.872282 134327148984064 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 65.205471s. I0419 06:39:44.872438 134325563549440 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0419 06:39:44.872539 134325563549440 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0419 06:39:44.872602 134325563549440 checkpoint_manager.py:2133] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is syncing with other hosts... I0419 06:39:44.875538 134325563549440 checkpoint_manager.py:2142] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is done on all hosts. I0419 06:39:44.875704 134452881033024 checkpoint_manager.py:2036] [process=0][thread=MainThread][step=0][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=0. W0419 06:39:44.875847 134452881033024 checkpoint_manager.py:1458] Waiting for previous save to complete took 44.534851 seconds. If this number is high, consider checkpointing less frequently. I0419 06:39:44.877743 134452881033024 checkpoint_manager.py:1518] [process=0] Saving checkpoint at step 9 I0419 06:39:44.881286 134452881033024 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9. I0419 06:39:46.244840 134327148984064 atomicity.py:137] Creating tmp directory gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9 I0419 06:39:46.257839 134452881033024 jax_array_handlers.py:358] Scheduling D2H of 39 prioritized jax.Array. I0419 06:39:46.257935 134452881033024 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 I0419 06:39:46.282829 134452881033024 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.025942s I0419 06:39:46.283649 134452881033024 base_pytree_checkpoint_handler.py:129] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 5.556 GiB/s (total gbytes: 1.5 GiB) (time elapsed: 0.2776501178741455 s) (per-host) I0419 06:39:46.283723 134452881033024 base_pytree_checkpoint_handler.py:737] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.277742s (batch_requests_ready=0.248774s, total_serialization_initiated=0.028197s, others=0.000771s) I0419 06:39:46.283801 134452881033024 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.281727s (all_items=0.000017s, per_item={'items': '0.00001693'}, temp_paths=0.281710) I0419 06:39:46.284630 134326092027648 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0419 06:39:46.284778 134452881033024 async_checkpointer.py:561] Finished blocking save. Time taken: 1.406965s. Continuing background save to gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9. I0419 06:39:46.751748 134452881033024 checkpoint_manager.py:1566] [process=0][thread=MainThread][step=9] Starting CheckpointManager Save Finalize thread=save_finalize I0419 06:39:46.752127 134325563549440 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0419 06:39:46.752289 134452881033024 standard_logger.py:34] {'step': 9, 'event_type': 'save', 'directory': 'gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776580740.3409698, 'wait_for_prev_duration_secs': 44.53485107421875, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776580784.8777807, 'checkpointer_blocking_duration_secs': 1.407114028930664, 'get_old_steps_start_time': 1776580786.2849147, 'get_old_steps_duration_secs': 2.5987625122070312e-05, 'checkpoint_manager_blocking_start_time': 1776580740.3388374, 'checkpoint_manager_blocking_duration_secs': 46.41341185569763} I0419 06:39:46.752451 134452881033024 checkpointing.py:408] Started an asynchronous checkpoint save for step 9 I0419 06:39:46.752501 134452881033024 checkpoint_manager.py:2024] [process=0][thread=MainThread][step=9][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0419 06:39:48.445027 134329866909440 atomicity.py:137] Creating tmp directory gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9/items I0419 06:39:52.930602 134326108813056 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 39 array_metadata.ArrayMetadata to gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9/items/array_metadatas/process_0 I0419 06:40:31.457045 134326100420352 base_pytree_checkpoint_handler.py:1232] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 41.902215s (commit=40.490711s, array_metadata_write=1.411505s) I0419 06:40:31.458460 134326092027648 base_pytree_checkpoint_handler.py:129] [process=0] /jax/checkpoint/write/gbytes_per_sec: 34.753 MiB/s (total gbytes: 1.5 GiB) (time elapsed: 45.452431440353394 s) (per-host) I0419 06:40:31.458578 134326092027648 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 45.173842s. I0419 06:40:33.522470 134326092027648 ocdbt_utils.py:49] Param validation support for Zarr3 will be added later (b/362328389). I0419 06:40:33.705031 134326092027648 array_metadata_store.py:411] [process=0][thread=async_save] Validated ArrayMetadata from all 8 hosts. Time taken: 0.000253s. I0419 06:40:34.245676 134326092027648 base_pytree_checkpoint_handler.py:1356] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 1.544834s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9/items I0419 06:40:34.247207 134326092027648 atomicity.py:608] Finalizing gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9/items I0419 06:40:35.103956 134326092027648 atomicity.py:608] Finalizing gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9 I0419 06:40:36.363990 134326092027648 atomicity.py:790] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9`. I0419 06:40:36.364790 134326092027648 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 51.486983s. directory=gs://lance-maxtext/linen_ckpt_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615/linen_xpk_feat_Migrate_Decoder_And_Tests_to_NNX_20260419_063615_01_base/checkpoints/9 I0419 06:40:36.366117 134326092027648 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 50.081380s. I0419 06:40:36.366324 134325563549440 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0419 06:40:36.366440 134325563549440 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0419 06:40:36.366503 134325563549440 checkpoint_manager.py:2133] [process=0][thread=save_finalize][step=9] CheckpointManager Save Finalize is syncing with other hosts... I0419 06:40:36.370059 134325563549440 checkpoint_manager.py:2142] [process=0][thread=save_finalize][step=9] CheckpointManager Save Finalize is done on all hosts. I0419 06:40:36.370238 134452881033024 checkpoint_manager.py:2036] [process=0][thread=MainThread][step=9][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=9. I0419 06:40:36.370386 134452881033024 checkpoint_manager.py:2013] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0419 06:40:36.371000 134452881033024 metric_logger.py:185] completed step: 9, seconds: 0.147, TFLOP/s/device: 92.121, Tokens/s/device: 13885.405, total_weights: 65536, loss: 8.179 Per train step: Total TFLOPs: 13.59 split as 93.93% learnable weight flops and 6.07% attention flops XPK End: Sun Apr 19 06:40:46 UTC 2026 EXIT_CODE=0