main (51c7f2b8e) vs feat/migrate-nnx-utils (dc4d58891)| Metric | main 51c7f2b8e | feat/migrate-nnx-utils dc4d58891 | Diff (feat/migrate-nnx-utils − main) |
|---|---|---|---|
| Parameters | 1.105 billion | 1.105 billion | — |
| Final loss | 10.8440 | 10.8510 | +0.007 |
| TFLOP/s | 2.942 | 122.790 | +119.848 |
| Tok/s | 443.5 | 18508.3 | +18064.839 |
| Avg s/step | — | 2.222 | — |
| Memory % | 5.15 | 5.15 | 0 |
| JAX | 0.9.2 | 0.9.2 | — |
Diff = branch value − main value. Green = branch improved. Red = branch regressed.
2026-04-15 21:22:12.864609: 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:22:12.982217 140311636520064 max_utils.py:238] Skipping jax distributed system due to skip_jax_distributed_system=True flag. I0415 21:23:04.649687 140311636520064 max_utils.py:800] System Information: Jax Version: 0.9.2 I0415 21:23:04.649809 140311636520064 max_utils.py:801] System Information: Jaxlib Version: 0.9.2 I0415 21:23:04.649846 140311636520064 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:23:04.649870 140311636520064 train_utils.py:334] WARNING: 'dataset_path' might be pointing your local file system I0415 21:23:04.649950 140311636520064 train.py:683] [DECOUPLED NO-OP] skipping cloud diagnostics wrapper. W0415 21:23:04.743341 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 21:23:05.142014 140311636520064 maxtext_utils.py:1517] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0415 21:23:05.142354 140311636520064 checkpointing.py:677] Setting up checkpoint logger... I0415 21:23:05.142402 140311636520064 checkpointing.py:233] Creating checkpoint manager with ocdbt=True and zarr3=True I0415 21:23:05.142463 140311636520064 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0415 21:23:05.142978 140311636520064 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 0x7f9c42500bf0>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0415 21:23:07.596238 140311636520064 checkpointing.py:265] Enabling policy for fixed interval checkpointing. I0415 21:23:07.596463 140311636520064 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 0x7f95c6716810>}, handler_registry=None I0415 21:23:07.596796 140311636520064 composite_checkpoint_handler.py:237] Deferred registration for item: "items". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7f95c6716810>` 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:23:07.596839 140311636520064 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7f95c6340b90>` 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:23:07.596868 140311636520064 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 0x7f95c6716810>, ('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7f95c6716810>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7f95c6340b90>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7f95c6340b90>}). I0415 21:23:07.597193 140311636520064 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.28 I0415 21:23:07.597252 140311636520064 async_checkpointer.py:177] [process=0][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>.<lambda> at 0x7f95c6338fe0> timeout: 600 secs and primary_host=0 for async checkpoint writes I0415 21:23:07.725083 140311636520064 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints I0415 21:23:07.725334 140311636520064 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=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), root_directory=gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7f95c6340290> I0415 21:23:07.725427 140311636520064 checkpointing.py:301] Checkpoint manager created! I0415 21:23:07.804394 140311636520064 dataset_info.py:707] Load dataset info from tests/assets/local_datasets/c4_en_dataset_minimal/c4/en/3.1.0 I0415 21:23:07.807415 140311636520064 reader.py:262] Creating a tf.data.Dataset reading 8 files located in folders: tests/assets/local_datasets/c4_en_dataset_minimal/c4/en/3.1.0. I0415 21:23:07.861209 140311636520064 logging_logger.py:49] Constructing tf.data.Dataset __local_c4_builder for split train, from tests/assets/local_datasets/c4_en_dataset_minimal/c4/en/3.1.0 I0415 21:23:07.888165 140311636520064 tokenizer.py:245] Tokenizer path: src/maxtext/assets/tokenizers/tokenizer.llama2 I0415 21:23:07.888232 140311636520064 tokenizer.py:187] Loading sentencepiece tokenizer: src/maxtext/assets/tokenizers/tokenizer.llama2 I0415 21:23:08.994372 140311636520064 nnx_wrappers.py:437] Unknown Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0415 21:23:08.994486 140311636520064 nnx_wrappers.py:437] Unknown Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0415 21:23:09.098851 140311636520064 attentions.py:1088] attentions/inputs_q Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0415 21:23:09.098935 140311636520064 attentions.py:1088] attentions/inputs_q Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0415 21:23:09.114009 140311636520064 attentions.py:1089] attentions/inputs_kv Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0415 21:23:09.114067 140311636520064 attentions.py:1089] attentions/inputs_kv Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0415 21:23:09.160683 140311636520064 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:23:09.160766 140311636520064 attentions.py:1154] attentions/query Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0415 21:23:09.175873 140311636520064 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:23:09.175928 140311636520064 attentions.py:1155] attentions/key Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0415 21:23:09.191189 140311636520064 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:23:09.191239 140311636520064 attentions.py:1156] attentions/value Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0415 21:23:09.213714 140311636520064 attentions.py:1197] attentions/out Logical: bfloat16[8,2048,16,128]..................................... ('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv'). I0415 21:23:09.213785 140311636520064 attentions.py:1197] attentions/out Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0415 21:23:09.259338 140311636520064 linears.py:525] linears/x Logical: bfloat16[8,2048,7168]....................................... ('activation_batch', 'activation_length', 'activation_mlp'). I0415 21:23:09.259407 140311636520064 linears.py:525] linears/x Physical: bfloat16[8,2048,7168]....................................... ('fsdp', None, None). I0415 21:23:09.639064 140311636520064 checkpointing.py:577] checkpoint manager exists so trying to load this run's existing checkpoint I0415 21:23:09.639168 140311636520064 checkpointing.py:665] No existing checkpoints found, not restoring checkpoint. W0415 21:23:09.707757 3236686 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:23:10.486254 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486362 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486399 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486429 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486454 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486479 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486502 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486524 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486544 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486566 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486587 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486607 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486628 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486650 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486674 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486694 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486714 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486734 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486753 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486773 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486792 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486811 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486831 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486851 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486871 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486890 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486910 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486929 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486948 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.486967 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.486998 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.487018 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.487038 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.487073 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.487093 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.487113 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.487133 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.487153 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.487172 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.487192 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.487211 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 21:23:10.487230 140311636520064 maxtext_utils.py:1620] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 21:23:10.487269 140311636520064 maxtext_utils.py:1620] params/params/decoder/decoder_norm/scale Shape: float32[2048] Logical: P('norm',) Physical: (None,) I0415 21:23:10.487310 140311636520064 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:23:10.487336 140311636520064 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:23:10.487369 140311636520064 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:23:10.487402 140311636520064 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:23:10.487426 140311636520064 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:23:10.487461 140311636520064 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:23:10.487496 140311636520064 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:23:10.487521 140311636520064 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:23:10.487544 140311636520064 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:23:10.487577 140311636520064 maxtext_utils.py:1620] params/params/decoder/logits_dense/kernel Shape: float32[2048,32000] Logical: P('embed_vocab', 'vocab') Physical: ('fsdp', None) I0415 21:23:10.487609 140311636520064 maxtext_utils.py:1620] params/params/token_embedder/embedding Shape: float32[32000,2048] Logical: P('vocab', 'embed_vocab') Physical: (None, 'fsdp') I0415 21:23:11.942757 140311636520064 train.py:155] train/xent Logical: float32[8,2048]............................................. ('activation_embed_and_logits_batch', 'activation_length'). I0415 21:23:11.942915 140311636520064 train.py:155] train/xent Physical: float32[8,2048]............................................. ('fsdp', None). I0415 21:23:11.957688 140311636520064 train.py:162] train/z_loss Logical: float32[8,2048]............................................. ('activation_embed_and_logits_batch', 'activation_length'). I0415 21:23:11.957744 140311636520064 train.py:162] train/z_loss Physical: float32[8,2048]............................................. ('fsdp', None). W0415 21:23:12.400703 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 21:23:12.579716 140311636520064 max_utils.py:791] Total memory size: 3.5 GB, Output size: 1.5 GB, Temp size: 1.9 GB, Argument size: 1.5 GB, Host temp size: 0.0 GB. I0415 21:23:12.580368 140311636520064 max_utils.py:194] tensorboardX not available; using no-op SummaryWriter. I0415 21:23:12.580603 140311636520064 metric_logger.py:289] number parameters: 1.105 billion W0415 21:23:16.996983 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 21:23:17.199283 140311636520064 checkpointing.py:772] Waiting for step 0 to finish before checkpoint... I0415 21:23:17.312263 140311636520064 checkpointing.py:776] Waited 0.1129615306854248 seconds for step 0 to finish before starting checkpointing. I0415 21:23:17.312764 140311636520064 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0415 21:23:17.312918 140311636520064 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 0 I0415 21:23:17.313335 140311636520064 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints/0. I0415 21:23:17.401493 140311636520064 signaling_client.py:373] Using ThreadSafeKeyValueSignalingClient I0415 21:23:17.484201 140198388172352 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints/0 I0415 21:23:17.508609 140311636520064 jax_array_handlers.py:347] Scheduling D2H of 81 prioritized jax.Array. I0415 21:23:17.508707 140311636520064 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:23:18.173768 140198377686592 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints/0/items W0415 21:23:18.216832 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.229578 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.238190 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.248386 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.260215 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.271386 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.277601 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.281847 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 21:23:18.390270 140198337840704 checkpoint.py:188] Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776288198078562506, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1776288198078562506, "commit_timestamp_nsecs": null, "custom_metadata": {}} to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0415 21:23:18.736185 3239351 google_auth_provider.cc:149] Using credentials at ~/.config/gcloud/application_default_credentials.json I0415 21:23:18.736248 3239351 google_auth_provider.cc:156] Using OAuth2 AuthProvider I0415 21:23:18.921861 140311636520064 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 1.413757s I0415 21:23:18.930513 140311636520064 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 8.081 GiB/s (total gbytes: 12.3 GiB) (time elapsed: a second) (per-host) I0415 21:23:18.930585 140311636520064 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.527348s (batch_requests_ready=0.098444s, total_serialization_initiated=1.420430s, others=0.008473s) I0415 21:23:18.930676 140311636520064 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 1.528642s (all_items=0.000020s, per_item={'items': '0.00002027'}, temp_paths=1.528621) I0415 21:23:18.932117 140198253954624 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0415 21:23:18.932211 140311636520064 async_checkpointer.py:561] Finished blocking save. Time taken: 1.619253s. Continuing background save to gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints/0. I0415 21:23:18.932435 140311636520064 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=0] Starting CheckpointManager Save Finalize thread=save_finalize I0415 21:23:18.932663 140198398658112 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0415 21:23:18.932782 140311636520064 standard_logger.py:34] {'step': 0, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776288197.3127515, 'wait_for_prev_duration_secs': 4.100799560546875e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776288197.3129387, 'checkpointer_blocking_duration_secs': 1.6193909645080566, 'get_old_steps_start_time': 1776288198.9323606, 'get_old_steps_duration_secs': 3.361701965332031e-05, 'checkpoint_manager_blocking_start_time': 1776288197.312663, 'checkpoint_manager_blocking_duration_secs': 1.6200878620147705} I0415 21:23:18.932919 140311636520064 checkpointing.py:408] Started an asynchronous checkpoint save for step 0 I0415 21:23:18.932968 140311636520064 max_utils.py:750] Memstats: After params initialized: I0415 21:23:18.933024 140311636520064 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_0(process=0,(0,0,0,0)) I0415 21:23:18.933062 140311636520064 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_1(process=0,(1,0,0,0)) I0415 21:23:18.933085 140311636520064 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_2(process=0,(0,1,0,0)) I0415 21:23:18.933104 140311636520064 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_3(process=0,(1,1,0,0)) I0415 21:23:18.933122 140311636520064 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_4(process=0,(0,2,0,0)) I0415 21:23:18.933140 140311636520064 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_5(process=0,(1,2,0,0)) I0415 21:23:18.933156 140311636520064 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_6(process=0,(0,3,0,0)) I0415 21:23:18.933172 140311636520064 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_7(process=0,(1,3,0,0)) W0415 21:23:18.940265 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.946471 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 21:23:18.951534 3236686 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 21:23:19.274335 140311636520064 metric_logger.py:185] completed step: 0, seconds: 4.618, TFLOP/s/device: 2.942, Tokens/s/device: 443.471, total_weights: 13328, loss: 10.844 I0415 21:23:19.275230 140311636520064 metric_logger.py:269] To see full metrics 'tensorboard --logdir=gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/tensorboard/' I0415 21:23:19.382153 140311636520064 metric_logger.py:185] completed step: 1, seconds: 2.072, TFLOP/s/device: 6.557, Tokens/s/device: 988.341, total_weights: 12332, loss: nan I0415 21:23:19.383253 140311636520064 metric_logger.py:225] Aborting training due to NaN loss. I0415 21:23:19.383364 140311636520064 metric_logger.py:185] completed step: 1, seconds: 2.072, TFLOP/s/device: 6.557, Tokens/s/device: 988.341, total_weights: 12332, loss: nan I0415 21:23:19.383466 140311636520064 metric_logger.py:225] Aborting training due to NaN loss. E0415 21:23:19.570272 140198285411904 future.py:315] [process=0][thread=array_type_handler][operation_id=1] _SignalingThread.run() raised an exception: cannot schedule new futures after shutdown Traceback (most recent call last): File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 312, in run super().run() File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/threading.py", line 1012, in run self._target(*self._args, **self._kwargs) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 257, in _target_setting_result self._result = target() ^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 391, in <lambda> target=lambda: asyncio_utils.run_sync(coro), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/asyncio_utils.py", line 36, in run_sync return asyncio.run(coro) ^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 195, in run return runner.run(main) ^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 118, in run return self._loop.run_until_complete(task) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete return future.result() ^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/jax_array_handlers.py", line 363, in _serialize_without_dispatcher await async_serialize_replica_slices_batch( File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/jax_array_handlers.py", line 608, in _async_serialize_replica_slices await asyncio.gather(*write_coros) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/array_metadata_store.py", line 195, in write if await async_path.exists(file_path): ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/path/async_path.py", line 67, in exists return await asyncio.to_thread(path.exists) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/threads.py", line 25, in to_thread return await loop.run_in_executor(None, func_call) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 867, in run_in_executor executor.submit(func, *args), loop=self) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/concurrent/futures/thread.py", line 171, in submit raise RuntimeError('cannot schedule new futures after shutdown') RuntimeError: cannot schedule new futures after shutdown E0415 21:23:19.573523 140198274926144 future.py:315] [process=0][thread=write_metadata_after_commits][operation_id=1] _SignalingThread.run() raised an exception: cannot schedule new futures after shutdown Traceback (most recent call last): File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 312, in run super().run() File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/threading.py", line 1012, in run self._target(*self._args, **self._kwargs) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 257, in _target_setting_result self._result = target() ^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 391, in <lambda> target=lambda: asyncio_utils.run_sync(coro), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/asyncio_utils.py", line 36, in run_sync return asyncio.run(coro) ^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 195, in run return runner.run(main) ^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 118, in run return self._loop.run_until_complete(task) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete return future.result() ^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py", line 1192, in _write_metadata_after_commits await asyncio.to_thread(commit_future.result) File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/threads.py", line 25, in to_thread return await loop.run_in_executor(None, func_call) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/concurrent/futures/thread.py", line 59, in run result = self.fn(*self.args, **self.kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 449, in result return self._f.result(timeout=timeout) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 398, in result return self._t.result(timeout=timeout) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 347, in result self.join(timeout=timeout) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 344, in join raise self._exception File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 312, in run super().run() File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/threading.py", line 1012, in run self._target(*self._args, **self._kwargs) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 257, in _target_setting_result self._result = target() ^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 391, in <lambda> target=lambda: asyncio_utils.run_sync(coro), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/asyncio_utils.py", line 36, in run_sync return asyncio.run(coro) ^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 195, in run return runner.run(main) ^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 118, in run return self._loop.run_until_complete(task) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete return future.result() ^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/jax_array_handlers.py", line 363, in _serialize_without_dispatcher await async_serialize_replica_slices_batch( File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/jax_array_handlers.py", line 608, in _async_serialize_replica_slices await asyncio.gather(*write_coros) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/array_metadata_store.py", line 195, in write if await async_path.exists(file_path): ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/path/async_path.py", line 67, in exists return await asyncio.to_thread(path.exists) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/threads.py", line 25, in to_thread return await loop.run_in_executor(None, func_call) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 867, in run_in_executor executor.submit(func, *args), loop=self) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/concurrent/futures/thread.py", line 171, in submit raise RuntimeError('cannot schedule new futures after shutdown') RuntimeError: cannot schedule new futures after shutdown E0415 21:23:19.575382 140198253954624 async_checkpointer.py:229] [process=0] Failed to run 3 Handler Commit operations or the Commit callback in background save thread, directory: gs://wanglance-maxtext/linen_ckpt_main_20260415_211126/linen_main_20260415_211126_05_fp8/checkpoints/0 Traceback (most recent call last): File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/checkpointers/async_checkpointer.py", line 215, in _thread_func _background_wait_for_commit_futures( File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/checkpointers/async_checkpointer.py", line 88, in _background_wait_for_commit_futures commit_future.result() File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 196, in result f.result(timeout=time_remaining) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 449, in result return self._f.result(timeout=timeout) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 398, in result return self._t.result(timeout=timeout) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 347, in result self.join(timeout=timeout) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 344, in join raise self._exception File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 312, in run super().run() File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/threading.py", line 1012, in run self._target(*self._args, **self._kwargs) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 257, in _target_setting_result self._result = target() ^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 391, in <lambda> target=lambda: asyncio_utils.run_sync(coro), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/asyncio_utils.py", line 36, in run_sync return asyncio.run(coro) ^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 195, in run return runner.run(main) ^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 118, in run return self._loop.run_until_complete(task) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete return future.result() ^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py", line 1192, in _write_metadata_after_commits await asyncio.to_thread(commit_future.result) File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/threads.py", line 25, in to_thread return await loop.run_in_executor(None, func_call) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/concurrent/futures/thread.py", line 59, in run result = self.fn(*self.args, **self.kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 449, in result return self._f.result(timeout=timeout) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 398, in result return self._t.result(timeout=timeout) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 347, in result self.join(timeout=timeout) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 344, in join raise self._exception File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 312, in run super().run() File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/threading.py", line 1012, in run self._target(*self._args, **self._kwargs) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 257, in _target_setting_result self._result = target() ^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/futures/future.py", line 391, in <lambda> target=lambda: asyncio_utils.run_sync(coro), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/asyncio_utils.py", line 36, in run_sync return asyncio.run(coro) ^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 195, in run return runner.run(main) ^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 118, in run return self._loop.run_until_complete(task) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete return future.result() ^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/jax_array_handlers.py", line 363, in _serialize_without_dispatcher await async_serialize_replica_slices_batch( File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/jax_array_handlers.py", line 608, in _async_serialize_replica_slices await asyncio.gather(*write_coros) File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/array_metadata_store.py", line 195, in write if await async_path.exists(file_path): ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/maxtext_venv/lib/python3.12/site-packages/orbax/checkpoint/_src/path/async_path.py", line 67, in exists return await asyncio.to_thread(path.exists) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/threads.py", line 25, in to_thread return await loop.run_in_executor(None, func_call) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 867, in run_in_executor executor.submit(func, *args), loop=self) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/concurrent/futures/thread.py", line 171, in submit raise RuntimeError('cannot schedule new futures after shutdown') RuntimeError: cannot schedule new futures after shutdown I0415 21:23:19.577101 140198398658112 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0415 21:23:19.577280 140198398658112 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is syncing with other hosts... I0415 21:23:19.577396 140198398658112 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is done on all hosts. Per train step: Total TFLOPs: 13.59 split as 93.93% learnable weight flops and 6.07% attention flops
2026-04-15 23:18:47.819746: 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 23:18:47.937408 127701190978688 max_utils.py:238] Skipping jax distributed system due to skip_jax_distributed_system=True flag. I0415 23:19:44.967139 127701190978688 max_utils.py:800] System Information: Jax Version: 0.9.2 I0415 23:19:44.967248 127701190978688 max_utils.py:801] System Information: Jaxlib Version: 0.9.2 I0415 23:19:44.967282 127701190978688 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 23:19:44.967307 127701190978688 train_utils.py:334] WARNING: 'dataset_path' might be pointing your local file system I0415 23:19:44.967388 127701190978688 train.py:715] [DECOUPLED NO-OP] skipping cloud diagnostics wrapper. W0415 23:19:45.058504 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 23:19:45.462851 127701190978688 maxtext_utils.py:1520] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0415 23:19:45.463169 127701190978688 checkpointing.py:677] Setting up checkpoint logger... I0415 23:19:45.463217 127701190978688 checkpointing.py:233] Creating checkpoint manager with ocdbt=True and zarr3=True I0415 23:19:45.463256 127701190978688 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0415 23:19:45.463705 127701190978688 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 0x7424293f8c50>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0415 23:19:47.861852 127701190978688 checkpointing.py:265] Enabling policy for fixed interval checkpointing. I0415 23:19:47.862037 127701190978688 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 0x741dad13a1e0>}, handler_registry=None I0415 23:19:47.862347 127701190978688 composite_checkpoint_handler.py:237] Deferred registration for item: "items". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x741dad13a1e0>` 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 23:19:47.862388 127701190978688 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x741dad139970>` 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 23:19:47.862417 127701190978688 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 0x741dad13a1e0>, ('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x741dad13a1e0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x741dad139970>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x741dad139970>}). I0415 23:19:47.863034 127701190978688 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.28 I0415 23:19:47.863101 127701190978688 async_checkpointer.py:177] [process=0][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>.<lambda> at 0x741dad1354e0> timeout: 600 secs and primary_host=0 for async checkpoint writes I0415 23:19:47.974970 127701190978688 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints I0415 23:19:47.975213 127701190978688 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=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), root_directory=gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x741dad510bc0> I0415 23:19:47.975298 127701190978688 checkpointing.py:301] Checkpoint manager created! I0415 23:19:48.058110 127701190978688 dataset_info.py:707] Load dataset info from tests/assets/local_datasets/c4_en_dataset_minimal/c4/en/3.1.0 I0415 23:19:48.061156 127701190978688 reader.py:262] Creating a tf.data.Dataset reading 8 files located in folders: tests/assets/local_datasets/c4_en_dataset_minimal/c4/en/3.1.0. I0415 23:19:48.115343 127701190978688 logging_logger.py:49] Constructing tf.data.Dataset __local_c4_builder for split train, from tests/assets/local_datasets/c4_en_dataset_minimal/c4/en/3.1.0 I0415 23:19:48.147462 127701190978688 tokenizer.py:245] Tokenizer path: src/maxtext/assets/tokenizers/tokenizer.llama2 I0415 23:19:48.147529 127701190978688 tokenizer.py:187] Loading sentencepiece tokenizer: src/maxtext/assets/tokenizers/tokenizer.llama2 I0415 23:19:49.276921 127701190978688 nnx_wrappers.py:437] Unknown Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0415 23:19:49.277036 127701190978688 nnx_wrappers.py:437] Unknown Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0415 23:19:49.380664 127701190978688 attentions.py:1088] attentions/inputs_q Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0415 23:19:49.380741 127701190978688 attentions.py:1088] attentions/inputs_q Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0415 23:19:49.395710 127701190978688 attentions.py:1089] attentions/inputs_kv Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0415 23:19:49.395758 127701190978688 attentions.py:1089] attentions/inputs_kv Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0415 23:19:49.442855 127701190978688 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 23:19:49.442923 127701190978688 attentions.py:1154] attentions/query Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0415 23:19:49.458015 127701190978688 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 23:19:49.458077 127701190978688 attentions.py:1155] attentions/key Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0415 23:19:49.473093 127701190978688 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 23:19:49.473143 127701190978688 attentions.py:1156] attentions/value Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0415 23:19:49.495578 127701190978688 attentions.py:1197] attentions/out Logical: bfloat16[8,2048,16,128]..................................... ('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv'). I0415 23:19:49.495638 127701190978688 attentions.py:1197] attentions/out Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0415 23:19:49.541207 127701190978688 linears.py:525] linears/x Logical: bfloat16[8,2048,7168]....................................... ('activation_batch', 'activation_length', 'activation_mlp'). I0415 23:19:49.541271 127701190978688 linears.py:525] linears/x Physical: bfloat16[8,2048,7168]....................................... ('fsdp', None, None). I0415 23:19:49.936945 127701190978688 checkpointing.py:577] checkpoint manager exists so trying to load this run's existing checkpoint I0415 23:19:49.937083 127701190978688 checkpointing.py:665] No existing checkpoints found, not restoring checkpoint. W0415 23:19:50.004399 3362481 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 23:19:50.778913 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779014 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779231 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779268 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779303 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779328 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779351 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779372 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779392 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779414 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779437 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779458 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779479 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779500 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779522 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779542 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779562 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779582 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779606 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779627 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779646 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779666 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779686 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779706 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779725 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779744 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779764 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779784 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779808 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779827 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779848 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779868 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779887 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779906 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779926 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779946 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.779964 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.779983 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.780002 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.780021 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.780040 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: P() Physical: () I0415 23:19:50.780072 127701190978688 maxtext_utils.py:1623] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: P() Physical: () I0415 23:19:50.780112 127701190978688 maxtext_utils.py:1623] params/params/decoder/decoder_norm/scale Shape: float32[2048] Logical: P('norm',) Physical: (None,) I0415 23:19:50.780155 127701190978688 maxtext_utils.py:1623] params/params/decoder/layers/mlp/wi_0/kernel Shape: float32[2048,16,7168] Logical: P('embed', 'layers', 'mlp') Physical: ('fsdp', None, None) I0415 23:19:50.780182 127701190978688 maxtext_utils.py:1623] params/params/decoder/layers/mlp/wi_1/kernel Shape: float32[2048,16,7168] Logical: P('embed', 'layers', 'mlp') Physical: ('fsdp', None, None) I0415 23:19:50.780216 127701190978688 maxtext_utils.py:1623] params/params/decoder/layers/mlp/wo/kernel Shape: float32[7168,16,2048] Logical: P('mlp', 'layers', 'embed') Physical: (None, None, 'fsdp') I0415 23:19:50.780249 127701190978688 maxtext_utils.py:1623] params/params/decoder/layers/post_self_attention_layer_norm/scale Shape: float32[2048,16] Logical: P('norm', 'layers') Physical: (None, None) I0415 23:19:50.780272 127701190978688 maxtext_utils.py:1623] params/params/decoder/layers/pre_self_attention_layer_norm/scale Shape: float32[2048,16] Logical: P('norm', 'layers') Physical: (None, None) I0415 23:19:50.780307 127701190978688 maxtext_utils.py:1623] 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 23:19:50.780342 127701190978688 maxtext_utils.py:1623] 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 23:19:50.780368 127701190978688 maxtext_utils.py:1623] 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 23:19:50.780392 127701190978688 maxtext_utils.py:1623] 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 23:19:50.780424 127701190978688 maxtext_utils.py:1623] params/params/decoder/logits_dense/kernel Shape: float32[2048,32000] Logical: P('embed', 'vocab') Physical: ('fsdp', None) I0415 23:19:50.780456 127701190978688 maxtext_utils.py:1623] params/params/token_embedder/embedding Shape: float32[32000,2048] Logical: P('vocab', 'embed') Physical: (None, 'fsdp') I0415 23:19:52.237690 127701190978688 train.py:155] train/xent Logical: float32[8,2048]............................................. ('activation_embed_and_logits_batch', 'activation_length'). I0415 23:19:52.237768 127701190978688 train.py:155] train/xent Physical: float32[8,2048]............................................. ('fsdp', None). I0415 23:19:52.252097 127701190978688 train.py:162] train/z_loss Logical: float32[8,2048]............................................. ('activation_embed_and_logits_batch', 'activation_length'). I0415 23:19:52.252146 127701190978688 train.py:162] train/z_loss Physical: float32[8,2048]............................................. ('fsdp', None). W0415 23:19:52.736221 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 23:19:52.921315 127701190978688 max_utils.py:791] Total memory size: 3.5 GB, Output size: 1.5 GB, Temp size: 1.9 GB, Argument size: 1.5 GB, Host temp size: 0.0 GB. I0415 23:19:52.921888 127701190978688 max_utils.py:194] tensorboardX not available; using no-op SummaryWriter. I0415 23:19:52.922108 127701190978688 metric_logger.py:289] number parameters: 1.105 billion W0415 23:19:57.183526 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 23:19:57.396159 127701190978688 checkpointing.py:772] Waiting for step 0 to finish before checkpoint... I0415 23:19:57.509158 127701190978688 checkpointing.py:776] Waited 0.11298537254333496 seconds for step 0 to finish before starting checkpointing. I0415 23:19:57.509633 127701190978688 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0415 23:19:57.509784 127701190978688 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 0 I0415 23:19:57.510160 127701190978688 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0. I0415 23:19:57.584296 127701190978688 signaling_client.py:373] Using ThreadSafeKeyValueSignalingClient I0415 23:19:57.673940 127701190978688 jax_array_handlers.py:347] Scheduling D2H of 81 prioritized jax.Array. I0415 23:19:57.674034 127587760211520 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0 I0415 23:19:57.674345 127701190978688 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 W0415 23:19:58.344208 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 23:19:58.354388 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 23:19:58.363152 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 23:19:58.388257 127587749725760 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0/items W0415 23:19:58.388743 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 23:19:58.396950 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 23:19:58.402482 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 23:19:58.408639 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 23:19:58.413533 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 23:19:58.521065 127587709879872 checkpoint.py:188] Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776295198272665666, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1776295198272665666, "commit_timestamp_nsecs": null, "custom_metadata": {}} to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0415 23:19:59.013299 3365268 google_auth_provider.cc:149] Using credentials at ~/.config/gcloud/application_default_credentials.json I0415 23:19:59.013352 3365268 google_auth_provider.cc:156] Using OAuth2 AuthProvider I0415 23:19:59.028148 127701190978688 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 1.354915s I0415 23:19:59.028676 127701190978688 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 8.555 GiB/s (total gbytes: 12.3 GiB) (time elapsed: a second) (per-host) I0415 23:19:59.028785 127701190978688 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.442822s (batch_requests_ready=0.081087s, total_serialization_initiated=1.361339s, others=0.000397s) I0415 23:19:59.028869 127701190978688 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 1.444120s (all_items=0.000022s, per_item={'items': '0.00002170'}, temp_paths=1.444098) I0415 23:19:59.038308 127587770697280 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0415 23:19:59.038458 127701190978688 async_checkpointer.py:561] Finished blocking save. Time taken: 1.528627s. Continuing background save to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0. I0415 23:19:59.038714 127701190978688 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=0] Starting CheckpointManager Save Finalize thread=save_finalize I0415 23:19:59.038958 127587615508032 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0415 23:19:59.039107 127701190978688 standard_logger.py:34] {'step': 0, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776295197.5096192, 'wait_for_prev_duration_secs': 4.1961669921875e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776295197.5098045, 'checkpointer_blocking_duration_secs': 1.528801679611206, 'get_old_steps_start_time': 1776295199.0386348, 'get_old_steps_duration_secs': 3.0994415283203125e-05, 'checkpoint_manager_blocking_start_time': 1776295197.5095298, 'checkpoint_manager_blocking_duration_secs': 1.5295405387878418} I0415 23:19:59.039224 127701190978688 checkpointing.py:408] Started an asynchronous checkpoint save for step 0 I0415 23:19:59.039275 127701190978688 max_utils.py:750] Memstats: After params initialized: I0415 23:19:59.039330 127701190978688 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_0(process=0,(0,0,0,0)) I0415 23:19:59.039359 127701190978688 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_1(process=0,(1,0,0,0)) I0415 23:19:59.039378 127701190978688 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_2(process=0,(0,1,0,0)) I0415 23:19:59.039397 127701190978688 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_3(process=0,(1,1,0,0)) I0415 23:19:59.039414 127701190978688 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_4(process=0,(0,2,0,0)) I0415 23:19:59.039432 127701190978688 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_5(process=0,(1,2,0,0)) I0415 23:19:59.039451 127701190978688 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_6(process=0,(0,3,0,0)) I0415 23:19:59.039467 127701190978688 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_7(process=0,(1,3,0,0)) W0415 23:19:59.045706 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 23:19:59.051881 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. W0415 23:19:59.056499 3362481 pjrt_executable.cc:642] Assume version compatibility. PjRt-IFRT does not track XLA executable versions. I0415 23:19:59.372128 127701190978688 metric_logger.py:185] completed step: 0, seconds: 4.474, TFLOP/s/device: 3.037, Tokens/s/device: 457.800, total_weights: 13328, loss: 10.844 I0415 23:19:59.373153 127701190978688 metric_logger.py:269] To see full metrics 'tensorboard --logdir=gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/tensorboard/' I0415 23:19:59.479840 127701190978688 metric_logger.py:185] completed step: 1, seconds: 1.973, TFLOP/s/device: 6.886, Tokens/s/device: 1037.950, total_weights: 12332, loss: 10.871 I0415 23:19:59.590301 127701190978688 metric_logger.py:185] completed step: 2, seconds: 0.011, TFLOP/s/device: 1251.693, Tokens/s/device: 188668.816, total_weights: 15161, loss: 10.871 I0415 23:19:59.700519 127701190978688 metric_logger.py:185] completed step: 3, seconds: 0.109, TFLOP/s/device: 125.195, Tokens/s/device: 18870.706, total_weights: 13327, loss: 10.857 I0415 23:19:59.997083 127587657451072 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 81 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0/items/array_metadatas/process_0 I0415 23:20:23.543423 127701190978688 metric_logger.py:185] completed step: 4, seconds: 0.111, TFLOP/s/device: 122.914, Tokens/s/device: 18526.895, total_weights: 11939, loss: 10.861 I0415 23:20:23.554070 127701190978688 metric_logger.py:185] completed step: 5, seconds: 0.110, TFLOP/s/device: 123.287, Tokens/s/device: 18583.212, total_weights: 15502, loss: 10.882 I0415 23:20:23.662772 127701190978688 metric_logger.py:185] completed step: 6, seconds: 23.843, TFLOP/s/device: 0.570, Tokens/s/device: 85.895, total_weights: 13864, loss: 10.886 I0415 23:20:23.772887 127701190978688 metric_logger.py:185] completed step: 7, seconds: 0.007, TFLOP/s/device: 1917.732, Tokens/s/device: 289061.397, total_weights: 12988, loss: 10.829 I0415 23:20:23.883594 127701190978688 metric_logger.py:185] completed step: 8, seconds: 0.109, TFLOP/s/device: 124.439, Tokens/s/device: 18756.812, total_weights: 13820, loss: 10.848 I0415 23:20:23.890149 127701190978688 checkpointing.py:772] Waiting for step 10 to finish before checkpoint... I0415 23:20:24.104357 127701190978688 checkpointing.py:776] Waited 0.21417546272277832 seconds for step 10 to finish before starting checkpointing. I0415 23:20:24.104830 127701190978688 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=0][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0415 23:20:26.983952 127587646965312 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 27.954582s (commit=27.496183s, array_metadata_write=0.458398s) I0415 23:20:26.985142 127587770697280 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 429.881 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 29 seconds) (per-host) I0415 23:20:26.985250 127587770697280 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 27.946729s. I0415 23:20:27.219876 127587770697280 checkpoint.py:228] Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776295198272665666, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0415 23:20:27.408278 127587770697280 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0415 23:20:27.561951 127587709879872 checkpoint.py:247] Updated Metadata={'item_handlers': {'items': 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776295198272665666, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0415 23:20:27.805839 127587770697280 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0415 23:20:27.806484 127587770697280 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.543141s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0/items I0415 23:20:27.807420 127587770697280 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0/items I0415 23:20:28.050282 127587770697280 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0 I0415 23:20:28.722709 127587770697280 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0`. I0415 23:20:28.723361 127587770697280 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 31.213541s. directory=gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/0 I0415 23:20:28.723446 127587770697280 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 29.684931s. I0415 23:20:28.723632 127587615508032 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0415 23:20:28.723739 127587615508032 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0415 23:20:28.723808 127587615508032 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is syncing with other hosts... I0415 23:20:28.723862 127587615508032 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is done on all hosts. I0415 23:20:28.723957 127701190978688 checkpoint_manager.py:2006] [process=0][thread=MainThread][step=0][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=0. W0415 23:20:28.724028 127701190978688 checkpoint_manager.py:1441] Waiting for previous save to complete took 4.619221 seconds. If this number is high, consider checkpointing less frequently. I0415 23:20:28.725244 127701190978688 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 10 I0415 23:20:28.725539 127701190978688 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10. I0415 23:20:28.909918 127701190978688 jax_array_handlers.py:347] Scheduling D2H of 81 prioritized jax.Array. I0415 23:20:28.910011 127701190978688 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 23:20:28.912611 127587615508032 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10 I0415 23:20:29.639226 127587625993792 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10/items I0415 23:20:38.816740 127701190978688 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 9.907252s I0415 23:20:38.823403 127701190978688 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1.233 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 10 seconds) (per-host) I0415 23:20:38.823466 127701190978688 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 10.013341s (batch_requests_ready=0.093693s, total_serialization_initiated=9.913085s, others=0.006563s) I0415 23:20:38.823545 127701190978688 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 10.013887s (all_items=0.000022s, per_item={'items': '0.00002193'}, temp_paths=10.013865) I0415 23:20:38.825176 127587678422592 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0415 23:20:38.825266 127701190978688 async_checkpointer.py:561] Finished blocking save. Time taken: 10.099970s. Continuing background save to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10. I0415 23:20:38.825466 127701190978688 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=10] Starting CheckpointManager Save Finalize thread=save_finalize I0415 23:20:38.825659 127587657451072 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0415 23:20:38.825763 127701190978688 standard_logger.py:34] {'step': 10, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776295224.104783, 'wait_for_prev_duration_secs': 4.619221210479736, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776295228.725268, 'checkpointer_blocking_duration_secs': 10.100087642669678, 'get_old_steps_start_time': 1776295238.8253806, 'get_old_steps_duration_secs': 4.315376281738281e-05, 'checkpoint_manager_blocking_start_time': 1776295224.1046965, 'checkpoint_manager_blocking_duration_secs': 14.721036911010742} I0415 23:20:38.825887 127701190978688 checkpointing.py:408] Started an asynchronous checkpoint save for step 10 I0415 23:20:38.826366 127701190978688 metric_logger.py:185] completed step: 9, seconds: 0.110, TFLOP/s/device: 123.598, Tokens/s/device: 18630.037, total_weights: 12300, loss: 10.857 I0415 23:20:38.835754 127701190978688 metric_logger.py:185] completed step: 10, seconds: 0.111, TFLOP/s/device: 122.295, Tokens/s/device: 18433.677, total_weights: 13436, loss: 10.876 I0415 23:20:38.943006 127701190978688 metric_logger.py:185] completed step: 11, seconds: 14.943, TFLOP/s/device: 0.909, Tokens/s/device: 137.055, total_weights: 13817, loss: 10.853 I0415 23:20:39.053348 127701190978688 metric_logger.py:185] completed step: 12, seconds: 0.008, TFLOP/s/device: 1672.880, Tokens/s/device: 252154.642, total_weights: 12076, loss: 10.852 I0415 23:20:39.163704 127701190978688 metric_logger.py:185] completed step: 13, seconds: 0.108, TFLOP/s/device: 125.279, Tokens/s/device: 18883.408, total_weights: 11048, loss: 10.864 I0415 23:20:39.273915 127701190978688 metric_logger.py:185] completed step: 14, seconds: 0.110, TFLOP/s/device: 123.224, Tokens/s/device: 18573.605, total_weights: 14267, loss: 10.834 I0415 23:20:39.384347 127701190978688 metric_logger.py:185] completed step: 15, seconds: 0.110, TFLOP/s/device: 123.062, Tokens/s/device: 18549.212, total_weights: 13004, loss: 10.832 I0415 23:20:39.494599 127701190978688 metric_logger.py:185] completed step: 16, seconds: 0.110, TFLOP/s/device: 123.299, Tokens/s/device: 18584.898, total_weights: 13811, loss: 10.867 I0415 23:20:39.565949 127587688908352 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 81 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10/items/array_metadatas/process_0 I0415 23:20:39.605121 127701190978688 metric_logger.py:185] completed step: 17, seconds: 0.110, TFLOP/s/device: 123.013, Tokens/s/device: 18541.823, total_weights: 13709, loss: 10.850 I0415 23:20:39.719389 127701190978688 metric_logger.py:185] completed step: 18, seconds: 0.110, TFLOP/s/device: 123.458, Tokens/s/device: 18608.877, total_weights: 12831, loss: 10.839 I0415 23:20:39.825460 127701190978688 checkpointing.py:772] Waiting for step 19 to finish before checkpoint... I0415 23:20:39.826967 127701190978688 checkpointing.py:776] Waited 0.0015239715576171875 seconds for step 19 to finish before starting checkpointing. I0415 23:20:39.827404 127701190978688 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=10][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0415 23:20:58.429606 127587615508032 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 19.605415s (commit=19.161349s, array_metadata_write=0.444066s) I0415 23:20:58.430602 127587678422592 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 426.668 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 29 seconds) (per-host) I0415 23:20:58.430708 127587678422592 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 19.605480s. I0415 23:20:58.864134 127587678422592 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0415 23:20:59.317289 127587678422592 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0415 23:20:59.318006 127587678422592 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.601541s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10/items I0415 23:20:59.318651 127587678422592 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10/items I0415 23:20:59.561346 127587678422592 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10 I0415 23:21:00.235229 127587678422592 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10`. I0415 23:21:00.235822 127587678422592 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 31.510540s. directory=gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/10 I0415 23:21:00.235899 127587678422592 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 21.410674s. I0415 23:21:00.236055 127587657451072 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0415 23:21:00.236153 127587657451072 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0415 23:21:00.236205 127587657451072 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=10] CheckpointManager Save Finalize is syncing with other hosts... I0415 23:21:00.236246 127587657451072 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=10] CheckpointManager Save Finalize is done on all hosts. I0415 23:21:00.236450 127701190978688 checkpoint_manager.py:2006] [process=0][thread=MainThread][step=10][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=10. W0415 23:21:00.236818 127701190978688 checkpoint_manager.py:1441] Waiting for previous save to complete took 20.409431 seconds. If this number is high, consider checkpointing less frequently. I0415 23:21:00.237030 127701190978688 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 19 I0415 23:21:00.237329 127701190978688 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19. I0415 23:21:00.415434 127587657451072 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19 I0415 23:21:00.423717 127701190978688 jax_array_handlers.py:347] Scheduling D2H of 81 prioritized jax.Array. I0415 23:21:00.423808 127701190978688 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 23:21:01.187728 127587615508032 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19/items I0415 23:21:08.700122 127701190978688 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 8.276812s I0415 23:21:08.707678 127701190978688 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1.472 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 8 seconds) (per-host) I0415 23:21:08.707801 127701190978688 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 8.383958s (batch_requests_ready=0.093964s, total_serialization_initiated=8.282513s, others=0.007480s) I0415 23:21:08.707925 127701190978688 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 8.384509s (all_items=0.000011s, per_item={'items': '0.00001073'}, temp_paths=8.384498) I0415 23:21:08.709328 127587688908352 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0415 23:21:08.709415 127701190978688 async_checkpointer.py:561] Finished blocking save. Time taken: 8.472327s. Continuing background save to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19. I0415 23:21:08.709626 127701190978688 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=19] Starting CheckpointManager Save Finalize thread=save_finalize I0415 23:21:08.709852 127587657451072 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0415 23:21:08.709947 127701190978688 standard_logger.py:34] {'step': 19, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776295239.827373, 'wait_for_prev_duration_secs': 20.409430503845215, 'time_between_consecutive_saves_sec': 11.103494644165039, 'checkpointer_blocking_start_time': 1776295260.237067, 'checkpointer_blocking_duration_secs': 8.472462892532349, 'get_old_steps_start_time': 1776295268.7095485, 'get_old_steps_duration_secs': 2.8371810913085938e-05, 'checkpoint_manager_blocking_start_time': 1776295239.8273382, 'checkpoint_manager_blocking_duration_secs': 28.88258123397827} I0415 23:21:08.710106 127701190978688 checkpointing.py:408] Started an asynchronous checkpoint save for step 19 I0415 23:21:08.710148 127701190978688 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=19][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0415 23:21:09.414724 127587126871616 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 81 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19/items/array_metadatas/process_0 I0415 23:21:33.780744 127587472901696 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 25.072432s (commit=24.646496s, array_metadata_write=0.425936s) I0415 23:21:33.781790 127587688908352 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 377.731 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 33 seconds) (per-host) I0415 23:21:33.781887 127587688908352 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 25.072426s. I0415 23:21:34.208529 127587688908352 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0415 23:21:34.608762 127587688908352 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0415 23:21:34.609437 127587688908352 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.549967s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19/items I0415 23:21:34.610076 127587688908352 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19/items I0415 23:21:34.850922 127587688908352 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19 I0415 23:21:35.538405 127587688908352 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19`. I0415 23:21:35.539036 127587688908352 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 35.301958s. directory=gs://wanglance-maxtext/linen_ckpt_feat_migrate_nnx_utils_20260415_230805/linen_feat_migrate_nnx_utils_20260415_230805_05_fp8/checkpoints/19 I0415 23:21:35.539110 127587688908352 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 26.829652s. I0415 23:21:35.539232 127587657451072 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0415 23:21:35.539280 127587657451072 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0415 23:21:35.539330 127587657451072 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=19] CheckpointManager Save Finalize is syncing with other hosts... I0415 23:21:35.539365 127587657451072 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=19] CheckpointManager Save Finalize is done on all hosts. I0415 23:21:35.539898 127701190978688 checkpoint_manager.py:2006] [process=0][thread=MainThread][step=19][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=19. I0415 23:21:35.540026 127701190978688 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0415 23:21:35.540668 127701190978688 metric_logger.py:185] completed step: 19, seconds: 0.111, TFLOP/s/device: 122.790, Tokens/s/device: 18508.310, total_weights: 13717, loss: 10.851 Per train step: Total TFLOPs: 13.59 split as 93.93% learnable weight flops and 6.07% attention flops