2026-04-15 21:54:02.584031: 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)
I0415 21:54:02.700196 123517920468096 max_utils.py:238] Skipping jax distributed system due to skip_jax_distributed_system=True flag.
I0415 21:54:45.642416 123517920468096 max_utils.py:800] System Information: Jax Version: 0.9.2
I0415 21:54:45.642522 123517920468096 max_utils.py:801] System Information: Jaxlib Version: 0.9.2
I0415 21:54:45.642557 123517920468096 max_utils.py:802] System Information: Jax Backend: PJRT C API
TFRT TPU v6 lite
Built on Apr 6 2026 20:48:10 (1775533690) cl/895581894
I0415 21:54:45.642581 123517920468096 train_utils.py:334] WARNING: 'dataset_path' might be pointing your local file system
I0415 21:54:45.642608 123517920468096 train_utils.py:347] WARNING: Sequence packing is essentially ignored for synthetic data. Please use a real dataset to use sequence packing.
I0415 21:54:45.642679 123517920468096 train.py:683] [DECOUPLED NO-OP] skipping cloud diagnostics wrapper.
W0415 21:54:45.732860 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
I0415 21:54:46.139198 123517920468096 maxtext_utils.py:1517] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1)
I0415 21:54:46.139452 123517920468096 checkpointing.py:677] Setting up checkpoint logger...
I0415 21:54:46.139497 123517920468096 checkpointing.py:233] Creating checkpoint manager with ocdbt=True and zarr3=True
I0415 21:54:46.139537 123517920468096 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0415 21:54:46.139974 123517920468096 base_pytree_checkpoint_handler.py:411] 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 0x70562ad08590>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0415 21:54:48.525811 123517920468096 checkpointing.py:265] Enabling policy for fixed interval checkpointing.
I0415 21:54:48.525985 123517920468096 checkpoint_manager.py:702] [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=('items',), item_handlers={'items': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x704fd2dbffe0>}, handler_registry=None
I0415 21:54:48.526272 123517920468096 composite_checkpoint_handler.py:237] Deferred registration for item: "items". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x704fd2dbffe0>` 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`.
I0415 21:54:48.526314 123517920468096 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x704faf3461e0>` 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`.
I0415 21:54:48.526344 123517920468096 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 0x704fd2dbffe0>, ('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x704fd2dbffe0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x704faf3461e0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x704faf3461e0>}).
I0415 21:54:48.526773 123517920468096 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.28
I0415 21:54:48.650093 123517920468096 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints
I0415 21:54:48.650314 123517920468096 checkpoint_manager.py:921] [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_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, 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), root_directory=gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x704faf345fd0>
I0415 21:54:48.650389 123517920468096 checkpointing.py:301] Checkpoint manager created!
W0415 21:54:48.666709 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
I0415 21:54:48.847874 123517920468096 nnx_wrappers.py:437] Unknown Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_norm_length', 'activation_embed').
I0415 21:54:48.847973 123517920468096 nnx_wrappers.py:437] Unknown Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None).
I0415 21:54:48.944364 123517920468096 attentions.py:1088] attentions/inputs_q Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed').
I0415 21:54:48.944445 123517920468096 attentions.py:1088] attentions/inputs_q Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None).
I0415 21:54:48.958684 123517920468096 attentions.py:1089] attentions/inputs_kv Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed').
I0415 21:54:48.958730 123517920468096 attentions.py:1089] attentions/inputs_kv Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None).
I0415 21:54:48.978983 123517920468096 attentions.py:1154] attentions/query Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim').
I0415 21:54:48.979037 123517920468096 attentions.py:1154] attentions/query Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None).
I0415 21:54:48.993180 123517920468096 attentions.py:1155] attentions/key Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim').
I0415 21:54:48.993228 123517920468096 attentions.py:1155] attentions/key Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None).
I0415 21:54:49.007363 123517920468096 attentions.py:1156] attentions/value Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim').
I0415 21:54:49.007411 123517920468096 attentions.py:1156] attentions/value Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None).
I0415 21:54:49.034016 123517920468096 attentions.py:1197] attentions/out Logical: bfloat16[8,2048,16,128]..................................... ('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv').
I0415 21:54:49.034157 123517920468096 attentions.py:1197] attentions/out Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None).
I0415 21:54:49.053382 123517920468096 linears.py:525] linears/x Logical: bfloat16[8,2048,7168]....................................... ('activation_batch', 'activation_length', 'activation_mlp').
I0415 21:54:49.053435 123517920468096 linears.py:525] linears/x Physical: bfloat16[8,2048,7168]....................................... ('fsdp', None, None).
I0415 21:54:49.237928 123517920468096 checkpointing.py:577] checkpoint manager exists so trying to load this run's existing checkpoint
I0415 21:54:49.238035 123517920468096 checkpointing.py:665] No existing checkpoints found, not restoring checkpoint.
W0415 21:54:49.290532 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
[DECOUPLED NO-OP] gcs_storage: using stubs.
[DECOUPLED NO-OP] mldiagnostics: using stub.
[DECOUPLED NO-OP] mldiagnostics: using stub.
[DECOUPLED NO-OP] mldiagnostics: using stub.
[DECOUPLED NO-OP] workload_monitor: using stub.
[DECOUPLED NO-OP] vertex_tensorboard: using stub.
fsdp: 8
I0415 21:54:49.670650 123517920468096 maxtext_utils.py:1620] params/params/decoder/decoder_norm/scale
Shape: float32[2048]
Logical: P('norm',)
Physical: (None,)
I0415 21:54:49.670760 123517920468096 maxtext_utils.py:1620] params/params/decoder/layers/mlp/wi_0/kernel
Shape: float32[2048,16,7168]
Logical: P('embed', 'layers', 'mlp')
Physical: ('fsdp', None, None)
I0415 21:54:49.670802 123517920468096 maxtext_utils.py:1620] params/params/decoder/layers/mlp/wi_1/kernel
Shape: float32[2048,16,7168]
Logical: P('embed', 'layers', 'mlp')
Physical: ('fsdp', None, None)
I0415 21:54:49.670847 123517920468096 maxtext_utils.py:1620] params/params/decoder/layers/mlp/wo/kernel
Shape: float32[7168,16,2048]
Logical: P('mlp', 'layers', 'embed')
Physical: (None, None, 'fsdp')
I0415 21:54:49.670886 123517920468096 maxtext_utils.py:1620] params/params/decoder/layers/post_self_attention_layer_norm/scale
Shape: float32[2048,16]
Logical: P('norm', 'layers')
Physical: (None, None)
I0415 21:54:49.670916 123517920468096 maxtext_utils.py:1620] params/params/decoder/layers/pre_self_attention_layer_norm/scale
Shape: float32[2048,16]
Logical: P('norm', 'layers')
Physical: (None, None)
I0415 21:54:49.670955 123517920468096 maxtext_utils.py:1620] 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)
I0415 21:54:49.670992 123517920468096 maxtext_utils.py:1620] 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')
I0415 21:54:49.671020 123517920468096 maxtext_utils.py:1620] 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)
I0415 21:54:49.671060 123517920468096 maxtext_utils.py:1620] 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)
I0415 21:54:49.671097 123517920468096 maxtext_utils.py:1620] params/params/decoder/logits_dense/kernel
Shape: float32[2048,32000]
Logical: P('embed_vocab', 'vocab')
Physical: ('fsdp', None)
I0415 21:54:49.671130 123517920468096 maxtext_utils.py:1620] params/params/token_embedder/embedding
Shape: float32[32000,2048]
Logical: P('vocab', 'embed_vocab')
Physical: (None, 'fsdp')
I0415 21:54:50.091156 123517920468096 train.py:155] train/xent Logical: float32[8,2048]............................................. ('activation_embed_and_logits_batch', 'activation_length').
I0415 21:54:50.091234 123517920468096 train.py:155] train/xent Physical: float32[8,2048]............................................. ('fsdp', None).
I0415 21:54:50.105030 123517920468096 train.py:162] train/z_loss Logical: float32[8,2048]............................................. ('activation_embed_and_logits_batch', 'activation_length').
I0415 21:54:50.105094 123517920468096 train.py:162] train/z_loss Physical: float32[8,2048]............................................. ('fsdp', None).
W0415 21:54:50.744348 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
I0415 21:54:50.860847 123517920468096 max_utils.py:791] Total memory size: 3.6 GB, Output size: 1.5 GB, Temp size: 2.0 GB, Argument size: 1.5 GB, Host temp size: 0.0 GB.
I0415 21:54:50.861328 123517920468096 max_utils.py:194] tensorboardX not available; using no-op SummaryWriter.
I0415 21:54:50.861464 123517920468096 metric_logger.py:289] number parameters: 1.104 billion
W0415 21:54:51.691634 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
I0415 21:54:51.788017 123517920468096 checkpointing.py:772] Waiting for step 0 to finish before checkpoint...
I0415 21:54:51.905162 123517920468096 checkpointing.py:776] Waited 0.11713385581970215 seconds for step 0 to finish before starting checkpointing.
I0415 21:54:51.905569 123517920468096 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
I0415 21:54:51.905713 123517920468096 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 0
I0415 21:54:51.905769 123517920468096 checkpointer.py:241] [process=0] Started saving checkpoint to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0.
I0415 21:54:52.091032 123517920468096 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0
I0415 21:54:52.887066 123517920468096 checkpoint.py:188] Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776290092648198647, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1776290092648198647, "commit_timestamp_nsecs": null, "custom_metadata": {}} to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0/_CHECKPOINT_METADATA
I0415 21:54:52.888079 123517920468096 signaling_client.py:373] Using ThreadSafeKeyValueSignalingClient
W0415 21:54:52.888397 123517920468096 future.py:328] [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
I0415 21:54:52.973686 123410441045568 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0/items
I0415 21:54:53.530998 123517920468096 jax_array_handlers.py:347] Scheduling D2H of 39 prioritized jax.Array.
I0415 21:54:53.531127 123517920468096 replica_slices.py:410] 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
I0415 21:54:53.547735 3274670 google_auth_provider.cc:149] Using credentials at ~/.config/gcloud/application_default_credentials.json
I0415 21:54:53.547782 3274670 google_auth_provider.cc:156] Using OAuth2 AuthProvider
W0415 21:54:53.760434 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:54:53.776449 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:54:53.792382 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:54:53.815108 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:54:53.826067 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:54:53.832171 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:54:53.838163 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:54:53.842320 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
I0415 21:54:54.997176 123517920468096 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 1.466536s
I0415 21:54:55.001589 123517920468096 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 7.868 GiB/s (total gbytes: 12.3 GiB) (time elapsed: a second) (per-host)
I0415 21:54:55.002003 123517920468096 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.568969s (batch_requests_ready=0.093959s, total_serialization_initiated=1.470361s, others=0.004649s)
I0415 21:54:55.002115 123517920468096 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 2.114110s (all_items=0.000031s, per_item={'items': '0.00003076'}, temp_paths=2.114079)
W0415 21:54:55.002156 123517920468096 future.py:328] [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
I0415 21:54:55.740778 123410430559808 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 39 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0/items/array_metadatas/process_0
I0415 21:55:26.354982 123416178853440 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 31.353042s (commit=30.910284s, array_metadata_write=0.442758s)
I0415 21:55:26.356056 123517920468096 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 383.829 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 32 seconds) (per-host)
I0415 21:55:26.618088 123517920468096 checkpoint.py:228] Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776290092648198647, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0/_CHECKPOINT_METADATA
I0415 21:55:27.020816 123517920468096 checkpoint.py:247] Updated Metadata={'item_handlers': {'items': 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776290092648198647, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0/_CHECKPOINT_METADATA
I0415 21:55:27.161459 123517920468096 array_metadata_store.py:367] [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
I0415 21:55:27.605604 123517920468096 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389).
I0415 21:55:27.605823 123517920468096 base_pytree_checkpoint_handler.py:1342] [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.584596s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0/items
I0415 21:55:27.606506 123517920468096 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0/items
I0415 21:55:27.854979 123517920468096 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0
I0415 21:55:28.542458 123517920468096 atomicity.py:794] [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0`.
I0415 21:55:28.543125 123517920468096 checkpointer.py:290] Finished synchronous save in 36.64 seconds to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/0
I0415 21:55:28.543309 123517920468096 checkpoint_manager.py:2103] [process=0][thread=MainThread][step=0] CheckpointManager Save Finalize is syncing with other hosts...
I0415 21:55:28.543354 123517920468096 checkpoint_manager.py:2112] [process=0][thread=MainThread][step=0] CheckpointManager Save Finalize is done on all hosts.
I0415 21:55:28.543377 123517920468096 checkpoint_manager.py:1570] [process=0][thread=MainThread][step=0] Finished synchronous save.
I0415 21:55:28.543416 123517920468096 standard_logger.py:34] {'step': 0, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1776290091.9055557, 'wait_for_prev_duration_secs': 4.172325134277344e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776290091.9057333, 'checkpointer_blocking_duration_secs': 36.63743710517883, 'get_old_steps_start_time': 1776290128.5432045, 'get_old_steps_duration_secs': 2.956390380859375e-05, 'checkpoint_manager_blocking_start_time': 1776290091.9054768, 'checkpoint_manager_blocking_duration_secs': 36.637914180755615}
I0415 21:55:28.543510 123517920468096 checkpointing.py:410] Saved a checkpoint at step 0.
I0415 21:55:28.543559 123517920468096 max_utils.py:750]
Memstats: After params initialized:
I0415 21:55:28.543614 123517920468096 max_utils.py:756] Using (GB) 1.59 / 31.25 (5.088000%) on TPU_0(process=0,(0,0,0,0))
I0415 21:55:28.543640 123517920468096 max_utils.py:756] Using (GB) 1.59 / 31.25 (5.088000%) on TPU_1(process=0,(1,0,0,0))
I0415 21:55:28.543659 123517920468096 max_utils.py:756] Using (GB) 1.59 / 31.25 (5.088000%) on TPU_2(process=0,(0,1,0,0))
I0415 21:55:28.543678 123517920468096 max_utils.py:756] Using (GB) 1.59 / 31.25 (5.088000%) on TPU_3(process=0,(1,1,0,0))
I0415 21:55:28.543696 123517920468096 max_utils.py:756] Using (GB) 1.59 / 31.25 (5.088000%) on TPU_4(process=0,(0,2,0,0))
I0415 21:55:28.543714 123517920468096 max_utils.py:756] Using (GB) 1.59 / 31.25 (5.088000%) on TPU_5(process=0,(1,2,0,0))
I0415 21:55:28.543730 123517920468096 max_utils.py:756] Using (GB) 1.59 / 31.25 (5.088000%) on TPU_6(process=0,(0,3,0,0))
I0415 21:55:28.543747 123517920468096 max_utils.py:756] Using (GB) 1.59 / 31.25 (5.088000%) on TPU_7(process=0,(1,3,0,0))
W0415 21:55:28.549899 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:55:28.555614 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
W0415 21:55:28.560189 3272824 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.
I0415 21:55:28.751264 123517920468096 metric_logger.py:185] completed step: 0, seconds: 0.926, TFLOP/s/device: 14.671, Tokens/s/device: 2211.443, total_weights: 16384, loss: 10.874
I0415 21:55:28.752099 123517920468096 metric_logger.py:269] To see full metrics 'tensorboard --logdir=gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/tensorboard/'
I0415 21:55:28.866091 123517920468096 metric_logger.py:185] completed step: 1, seconds: 36.960, TFLOP/s/device: 0.368, Tokens/s/device: 55.411, total_weights: 16384, loss: 10.874
I0415 21:55:28.984775 123517920468096 metric_logger.py:185] completed step: 2, seconds: 0.007, TFLOP/s/device: 2030.354, Tokens/s/device: 306037.059, total_weights: 16384, loss: 10.276
I0415 21:55:29.103062 123517920468096 metric_logger.py:185] completed step: 3, seconds: 0.116, TFLOP/s/device: 117.357, Tokens/s/device: 17689.331, total_weights: 16384, loss: 9.282
I0415 21:55:29.107421 123517920468096 checkpointing.py:772] Waiting for step 5 to finish before checkpoint...
I0415 21:55:29.339282 123517920468096 checkpointing.py:776] Waited 0.23183345794677734 seconds for step 5 to finish before starting checkpointing.
I0415 21:55:29.339653 123517920468096 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
I0415 21:55:29.339801 123517920468096 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 5
I0415 21:55:29.339857 123517920468096 checkpointer.py:241] [process=0] Started saving checkpoint to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5.
I0415 21:55:29.513097 123517920468096 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5
W0415 21:55:30.364597 123517920468096 future.py:328] [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
I0415 21:55:30.448101 123416178853440 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5/items
I0415 21:55:30.999702 123517920468096 jax_array_handlers.py:347] Scheduling D2H of 39 prioritized jax.Array.
I0415 21:55:30.999810 123517920468096 replica_slices.py:410] 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
I0415 21:55:31.157889 123517920468096 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.158532s
I0415 21:55:31.162073 123517920468096 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 49.055 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 251 milliseconds) (per-host)
I0415 21:55:31.162135 123517920468096 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.251680s (batch_requests_ready=0.085541s, total_serialization_initiated=0.162081s, others=0.004059s)
I0415 21:55:31.162203 123517920468096 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.797929s (all_items=0.000032s, per_item={'items': '0.00003219'}, temp_paths=0.797897)
W0415 21:55:31.162235 123517920468096 future.py:328] [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
I0415 21:55:31.887923 123410805950016 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 39 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5/items/array_metadatas/process_0
I0415 21:56:03.082136 123410430559808 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 31.919584s (commit=31.486465s, array_metadata_write=0.433119s)
I0415 21:56:03.083255 123517920468096 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 392.779 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 32 seconds) (per-host)
I0415 21:56:03.855791 123517920468096 array_metadata_store.py:367] [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
I0415 21:56:04.268642 123517920468096 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389).
I0415 21:56:04.269195 123517920468096 base_pytree_checkpoint_handler.py:1342] [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.558000s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5/items
I0415 21:56:04.269859 123517920468096 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5/items
I0415 21:56:04.509345 123517920468096 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5
I0415 21:56:05.164222 123517920468096 atomicity.py:794] [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5`.
I0415 21:56:05.164869 123517920468096 checkpointer.py:290] Finished synchronous save in 35.83 seconds to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/5
I0415 21:56:05.165037 123517920468096 checkpoint_manager.py:2103] [process=0][thread=MainThread][step=5] CheckpointManager Save Finalize is syncing with other hosts...
I0415 21:56:05.165100 123517920468096 checkpoint_manager.py:2112] [process=0][thread=MainThread][step=5] CheckpointManager Save Finalize is done on all hosts.
I0415 21:56:05.165124 123517920468096 checkpoint_manager.py:1570] [process=0][thread=MainThread][step=5] Finished synchronous save.
I0415 21:56:05.165161 123517920468096 standard_logger.py:34] {'step': 5, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1776290129.3396308, 'wait_for_prev_duration_secs': 4.76837158203125e-05, 'time_between_consecutive_saves_sec': 0.7962620258331299, 'checkpointer_blocking_start_time': 1776290129.3398209, 'checkpointer_blocking_duration_secs': 35.825098514556885, 'get_old_steps_start_time': 1776290165.164949, 'get_old_steps_duration_secs': 2.7894973754882812e-05, 'checkpoint_manager_blocking_start_time': 1776290129.3395696, 'checkpoint_manager_blocking_duration_secs': 35.82556962966919}
I0415 21:56:05.165240 123517920468096 checkpointing.py:410] Saved a checkpoint at step 5.
I0415 21:56:05.165764 123517920468096 metric_logger.py:185] completed step: 4, seconds: 0.119, TFLOP/s/device: 114.515, Tokens/s/device: 17260.997, total_weights: 16384, loss: 8.327
I0415 21:56:05.174315 123517920468096 metric_logger.py:185] completed step: 5, seconds: 0.118, TFLOP/s/device: 115.127, Tokens/s/device: 17353.138, total_weights: 16384, loss: 7.367
I0415 21:56:05.289135 123517920468096 metric_logger.py:185] completed step: 6, seconds: 36.064, TFLOP/s/device: 0.377, Tokens/s/device: 56.789, total_weights: 16384, loss: 6.459
I0415 21:56:05.407774 123517920468096 metric_logger.py:185] completed step: 7, seconds: 0.007, TFLOP/s/device: 1900.829, Tokens/s/device: 286513.710, total_weights: 16384, loss: 5.622
I0415 21:56:05.526539 123517920468096 metric_logger.py:185] completed step: 8, seconds: 0.116, TFLOP/s/device: 117.395, Tokens/s/device: 17694.986, total_weights: 16384, loss: 4.864
I0415 21:56:05.644713 123517920468096 checkpointing.py:772] Waiting for step 9 to finish before checkpoint...
I0415 21:56:05.645537 123517920468096 checkpointing.py:776] Waited 0.000843048095703125 seconds for step 9 to finish before starting checkpointing.
I0415 21:56:05.645771 123517920468096 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
I0415 21:56:05.645910 123517920468096 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 9
I0415 21:56:05.645970 123517920468096 checkpointer.py:241] [process=0] Started saving checkpoint to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9.
I0415 21:56:05.808576 123517920468096 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9
W0415 21:56:06.640547 123517920468096 future.py:328] [process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
I0415 21:56:06.740309 123416178853440 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9/items
I0415 21:56:07.301529 123517920468096 jax_array_handlers.py:347] Scheduling D2H of 39 prioritized jax.Array.
I0415 21:56:07.301637 123517920468096 replica_slices.py:410] 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
I0415 21:56:07.463732 123517920468096 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.162540s
I0415 21:56:07.468206 123517920468096 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 48.132 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 256 milliseconds) (per-host)
I0415 21:56:07.468284 123517920468096 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.256540s (batch_requests_ready=0.086262s, total_serialization_initiated=0.165999s, others=0.004279s)
I0415 21:56:07.468376 123517920468096 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.828145s (all_items=0.000019s, per_item={'items': '0.00001907'}, temp_paths=0.828126)
W0415 21:56:07.468411 123517920468096 future.py:328] [process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
I0415 21:56:08.229360 123410552194624 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 39 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9/items/array_metadatas/process_0
I0415 21:56:43.518459 123410805950016 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 36.049702s (commit=35.584029s, array_metadata_write=0.465673s)
I0415 21:56:43.519620 123517920468096 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 348.047 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 36 seconds) (per-host)
I0415 21:56:44.311070 123517920468096 array_metadata_store.py:367] [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
I0415 21:56:44.752323 123517920468096 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389).
I0415 21:56:44.752806 123517920468096 base_pytree_checkpoint_handler.py:1342] [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.580091s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9/items
I0415 21:56:44.753471 123517920468096 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9/items
I0415 21:56:45.006982 123517920468096 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9
I0415 21:56:45.759462 123517920468096 atomicity.py:794] [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9`.
I0415 21:56:45.760118 123517920468096 checkpointer.py:290] Finished synchronous save in 40.11 seconds to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints/9
I0415 21:56:45.760250 123517920468096 checkpoint_manager.py:2103] [process=0][thread=MainThread][step=9] CheckpointManager Save Finalize is syncing with other hosts...
I0415 21:56:45.760293 123517920468096 checkpoint_manager.py:2112] [process=0][thread=MainThread][step=9] CheckpointManager Save Finalize is done on all hosts.
I0415 21:56:45.760315 123517920468096 checkpoint_manager.py:1570] [process=0][thread=MainThread][step=9] Finished synchronous save.
I0415 21:56:45.760347 123517920468096 standard_logger.py:34] {'step': 9, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_14_async_ckpt_false/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1776290165.645751, 'wait_for_prev_duration_secs': 4.363059997558594e-05, 'time_between_consecutive_saves_sec': 0.4806342124938965, 'checkpointer_blocking_start_time': 1776290165.64593, 'checkpointer_blocking_duration_secs': 40.11423087120056, 'get_old_steps_start_time': 1776290205.760175, 'get_old_steps_duration_secs': 1.8596649169921875e-05, 'checkpoint_manager_blocking_start_time': 1776290165.6457186, 'checkpoint_manager_blocking_duration_secs': 40.11461067199707}
I0415 21:56:45.760422 123517920468096 checkpointing.py:410] Saved a checkpoint at step 9.
I0415 21:56:45.760445 123517920468096 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
I0415 21:56:45.760464 123517920468096 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
I0415 21:56:45.760976 123517920468096 metric_logger.py:185] completed step: 9, seconds: 0.119, TFLOP/s/device: 114.603, Tokens/s/device: 17274.246, total_weights: 16384, loss: 4.189
Per train step:
Total TFLOPs: 13.59
split as 93.93% learnable weight flops and 6.07% attention flops