feat/nnx-trainstate-and-training-loop| Metric | Linen 29bbe47ed | NNX 29bbe47ed | Diff (NNX − Linen) |
|---|---|---|---|
| Parameters | 1.105 billion | 1.104 billion | — |
| Final loss | 10.8510 | 10.6390 | -0.212 |
| TFLOP/s | 121.624 | 122.873 | +1.249 |
| Tok/s | 18332.5 | 18520.7 | +188.168 |
| Avg s/step | 2.141 | 2.882 | +0.741 |
| Memory % | 5.15 | 5.15 | 0 |
| JAX | 0.8.1 | 0.8.1 | — |
Diff = NNX value − Linen value. Green = NNX improved. Red = NNX regressed.
2026-04-21 16:30:33.630953: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered WARNING: All log messages before absl::InitializeLog() is called are written to STDERR E0000 00:00:1776789033.641007 1603767 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered E0000 00:00:1776789033.644146 1603767 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered W0000 00:00:1776789033.653897 1603767 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776789033.653908 1603767 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776789033.653910 1603767 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776789033.653911 1603767 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'} `rope_scaling`'s factor field must be a float >= 1, got 40 `rope_scaling`'s beta_fast field must be a float, got 32 `rope_scaling`'s beta_slow field must be a float, got 1 Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'} Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'} Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'} 2026-04-21 16:30:37.290169: 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) I0421 16:30:37.425600 130724548144256 max_utils.py:238] Skipping jax distributed system due to skip_jax_distributed_system=True flag. I0421 16:31:06.437083 130724548144256 max_utils.py:800] System Information: Jax Version: 0.8.1 I0421 16:31:06.437268 130724548144256 max_utils.py:801] System Information: Jaxlib Version: 0.8.1 I0421 16:31:06.437308 130724548144256 max_utils.py:802] System Information: Jax Backend: PJRT C API TFRT TPU v6 lite Built on Nov 12 2025 14:16:36 (1762985796) cl/831091709 I0421 16:31:06.437338 130724548144256 train_utils.py:378] WARNING: 'dataset_path' might be pointing your local file system I0421 16:31:06.437450 130724548144256 train.py:806] [DECOUPLED NO-OP] skipping cloud diagnostics wrapper. I0421 16:31:07.066428 130724548144256 maxtext_utils.py:1732] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0421 16:31:07.066657 130724548144256 maxtext_utils.py:1732] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0421 16:31:07.066825 130724548144256 checkpointing.py:688] Setting up checkpoint logger... I0421 16:31:07.066864 130724548144256 checkpointing.py:234] Creating checkpoint manager with ocdbt=True and zarr3=True I0421 16:31:07.066898 130724548144256 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0421 16:31:07.067368 130724548144256 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 0x76e461fa20c0>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0421 16:31:09.518857 130724548144256 checkpointing.py:266] Enabling policy for fixed interval checkpointing. I0421 16:31:09.519082 130724548144256 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 0x76e2b9396270>}, handler_registry=None I0421 16:31:09.519310 130724548144256 composite_checkpoint_handler.py:237] Deferred registration for item: "items". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x76e2b9396270>` 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`. I0421 16:31:09.519350 130724548144256 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x76e2b9396840>` 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`. I0421 16:31:09.519378 130724548144256 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 0x76e2b9396270>, ('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x76e2b9396270>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x76e2b9396840>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x76e2b9396840>}). I0421 16:31:09.519638 130724548144256 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.28 I0421 16:31:09.519688 130724548144256 async_checkpointer.py:177] [process=0][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>.<lambda> at 0x76e2b93b8ae0> timeout: 600 secs and primary_host=0 for async checkpoint writes I0421 16:31:09.655138 130724548144256 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints I0421 16:31:09.655378 130724548144256 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_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x76e30bb3c890> I0421 16:31:09.655457 130724548144256 checkpointing.py:302] Checkpoint manager created! I0421 16:31:09.733214 130724548144256 dataset_info.py:707] Load dataset info from tests/assets/local_datasets/c4_en_dataset_minimal/c4/en/3.1.0 I0421 16:31:09.736197 130724548144256 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. I0421 16:31:09.790675 130724548144256 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 I0421 16:31:09.813524 130724548144256 tokenizer.py:245] Tokenizer path: src/maxtext/assets/tokenizers/tokenizer.llama2 I0421 16:31:09.813591 130724548144256 tokenizer.py:187] Loading sentencepiece tokenizer: src/maxtext/assets/tokenizers/tokenizer.llama2 I0421 16:31:10.574517 130724548144256 nnx_wrappers.py:437] Unknown Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0421 16:31:10.574619 130724548144256 nnx_wrappers.py:437] Unknown Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0421 16:31:10.878263 130724548144256 attentions.py:1088] attentions/inputs_q Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0421 16:31:10.878351 130724548144256 attentions.py:1088] attentions/inputs_q Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0421 16:31:10.891457 130724548144256 attentions.py:1089] attentions/inputs_kv Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0421 16:31:10.891505 130724548144256 attentions.py:1089] attentions/inputs_kv Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0421 16:31:10.941151 130724548144256 attentions.py:1154] attentions/query Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0421 16:31:10.941227 130724548144256 attentions.py:1154] attentions/query Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0421 16:31:10.954130 130724548144256 attentions.py:1155] attentions/key Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0421 16:31:10.954183 130724548144256 attentions.py:1155] attentions/key Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0421 16:31:10.967084 130724548144256 attentions.py:1156] attentions/value Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0421 16:31:10.967133 130724548144256 attentions.py:1156] attentions/value Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0421 16:31:10.987031 130724548144256 attentions.py:1198] attentions/out Logical: bfloat16[8,2048,16,128]..................................... ('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv'). I0421 16:31:10.987105 130724548144256 attentions.py:1198] attentions/out Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0421 16:31:11.052456 130724548144256 linears.py:525] linears/x Logical: bfloat16[8,2048,7168]....................................... ('activation_batch', 'activation_length', 'activation_mlp'). I0421 16:31:11.052545 130724548144256 linears.py:525] linears/x Physical: bfloat16[8,2048,7168]....................................... ('fsdp', None, None). I0421 16:31:11.716440 130724548144256 checkpointing.py:578] checkpoint manager exists so trying to load this run's existing checkpoint I0421 16:31:11.716547 130724548144256 checkpointing.py:676] No existing checkpoints found, not restoring checkpoint. [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 I0421 16:31:12.790654 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.790766 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.790804 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.790834 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.790859 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.790883 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_0/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.790908 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.790931 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.790954 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.790981 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791003 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791024 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wi_1/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791059 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791083 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791106 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791127 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791148 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791168 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/mlp/wo/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791189 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791209 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791229 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791248 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791268 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791288 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/key/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791308 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791328 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791348 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791368 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791387 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791407 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/out/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791428 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791448 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791468 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791488 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791507 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791529 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/query/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791549 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/input_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791568 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/input_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791589 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/kernel_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791612 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/kernel_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791631 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/output_grad_amax_history Shape: float32[16,1024] Logical: PartitionSpec() Physical: () I0421 16:31:12.791651 130724548144256 maxtext_utils.py:1835] params/_overwrite_with_gradient/decoder/layers/self_attention/value/Fp8DirectDotGeneralOp_0/output_grad_scale Shape: float32[16,1] Logical: PartitionSpec() Physical: () I0421 16:31:12.791691 130724548144256 maxtext_utils.py:1835] params/params/decoder/decoder_norm/scale Shape: float32[2048] Logical: PartitionSpec('norm',) Physical: (None,) I0421 16:31:12.791733 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/mlp/wi_0/kernel Shape: float32[2048,16,7168] Logical: PartitionSpec('embed', 'layers', 'mlp') Physical: ('fsdp', None, None) I0421 16:31:12.791759 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/mlp/wi_1/kernel Shape: float32[2048,16,7168] Logical: PartitionSpec('embed', 'layers', 'mlp') Physical: ('fsdp', None, None) I0421 16:31:12.791795 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/mlp/wo/kernel Shape: float32[7168,16,2048] Logical: PartitionSpec('mlp', 'layers', 'embed') Physical: (None, None, 'fsdp') I0421 16:31:12.791827 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/post_self_attention_layer_norm/scale Shape: float32[2048,16] Logical: PartitionSpec('norm', 'layers') Physical: (None, None) I0421 16:31:12.791851 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/pre_self_attention_layer_norm/scale Shape: float32[2048,16] Logical: PartitionSpec('norm', 'layers') Physical: (None, None) I0421 16:31:12.791886 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/self_attention/key/kernel Shape: float32[2048,16,16,128] Logical: PartitionSpec('embed', 'layers', 'kv_heads', 'kv_head_dim') Physical: ('fsdp', None, None, None) I0421 16:31:12.791923 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/self_attention/out/kernel Shape: float32[16,16,128,2048] Logical: PartitionSpec('heads', 'layers', 'kv', 'embed') Physical: (None, None, None, 'fsdp') I0421 16:31:12.791949 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/self_attention/query/kernel Shape: float32[2048,16,16,128] Logical: PartitionSpec('embed', 'layers', 'q_heads', 'kv') Physical: ('fsdp', None, None, None) I0421 16:31:12.791977 130724548144256 maxtext_utils.py:1835] params/params/decoder/layers/self_attention/value/kernel Shape: float32[2048,16,16,128] Logical: PartitionSpec('embed', 'layers', 'kv_heads', 'kv_head_dim') Physical: ('fsdp', None, None, None) I0421 16:31:12.792011 130724548144256 maxtext_utils.py:1835] params/params/decoder/logits_dense/kernel Shape: float32[2048,32000] Logical: PartitionSpec('embed_vocab', 'vocab') Physical: ('fsdp', None) I0421 16:31:12.792052 130724548144256 maxtext_utils.py:1835] params/params/token_embedder/embedding Shape: float32[32000,2048] Logical: PartitionSpec('vocab', 'embed_vocab') Physical: (None, 'fsdp') I0421 16:31:14.222369 130724548144256 train.py:157] train/xent Logical: float32[8,2048]............................................. ('activation_embed_and_logits_batch', 'activation_length'). I0421 16:31:14.222457 130724548144256 train.py:157] train/xent Physical: float32[8,2048]............................................. ('fsdp', None). I0421 16:31:14.234538 130724548144256 train.py:164] train/z_loss Logical: float32[8,2048]............................................. ('activation_embed_and_logits_batch', 'activation_length'). I0421 16:31:14.234589 130724548144256 train.py:164] train/z_loss Physical: float32[8,2048]............................................. ('fsdp', None). I0421 16:31:15.483569 130724548144256 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. I0421 16:31:15.484218 130724548144256 max_utils.py:194] tensorboardX not available; using no-op SummaryWriter. I0421 16:31:15.484433 130724548144256 metric_logger.py:301] number parameters: 1.105 billion I0421 16:31:21.885712 130724548144256 checkpointing.py:794] Waiting for step 0 to finish before checkpoint... I0421 16:31:22.002175 130724548144256 checkpointing.py:798] Waited 0.1164391040802002 seconds for step 0 to finish before starting checkpointing. I0421 16:31:22.002682 130724548144256 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0421 16:31:22.002842 130724548144256 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 0 I0421 16:31:22.003249 130724548144256 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0. I0421 16:31:22.085035 130724548144256 signaling_client.py:373] Using ThreadSafeKeyValueSignalingClient I0421 16:31:22.180531 130724548144256 jax_array_handlers.py:347] Scheduling D2H of 81 prioritized jax.Array. I0421 16:31:22.180633 130724548144256 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 I0421 16:31:22.188678 130633783313984 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0 I0421 16:31:22.880928 130633772828224 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/items I0421 16:31:23.023537 130633741370944 checkpoint.py:188] Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776789082779483016, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1776789082779483016, "commit_timestamp_nsecs": null, "custom_metadata": {}} to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0421 16:31:23.479125 1606169 google_auth_provider.cc:149] Using credentials at ~/.config/gcloud/application_default_credentials.json I0421 16:31:23.479185 1606169 google_auth_provider.cc:156] Using OAuth2 AuthProvider I0421 16:31:23.993320 130724548144256 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 1.813647s I0421 16:31:24.000263 130724548144256 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 6.449 GiB/s (total gbytes: 12.3 GiB) (time elapsed: a second) (per-host) I0421 16:31:24.000332 130724548144256 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.913869s (batch_requests_ready=0.087050s, total_serialization_initiated=1.819984s, others=0.006835s) I0421 16:31:24.000420 130724548144256 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 1.914842s (all_items=0.000020s, per_item={'items': '0.00001955'}, temp_paths=1.914822) I0421 16:31:24.002124 130633793799744 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0421 16:31:24.002231 130724548144256 async_checkpointer.py:561] Finished blocking save. Time taken: 1.999349s. Continuing background save to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0. I0421 16:31:24.002418 130724548144256 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=0] Starting CheckpointManager Save Finalize thread=save_finalize I0421 16:31:24.002593 130633678456384 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0421 16:31:24.002694 130724548144256 standard_logger.py:34] {'step': 0, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776789082.0026703, 'wait_for_prev_duration_secs': 4.029273986816406e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776789082.0028636, 'checkpointer_blocking_duration_secs': 1.9994807243347168, 'get_old_steps_start_time': 1776789084.002362, 'get_old_steps_duration_secs': 2.193450927734375e-05, 'checkpoint_manager_blocking_start_time': 1776789082.0025892, 'checkpoint_manager_blocking_duration_secs': 2.000077962875366} I0421 16:31:24.002808 130724548144256 checkpointing.py:409] Started an asynchronous checkpoint save for step 0 I0421 16:31:24.002847 130724548144256 max_utils.py:750] Memstats: After params initialized: I0421 16:31:24.002891 130724548144256 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_0(process=0,(0,0,0,0)) I0421 16:31:24.002916 130724548144256 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_1(process=0,(1,0,0,0)) I0421 16:31:24.002934 130724548144256 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_2(process=0,(0,1,0,0)) I0421 16:31:24.002953 130724548144256 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_3(process=0,(1,1,0,0)) I0421 16:31:24.002971 130724548144256 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_4(process=0,(0,2,0,0)) I0421 16:31:24.002988 130724548144256 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_5(process=0,(1,2,0,0)) I0421 16:31:24.003004 130724548144256 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_6(process=0,(0,3,0,0)) I0421 16:31:24.003020 130724548144256 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_7(process=0,(1,3,0,0)) I0421 16:31:24.417341 130724548144256 metric_logger.py:196] completed step: 0, seconds: 6.401, TFLOP/s/device: 2.123, Tokens/s/device: 319.963, total_weights: 13328, loss: 10.844, lm_loss: 10.844, perplexity: 51212.484 I0421 16:31:24.418372 130724548144256 metric_logger.py:281] To see full metrics 'tensorboard --logdir=gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/tensorboard/' I0421 16:31:24.526331 130724548144256 metric_logger.py:196] completed step: 1, seconds: 2.528, TFLOP/s/device: 5.374, Tokens/s/device: 810.008, total_weights: 12332, loss: 10.871, lm_loss: 10.871, perplexity: 52634.344 I0421 16:31:24.638418 130724548144256 metric_logger.py:196] completed step: 2, seconds: 0.011, TFLOP/s/device: 1202.614, Tokens/s/device: 181271.021, total_weights: 15161, loss: 10.871, lm_loss: 10.871, perplexity: 52617.816 I0421 16:31:24.750341 130724548144256 metric_logger.py:196] completed step: 3, seconds: 0.111, TFLOP/s/device: 122.434, Tokens/s/device: 18454.607, total_weights: 13327, loss: 10.857, lm_loss: 10.857, perplexity: 51888.457 I0421 16:31:24.856770 130633751856704 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 81 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/items/array_metadatas/process_0 I0421 16:31:42.514036 130724548144256 metric_logger.py:196] completed step: 4, seconds: 0.111, TFLOP/s/device: 122.171, Tokens/s/device: 18414.948, total_weights: 11939, loss: 10.861, lm_loss: 10.861, perplexity: 52094.562 I0421 16:31:42.526863 130724548144256 metric_logger.py:196] completed step: 5, seconds: 0.112, TFLOP/s/device: 121.601, Tokens/s/device: 18329.082, total_weights: 15502, loss: 10.882, lm_loss: 10.882, perplexity: 53230.324 I0421 16:31:42.637370 130724548144256 metric_logger.py:196] completed step: 6, seconds: 17.766, TFLOP/s/device: 0.765, Tokens/s/device: 115.279, total_weights: 13864, loss: 10.886, lm_loss: 10.886, perplexity: 53423.605 I0421 16:31:42.749301 130724548144256 metric_logger.py:196] completed step: 7, seconds: 0.007, TFLOP/s/device: 1835.850, Tokens/s/device: 276719.362, total_weights: 12988, loss: 10.829, lm_loss: 10.829, perplexity: 50459.859 I0421 16:31:42.861098 130724548144256 metric_logger.py:196] completed step: 8, seconds: 0.112, TFLOP/s/device: 121.711, Tokens/s/device: 18345.665, total_weights: 13820, loss: 10.848, lm_loss: 10.848, perplexity: 51456.266 I0421 16:31:42.868289 130724548144256 checkpointing.py:794] Waiting for step 10 to finish before checkpoint... I0421 16:31:43.084928 130724548144256 checkpointing.py:798] Waited 0.21660280227661133 seconds for step 10 to finish before starting checkpointing. I0421 16:31:43.085489 130724548144256 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=0][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0421 16:31:56.642745 130633699427904 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 32.641659s (commit=32.206351s, array_metadata_write=0.435308s) I0421 16:31:56.644157 130633793799744 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 365.711 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 34 seconds) (per-host) I0421 16:31:56.644313 130633793799744 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 32.642006s. I0421 16:31:56.853263 130633793799744 checkpoint.py:228] Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776789082779483016, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0421 16:31:57.049852 130633793799744 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0421 16:31:57.243819 130633741370944 checkpoint.py:247] Updated Metadata={'item_handlers': {'items': 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776789082779483016, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0421 16:31:57.453285 130633793799744 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0421 16:31:57.453870 130633793799744 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.557165s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/items I0421 16:31:57.454657 130633793799744 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/items I0421 16:31:57.702055 130633793799744 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0 I0421 16:31:58.367659 130633793799744 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0`. I0421 16:31:58.368437 130633793799744 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 36.365557s. directory=gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0 I0421 16:31:58.368525 130633793799744 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 34.366247s. I0421 16:31:58.368816 130633678456384 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0421 16:31:58.368945 130633678456384 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0421 16:31:58.369014 130633678456384 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is syncing with other hosts... I0421 16:31:58.369088 130633678456384 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is done on all hosts. I0421 16:31:58.369427 130724548144256 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. W0421 16:31:58.369577 130724548144256 checkpoint_manager.py:1441] Waiting for previous save to complete took 15.284113 seconds. If this number is high, consider checkpointing less frequently. I0421 16:31:58.370851 130724548144256 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 10 I0421 16:31:58.371210 130724548144256 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10. I0421 16:31:58.543489 130633678456384 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10 I0421 16:31:58.553750 130724548144256 jax_array_handlers.py:347] Scheduling D2H of 81 prioritized jax.Array. I0421 16:31:58.553864 130724548144256 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 I0421 16:31:59.330957 130633667970624 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10/items I0421 16:32:01.752746 130724548144256 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 3.199444s I0421 16:32:01.759650 130724548144256 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 3.736 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 3 seconds) (per-host) I0421 16:32:01.759718 130724548144256 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 3.303768s (batch_requests_ready=0.091388s, total_serialization_initiated=3.205603s, others=0.006778s) I0421 16:32:01.759804 130724548144256 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 3.304371s (all_items=0.000024s, per_item={'items': '0.00002408'}, temp_paths=3.304347) I0421 16:32:01.761356 130633678456384 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0421 16:32:01.761547 130724548144256 async_checkpointer.py:561] Finished blocking save. Time taken: 3.390649s. Continuing background save to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10. I0421 16:32:01.761728 130724548144256 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=10] Starting CheckpointManager Save Finalize thread=save_finalize I0421 16:32:01.761900 130633793799744 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0421 16:32:01.761991 130724548144256 standard_logger.py:34] {'step': 10, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776789103.085424, 'wait_for_prev_duration_secs': 15.28411316871643, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776789118.3708758, 'checkpointer_blocking_duration_secs': 3.3907480239868164, 'get_old_steps_start_time': 1776789121.7616434, 'get_old_steps_duration_secs': 4.1484832763671875e-05, 'checkpoint_manager_blocking_start_time': 1776789103.0853198, 'checkpoint_manager_blocking_duration_secs': 18.676640033721924} I0421 16:32:01.762131 130724548144256 checkpointing.py:409] Started an asynchronous checkpoint save for step 10 I0421 16:32:01.762729 130724548144256 metric_logger.py:196] completed step: 9, seconds: 0.112, TFLOP/s/device: 121.072, Tokens/s/device: 18249.216, total_weights: 12300, loss: 10.857, lm_loss: 10.857, perplexity: 51900.086 I0421 16:32:01.771829 130724548144256 metric_logger.py:196] completed step: 10, seconds: 0.112, TFLOP/s/device: 121.548, Tokens/s/device: 18321.048, total_weights: 13436, loss: 10.876, lm_loss: 10.876, perplexity: 52905.383 I0421 16:32:01.881299 130724548144256 metric_logger.py:196] completed step: 11, seconds: 18.901, TFLOP/s/device: 0.719, Tokens/s/device: 108.353, total_weights: 13817, loss: 10.853, lm_loss: 10.853, perplexity: 51714.605 I0421 16:32:01.993318 130724548144256 metric_logger.py:196] completed step: 12, seconds: 0.008, TFLOP/s/device: 1733.052, Tokens/s/device: 261224.490, total_weights: 12076, loss: 10.852, lm_loss: 10.852, perplexity: 51622.902 I0421 16:32:02.105037 130724548144256 metric_logger.py:196] completed step: 13, seconds: 0.111, TFLOP/s/device: 122.507, Tokens/s/device: 18465.589, total_weights: 11048, loss: 10.864, lm_loss: 10.864, perplexity: 52261.953 I0421 16:32:02.216873 130724548144256 metric_logger.py:196] completed step: 14, seconds: 0.112, TFLOP/s/device: 121.539, Tokens/s/device: 18319.737, total_weights: 14267, loss: 10.834, lm_loss: 10.834, perplexity: 50738.680 I0421 16:32:02.328944 130724548144256 metric_logger.py:196] completed step: 15, seconds: 0.112, TFLOP/s/device: 121.569, Tokens/s/device: 18324.162, total_weights: 13004, loss: 10.832, lm_loss: 10.832, perplexity: 50634.074 I0421 16:32:02.440886 130724548144256 metric_logger.py:196] completed step: 16, seconds: 0.112, TFLOP/s/device: 121.471, Tokens/s/device: 18309.418, total_weights: 13811, loss: 10.867, lm_loss: 10.867, perplexity: 52422.980 I0421 16:32:02.533606 130633751856704 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 81 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10/items/array_metadatas/process_0 I0421 16:32:02.552876 130724548144256 metric_logger.py:196] completed step: 17, seconds: 0.112, TFLOP/s/device: 121.225, Tokens/s/device: 18272.336, total_weights: 13709, loss: 10.850, lm_loss: 10.850, perplexity: 51525.906 I0421 16:32:02.664551 130724548144256 metric_logger.py:196] completed step: 18, seconds: 0.112, TFLOP/s/device: 120.877, Tokens/s/device: 18219.830, total_weights: 12831, loss: 10.839, lm_loss: 10.839, perplexity: 50989.945 I0421 16:32:02.776359 130724548144256 checkpointing.py:794] Waiting for step 19 to finish before checkpoint... I0421 16:32:02.777883 130724548144256 checkpointing.py:798] Waited 0.0015439987182617188 seconds for step 19 to finish before starting checkpointing. I0421 16:32:02.779389 130724548144256 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=10][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0421 16:32:35.839029 130633667970624 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 34.078555s (commit=33.664925s, array_metadata_write=0.413630s) I0421 16:32:35.840253 130633678456384 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 338.060 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 37 seconds) (per-host) I0421 16:32:35.840374 130633678456384 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 34.078926s. I0421 16:32:36.256565 130633678456384 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0421 16:32:36.682810 130633678456384 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0421 16:32:36.683504 130633678456384 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.571593s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10/items I0421 16:32:36.684231 130633678456384 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10/items I0421 16:32:36.926614 130633678456384 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10 I0421 16:32:37.609967 130633678456384 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10`. I0421 16:32:37.610716 130633678456384 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 39.239824s. directory=gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10 I0421 16:32:37.610796 130633678456384 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 35.849337s. I0421 16:32:37.610901 130633793799744 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0421 16:32:37.610951 130633793799744 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0421 16:32:37.610997 130633793799744 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=10] CheckpointManager Save Finalize is syncing with other hosts... I0421 16:32:37.611034 130633793799744 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=10] CheckpointManager Save Finalize is done on all hosts. I0421 16:32:37.611432 130724548144256 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. W0421 16:32:37.611550 130724548144256 checkpoint_manager.py:1441] Waiting for previous save to complete took 34.832174 seconds. If this number is high, consider checkpointing less frequently. I0421 16:32:37.611764 130724548144256 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 19 I0421 16:32:37.612080 130724548144256 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19. I0421 16:32:37.780269 130633793799744 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19 I0421 16:32:37.793523 130724548144256 jax_array_handlers.py:347] Scheduling D2H of 81 prioritized jax.Array. I0421 16:32:37.793623 130724548144256 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 I0421 16:32:38.579533 130633720399424 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19/items I0421 16:32:48.921901 130724548144256 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 11.128865s I0421 16:32:48.930210 130724548144256 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1.099 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 11 seconds) (per-host) I0421 16:32:48.930323 130724548144256 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 11.233214s (batch_requests_ready=0.090372s, total_serialization_initiated=11.134646s, others=0.008196s) I0421 16:32:48.930417 130724548144256 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 11.233798s (all_items=0.000012s, per_item={'items': '0.00001168'}, temp_paths=11.233786) I0421 16:32:48.932523 130633699427904 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0421 16:32:48.932729 130724548144256 async_checkpointer.py:561] Finished blocking save. Time taken: 11.320912s. Continuing background save to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19. I0421 16:32:48.932970 130724548144256 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=19] Starting CheckpointManager Save Finalize thread=save_finalize I0421 16:32:48.933233 130633678456384 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0421 16:32:48.933391 130724548144256 standard_logger.py:34] {'step': 19, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776789122.7793615, 'wait_for_prev_duration_secs': 34.832173585891724, 'time_between_consecutive_saves_sec': 4.410256862640381, 'checkpointer_blocking_start_time': 1776789157.611787, 'checkpointer_blocking_duration_secs': 11.321102380752563, 'get_old_steps_start_time': 1776789168.9329047, 'get_old_steps_duration_secs': 2.0742416381835938e-05, 'checkpoint_manager_blocking_start_time': 1776789122.7793226, 'checkpoint_manager_blocking_duration_secs': 46.15403437614441} I0421 16:32:48.933714 130724548144256 checkpointing.py:409] Started an asynchronous checkpoint save for step 19 I0421 16:32:48.933758 130724548144256 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=19][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0421 16:32:49.688141 130633751856704 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 81 array_metadata.ArrayMetadata to gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19/items/array_metadatas/process_0 I0421 16:33:12.626008 130633793799744 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 23.695011s (commit=23.217337s, array_metadata_write=0.477674s) I0421 16:33:12.627464 130633699427904 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 361.809 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 34 seconds) (per-host) I0421 16:33:12.627612 130633699427904 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 23.694825s. I0421 16:33:13.036959 130633699427904 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0421 16:33:13.456723 130633699427904 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0421 16:33:13.457762 130633699427904 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.569226s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19/items I0421 16:33:13.458559 130633699427904 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19/items I0421 16:33:13.704083 130633699427904 atomicity.py:608] Finalizing gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19 I0421 16:33:14.400834 130633699427904 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19`. I0421 16:33:14.401614 130633699427904 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 36.789814s. directory=gs://wanglance-maxtext/linen_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/linen_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19 I0421 16:33:14.401685 130633699427904 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 25.468902s. I0421 16:33:14.401914 130633678456384 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0421 16:33:14.402072 130633678456384 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0421 16:33:14.402127 130633678456384 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=19] CheckpointManager Save Finalize is syncing with other hosts... I0421 16:33:14.402173 130633678456384 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=19] CheckpointManager Save Finalize is done on all hosts. I0421 16:33:14.402351 130724548144256 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. I0421 16:33:14.402539 130724548144256 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0421 16:33:14.403581 130724548144256 metric_logger.py:196] completed step: 19, seconds: 0.112, TFLOP/s/device: 121.624, Tokens/s/device: 18332.528, total_weights: 13717, loss: 10.851, lm_loss: 10.851, perplexity: 51591.992 Per train step: Total TFLOPs: 13.59 split as 93.93% learnable weight flops and 6.07% attention flops
2026-04-21 17:52:35.905003: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered WARNING: All log messages before absl::InitializeLog() is called are written to STDERR E0000 00:00:1776793955.914973 1677386 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered E0000 00:00:1776793955.918081 1677386 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered W0000 00:00:1776793955.927938 1677386 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776793955.927948 1677386 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776793955.927950 1677386 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. W0000 00:00:1776793955.927951 1677386 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once. Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'} `rope_scaling`'s factor field must be a float >= 1, got 40 `rope_scaling`'s beta_fast field must be a float, got 32 `rope_scaling`'s beta_slow field must be a float, got 1 Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'} Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'} Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'} 2026-04-21 17:52:39.770653: 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) I0421 17:52:39.906843 138547380145280 max_utils.py:238] Skipping jax distributed system due to skip_jax_distributed_system=True flag. I0421 17:53:07.737950 138547380145280 max_utils.py:800] System Information: Jax Version: 0.8.1 I0421 17:53:07.738089 138547380145280 max_utils.py:801] System Information: Jaxlib Version: 0.8.1 I0421 17:53:07.738140 138547380145280 max_utils.py:802] System Information: Jax Backend: PJRT C API TFRT TPU v6 lite Built on Nov 12 2025 14:16:36 (1762985796) cl/831091709 I0421 17:53:07.738166 138547380145280 train_utils.py:378] WARNING: 'dataset_path' might be pointing your local file system I0421 17:53:07.738280 138547380145280 train.py:806] [DECOUPLED NO-OP] skipping cloud diagnostics wrapper. I0421 17:53:08.427786 138547380145280 maxtext_utils.py:1732] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0421 17:53:08.822625 138547380145280 checkpointing.py:688] Setting up checkpoint logger... I0421 17:53:08.822738 138547380145280 checkpointing.py:234] Creating checkpoint manager with ocdbt=True and zarr3=True I0421 17:53:08.822778 138547380145280 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0421 17:53:08.822984 138547380145280 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 0x7e01c9bab2f0>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0421 17:53:11.235898 138547380145280 checkpointing.py:266] Enabling policy for fixed interval checkpointing. I0421 17:53:11.236114 138547380145280 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 0x7e00717acb30>}, handler_registry=None I0421 17:53:11.236356 138547380145280 composite_checkpoint_handler.py:237] Deferred registration for item: "items". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7e00717acb30>` 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`. I0421 17:53:11.236396 138547380145280 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7e001ed242c0>` 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`. I0421 17:53:11.236425 138547380145280 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 0x7e00717acb30>, ('items', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7e00717acb30>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7e001ed242c0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7e001ed242c0>}). I0421 17:53:11.236722 138547380145280 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.28 I0421 17:53:11.236779 138547380145280 async_checkpointer.py:177] [process=0][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>.<lambda> at 0x7e001f58e840> timeout: 600 secs and primary_host=0 for async checkpoint writes I0421 17:53:11.356279 138547380145280 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints I0421 17:53:11.356520 138547380145280 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/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7e001ed27ef0> I0421 17:53:11.356603 138547380145280 checkpointing.py:302] Checkpoint manager created! I0421 17:53:11.435202 138547380145280 dataset_info.py:707] Load dataset info from tests/assets/local_datasets/c4_en_dataset_minimal/c4/en/3.1.0 I0421 17:53:11.438282 138547380145280 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. I0421 17:53:11.491579 138547380145280 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 I0421 17:53:11.515356 138547380145280 tokenizer.py:245] Tokenizer path: src/maxtext/assets/tokenizers/tokenizer.llama2 I0421 17:53:11.515421 138547380145280 tokenizer.py:187] Loading sentencepiece tokenizer: src/maxtext/assets/tokenizers/tokenizer.llama2 I0421 17:53:12.817840 138547380145280 checkpointing.py:578] checkpoint manager exists so trying to load this run's existing checkpoint I0421 17:53:12.817970 138547380145280 checkpointing.py:676] No existing checkpoints found, not restoring checkpoint. [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 I0421 17:53:16.322550 138547380145280 nnx_decoders.py:465] nnx_decoders/carry Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0421 17:53:16.322778 138547380145280 nnx_decoders.py:465] nnx_decoders/carry Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0421 17:53:16.328033 138547380145280 nnx_decoders.py:465] Unknown Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_norm_length', 'activation_embed'). I0421 17:53:16.328200 138547380145280 nnx_decoders.py:465] Unknown Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0421 17:53:16.341598 138547380145280 attentions.py:1088] attentions/inputs_q Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0421 17:53:16.341679 138547380145280 attentions.py:1088] attentions/inputs_q Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0421 17:53:16.354268 138547380145280 attentions.py:1089] attentions/inputs_kv Logical: bfloat16[8,2048,2048]....................................... ('activation_batch', 'activation_attn_length', 'activation_attn_embed'). I0421 17:53:16.354350 138547380145280 attentions.py:1089] attentions/inputs_kv Physical: bfloat16[8,2048,2048]....................................... ('fsdp', None, None). I0421 17:53:16.406295 138547380145280 attentions.py:1154] attentions/query Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0421 17:53:16.406408 138547380145280 attentions.py:1154] attentions/query Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0421 17:53:16.419217 138547380145280 attentions.py:1155] attentions/key Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0421 17:53:16.419296 138547380145280 attentions.py:1155] attentions/key Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0421 17:53:16.432000 138547380145280 attentions.py:1156] attentions/value Logical: bfloat16[8,2048,16,128]..................................... ('activation_kv_batch', 'activation_attn_length', 'activation_kv_heads', 'activation_kv_head_dim'). I0421 17:53:16.432106 138547380145280 attentions.py:1156] attentions/value Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0421 17:53:16.453998 138547380145280 attentions.py:1198] attentions/out Logical: bfloat16[8,2048,16,128]..................................... ('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv'). I0421 17:53:16.454107 138547380145280 attentions.py:1198] attentions/out Physical: bfloat16[8,2048,16,128]..................................... ('fsdp', None, None, None). I0421 17:53:16.541869 138547380145280 linears.py:525] linears/x Logical: bfloat16[8,2048,7168]....................................... ('activation_batch', 'activation_length', 'activation_mlp'). I0421 17:53:16.541987 138547380145280 linears.py:525] linears/x Physical: bfloat16[8,2048,7168]....................................... ('fsdp', None, None). I0421 17:53:27.621690 138547380145280 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. I0421 17:53:27.625413 138547380145280 max_utils.py:194] tensorboardX not available; using no-op SummaryWriter. I0421 17:53:27.628185 138547380145280 metric_logger.py:301] number parameters: 1.104 billion I0421 17:53:42.287547 138547380145280 checkpointing.py:794] Waiting for step 0 to finish before checkpoint... I0421 17:53:42.424866 138547380145280 checkpointing.py:798] Waited 0.1372966766357422 seconds for step 0 to finish before starting checkpointing. I0421 17:53:42.425446 138547380145280 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0421 17:53:42.425618 138547380145280 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 0 I0421 17:53:42.426037 138547380145280 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0. I0421 17:53:42.513439 138547380145280 signaling_client.py:373] Using ThreadSafeKeyValueSignalingClient I0421 17:53:42.628107 138455090726464 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0 I0421 17:53:42.632356 138547380145280 jax_array_handlers.py:347] Scheduling D2H of 153 prioritized jax.Array. I0421 17:53:42.632462 138547380145280 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 I0421 17:53:43.366388 138455698900544 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/items I0421 17:53:43.497641 138455667443264 checkpoint.py:188] Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776794023267538338, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1776794023267538338, "commit_timestamp_nsecs": null, "custom_metadata": {}} to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0421 17:53:44.342731 1680802 google_auth_provider.cc:149] Using credentials at ~/.config/gcloud/application_default_credentials.json I0421 17:53:44.342806 1680802 google_auth_provider.cc:156] Using OAuth2 AuthProvider I0421 17:53:44.969084 138547380145280 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 2.339313s I0421 17:53:44.969700 138547380145280 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 5.027 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 2 seconds) (per-host) I0421 17:53:44.969766 138547380145280 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 2.455242s (batch_requests_ready=0.096367s, total_serialization_initiated=2.358413s, others=0.000462s) I0421 17:53:44.969865 138547380145280 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 2.455929s (all_items=0.000021s, per_item={'items': '0.00002122'}, temp_paths=2.455908) I0421 17:53:44.984807 138455080240704 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0421 17:53:44.984948 138547380145280 async_checkpointer.py:561] Finished blocking save. Time taken: 2.559274s. Continuing background save to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0. I0421 17:53:44.985232 138547380145280 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=0] Starting CheckpointManager Save Finalize thread=save_finalize I0421 17:53:44.985538 138455552099904 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0421 17:53:44.985661 138547380145280 standard_logger.py:34] {'step': 0, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776794022.4254303, 'wait_for_prev_duration_secs': 4.5299530029296875e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776794022.4256415, 'checkpointer_blocking_duration_secs': 2.559495687484741, 'get_old_steps_start_time': 1776794024.9851568, 'get_old_steps_duration_secs': 2.5987625122070312e-05, 'checkpoint_manager_blocking_start_time': 1776794022.4253411, 'checkpoint_manager_blocking_duration_secs': 2.560288667678833} I0421 17:53:44.985813 138547380145280 checkpointing.py:409] Started an asynchronous checkpoint save for step 0 I0421 17:53:44.985885 138547380145280 max_utils.py:750] Memstats: After params initialized: I0421 17:53:44.985945 138547380145280 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_0(process=0,(0,0,0,0)) I0421 17:53:44.985980 138547380145280 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_1(process=0,(1,0,0,0)) I0421 17:53:44.986003 138547380145280 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_2(process=0,(0,1,0,0)) I0421 17:53:44.986025 138547380145280 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_3(process=0,(1,1,0,0)) I0421 17:53:44.986058 138547380145280 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_4(process=0,(0,2,0,0)) I0421 17:53:44.986079 138547380145280 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_5(process=0,(1,2,0,0)) I0421 17:53:44.986098 138547380145280 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_6(process=0,(0,3,0,0)) I0421 17:53:44.986116 138547380145280 max_utils.py:756] Using (GB) 1.61 / 31.25 (5.152000%) on TPU_7(process=0,(1,3,0,0)) I0421 17:53:45.410087 138547380145280 metric_logger.py:196] completed step: 0, seconds: 14.656, TFLOP/s/device: 0.927, Tokens/s/device: 139.734, total_weights: 13328, loss: 10.885, lm_loss: 10.885, perplexity: 53359.055 I0421 17:53:45.411037 138547380145280 metric_logger.py:281] To see full metrics 'tensorboard --logdir=gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/tensorboard/' I0421 17:53:45.508493 138547380145280 metric_logger.py:196] completed step: 1, seconds: 3.112, TFLOP/s/device: 4.365, Tokens/s/device: 658.013, total_weights: 12332, loss: 10.851, lm_loss: 10.851, perplexity: 51588.516 I0421 17:53:45.619287 138547380145280 metric_logger.py:196] completed step: 2, seconds: 0.044, TFLOP/s/device: 309.002, Tokens/s/device: 46576.152, total_weights: 15161, loss: 10.853, lm_loss: 10.853, perplexity: 51692.324 I0421 17:53:45.730500 138547380145280 metric_logger.py:196] completed step: 3, seconds: 0.079, TFLOP/s/device: 171.791, Tokens/s/device: 25894.223, total_weights: 13327, loss: 10.861, lm_loss: 10.861, perplexity: 52098.078 I0421 17:53:45.748745 138455677929024 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 153 array_metadata.ArrayMetadata to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/items/array_metadatas/process_0 I0421 17:54:03.160729 138547380145280 metric_logger.py:196] completed step: 4, seconds: 0.110, TFLOP/s/device: 124.048, Tokens/s/device: 18697.903, total_weights: 11939, loss: 10.851, lm_loss: 10.851, perplexity: 51575.695 I0421 17:54:03.210970 138547380145280 metric_logger.py:196] completed step: 5, seconds: 0.111, TFLOP/s/device: 122.257, Tokens/s/device: 18427.872, total_weights: 15502, loss: 10.795, lm_loss: 10.795, perplexity: 48797.516 I0421 17:54:03.631375 138547380145280 metric_logger.py:196] completed step: 6, seconds: 17.463, TFLOP/s/device: 0.778, Tokens/s/device: 117.274, total_weights: 13864, loss: 10.785, lm_loss: 10.785, perplexity: 48309.668 I0421 17:54:03.641772 138547380145280 metric_logger.py:196] completed step: 7, seconds: 0.422, TFLOP/s/device: 32.206, Tokens/s/device: 4854.507, total_weights: 12988, loss: 10.765, lm_loss: 10.765, perplexity: 47319.543 I0421 17:54:03.749381 138547380145280 metric_logger.py:196] completed step: 8, seconds: 0.011, TFLOP/s/device: 1190.392, Tokens/s/device: 179428.772, total_weights: 13820, loss: 10.738, lm_loss: 10.738, perplexity: 46060.195 I0421 17:54:03.760277 138547380145280 checkpointing.py:794] Waiting for step 10 to finish before checkpoint... I0421 17:54:03.973368 138547380145280 checkpointing.py:798] Waited 0.213057279586792 seconds for step 10 to finish before starting checkpointing. I0421 17:54:03.973924 138547380145280 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=0][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0421 17:54:27.281961 138455562585664 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 42.306672s (commit=41.836441s, array_metadata_write=0.470232s) I0421 17:54:27.283436 138455080240704 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 282.297 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 44 seconds) (per-host) I0421 17:54:27.283546 138455080240704 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 42.298531s. I0421 17:54:27.579152 138455080240704 checkpoint.py:228] Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776794023267538338, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0421 17:54:27.753490 138455080240704 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0421 17:54:27.959246 138455667443264 checkpoint.py:247] Updated Metadata={'item_handlers': {'items': 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776794023267538338, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/_CHECKPOINT_METADATA I0421 17:54:28.179854 138455080240704 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0421 17:54:28.180894 138455080240704 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.562032s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/items I0421 17:54:28.181705 138455080240704 atomicity.py:608] Finalizing gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0/items I0421 17:54:28.415660 138455080240704 atomicity.py:608] Finalizing gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0 I0421 17:54:29.080947 138455080240704 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0`. I0421 17:54:29.081709 138455080240704 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 46.656056s. directory=gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/0 I0421 17:54:29.081782 138455080240704 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 44.096768s. I0421 17:54:29.081954 138455552099904 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0421 17:54:29.082091 138455552099904 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0421 17:54:29.082143 138455552099904 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is syncing with other hosts... I0421 17:54:29.082188 138455552099904 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=0] CheckpointManager Save Finalize is done on all hosts. I0421 17:54:29.082274 138547380145280 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. W0421 17:54:29.082341 138547380145280 checkpoint_manager.py:1441] Waiting for previous save to complete took 25.108424 seconds. If this number is high, consider checkpointing less frequently. I0421 17:54:29.082820 138547380145280 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 10 I0421 17:54:29.083139 138547380145280 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10. I0421 17:54:29.263632 138455552099904 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10 I0421 17:54:29.288801 138547380145280 jax_array_handlers.py:347] Scheduling D2H of 153 prioritized jax.Array. I0421 17:54:29.288906 138547380145280 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 I0421 17:54:29.978788 138455541614144 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10/items I0421 17:54:36.240950 138547380145280 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 6.954561s I0421 17:54:36.248128 138547380145280 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1.744 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 7 seconds) (per-host) I0421 17:54:36.248202 138547380145280 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 7.075807s (batch_requests_ready=0.097450s, total_serialization_initiated=6.971318s, others=0.007040s) I0421 17:54:36.248315 138547380145280 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 7.076431s (all_items=0.000017s, per_item={'items': '0.00001669'}, temp_paths=7.076414) I0421 17:54:36.254688 138455562585664 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0421 17:54:36.254801 138547380145280 async_checkpointer.py:561] Finished blocking save. Time taken: 7.171937s. Continuing background save to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10. I0421 17:54:36.255005 138547380145280 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=10] Starting CheckpointManager Save Finalize thread=save_finalize I0421 17:54:36.255197 138455080240704 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0421 17:54:36.255314 138547380145280 standard_logger.py:34] {'step': 10, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776794043.9738963, 'wait_for_prev_duration_secs': 25.108424186706543, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776794069.0828443, 'checkpointer_blocking_duration_secs': 7.1720826625823975, 'get_old_steps_start_time': 1776794076.2549427, 'get_old_steps_duration_secs': 2.2649765014648438e-05, 'checkpoint_manager_blocking_start_time': 1776794043.9738276, 'checkpoint_manager_blocking_duration_secs': 32.281458377838135} I0421 17:54:36.255475 138547380145280 checkpointing.py:409] Started an asynchronous checkpoint save for step 10 I0421 17:54:36.256180 138547380145280 metric_logger.py:196] completed step: 9, seconds: 0.009, TFLOP/s/device: 1431.731, Tokens/s/device: 215806.112, total_weights: 12300, loss: 10.730, lm_loss: 10.730, perplexity: 45727.363 I0421 17:54:36.271422 138547380145280 metric_logger.py:196] completed step: 10, seconds: 0.110, TFLOP/s/device: 123.830, Tokens/s/device: 18665.014, total_weights: 13436, loss: 10.691, lm_loss: 10.691, perplexity: 43967.840 I0421 17:54:36.375127 138547380145280 metric_logger.py:196] completed step: 11, seconds: 32.507, TFLOP/s/device: 0.418, Tokens/s/device: 63.003, total_weights: 13817, loss: 10.683, lm_loss: 10.683, perplexity: 43623.184 I0421 17:54:36.486305 138547380145280 metric_logger.py:196] completed step: 12, seconds: 0.013, TFLOP/s/device: 1024.516, Tokens/s/device: 154426.180, total_weights: 12076, loss: 10.692, lm_loss: 10.692, perplexity: 43984.438 I0421 17:54:36.597374 138547380145280 metric_logger.py:196] completed step: 13, seconds: 0.106, TFLOP/s/device: 128.779, Tokens/s/device: 19411.034, total_weights: 11048, loss: 10.703, lm_loss: 10.703, perplexity: 44472.184 I0421 17:54:36.708342 138547380145280 metric_logger.py:196] completed step: 14, seconds: 0.111, TFLOP/s/device: 122.435, Tokens/s/device: 18454.773, total_weights: 14267, loss: 10.682, lm_loss: 10.682, perplexity: 43586.086 I0421 17:54:36.819472 138547380145280 metric_logger.py:196] completed step: 15, seconds: 0.112, TFLOP/s/device: 121.455, Tokens/s/device: 18306.963, total_weights: 13004, loss: 10.661, lm_loss: 10.661, perplexity: 42643.445 I0421 17:54:36.930867 138547380145280 metric_logger.py:196] completed step: 16, seconds: 0.110, TFLOP/s/device: 123.349, Tokens/s/device: 18592.490, total_weights: 13811, loss: 10.659, lm_loss: 10.659, perplexity: 42569.543 I0421 17:54:37.041797 138547380145280 metric_logger.py:196] completed step: 17, seconds: 0.111, TFLOP/s/device: 122.241, Tokens/s/device: 18425.551, total_weights: 13709, loss: 10.658, lm_loss: 10.658, perplexity: 42516.812 I0421 17:54:37.090610 138455677929024 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 153 array_metadata.ArrayMetadata to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10/items/array_metadatas/process_0 I0421 17:54:37.155637 138547380145280 metric_logger.py:196] completed step: 18, seconds: 0.111, TFLOP/s/device: 121.862, Tokens/s/device: 18368.372, total_weights: 12831, loss: 10.659, lm_loss: 10.659, perplexity: 42556.516 I0421 17:54:37.265748 138547380145280 checkpointing.py:794] Waiting for step 19 to finish before checkpoint... I0421 17:54:37.268561 138547380145280 checkpointing.py:798] Waited 0.002813100814819336 seconds for step 19 to finish before starting checkpointing. I0421 17:54:37.268921 138547380145280 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=10][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0421 17:55:05.004736 138455304635968 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 28.751005s (commit=28.284742s, array_metadata_write=0.466263s) I0421 17:55:05.006103 138455562585664 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 352.689 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 35 seconds) (per-host) I0421 17:55:05.006226 138455562585664 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 28.751378s. I0421 17:55:05.425837 138455562585664 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0421 17:55:05.820654 138455562585664 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0421 17:55:05.821794 138455562585664 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.551805s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10/items I0421 17:55:05.822516 138455562585664 atomicity.py:608] Finalizing gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10/items I0421 17:55:06.066101 138455562585664 atomicity.py:608] Finalizing gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10 I0421 17:55:06.765116 138455562585664 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10`. I0421 17:55:06.765846 138455562585664 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 37.682991s. directory=gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/10 I0421 17:55:06.765930 138455562585664 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 30.511083s. I0421 17:55:06.766090 138455080240704 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0421 17:55:06.766206 138455080240704 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0421 17:55:06.766255 138455080240704 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=10] CheckpointManager Save Finalize is syncing with other hosts... I0421 17:55:06.766295 138455080240704 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=10] CheckpointManager Save Finalize is done on all hosts. I0421 17:55:06.766383 138547380145280 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. W0421 17:55:06.766437 138547380145280 checkpoint_manager.py:1441] Waiting for previous save to complete took 29.497530 seconds. If this number is high, consider checkpointing less frequently. I0421 17:55:06.766622 138547380145280 checkpoint_manager.py:1501] [process=0] Saving checkpoint at step 19 I0421 17:55:06.766903 138547380145280 async_checkpointer.py:452] [process=0] Started async saving checkpoint to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19. I0421 17:55:06.941726 138455562585664 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19 I0421 17:55:06.969107 138547380145280 jax_array_handlers.py:347] Scheduling D2H of 153 prioritized jax.Array. I0421 17:55:06.969212 138547380145280 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 I0421 17:55:07.661836 138455646471744 atomicity.py:137] Creating tmp directory gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19/items I0421 17:55:15.381855 138547380145280 base_pytree_checkpoint_handler.py:153] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 8.415109s I0421 17:55:15.393781 138547380145280 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1.445 GiB/s (total gbytes: 12.3 GiB) (time elapsed: 8 seconds) (per-host) I0421 17:55:15.393852 138547380145280 base_pytree_checkpoint_handler.py:732] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 8.541896s (batch_requests_ready=0.098491s, total_serialization_initiated=8.431630s, others=0.011775s) I0421 17:55:15.393940 138547380145280 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 8.542469s (all_items=0.000012s, per_item={'items': '0.00001216'}, temp_paths=8.542457) I0421 17:55:15.395557 138464051856960 async_checkpointer.py:79] [process=0][thread=async_save] Background save thread started. I0421 17:55:15.395668 138547380145280 async_checkpointer.py:561] Finished blocking save. Time taken: 8.629005s. Continuing background save to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19. I0421 17:55:15.395844 138547380145280 checkpoint_manager.py:1549] [process=0][thread=MainThread][step=19] Starting CheckpointManager Save Finalize thread=save_finalize I0421 17:55:15.396012 138455080240704 async_checkpointer.py:265] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0421 17:55:15.396128 138547380145280 standard_logger.py:34] {'step': 19, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776794077.2688954, 'wait_for_prev_duration_secs': 29.497530460357666, 'time_between_consecutive_saves_sec': 8.186691761016846, 'checkpointer_blocking_start_time': 1776794106.766644, 'checkpointer_blocking_duration_secs': 8.629139423370361, 'get_old_steps_start_time': 1776794115.395794, 'get_old_steps_duration_secs': 1.8358230590820312e-05, 'checkpoint_manager_blocking_start_time': 1776794077.2688618, 'checkpoint_manager_blocking_duration_secs': 38.12724208831787} I0421 17:55:15.396336 138547380145280 checkpointing.py:409] Started an asynchronous checkpoint save for step 19 I0421 17:55:15.396370 138547380145280 checkpoint_manager.py:1994] [process=0][thread=MainThread][step=19][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0421 17:55:16.154148 138455552099904 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 153 array_metadata.ArrayMetadata to gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19/items/array_metadatas/process_0 I0421 17:55:44.398872 138455562585664 base_pytree_checkpoint_handler.py:1217] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 29.004348s (commit=28.485673s, array_metadata_write=0.518675s) I0421 17:55:44.400216 138464051856960 base_pytree_checkpoint_handler.py:128] [process=0] /jax/checkpoint/write/gbytes_per_sec: 336.584 MiB/s (total gbytes: 12.3 GiB) (time elapsed: 37 seconds) (per-host) I0421 17:55:44.400325 138464051856960 async_checkpointer.py:90] [process=0][thread=async_save] 3 Handler Commit operations completed. Time taken: 29.004611s. I0421 17:55:44.840378 138464051856960 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0421 17:55:45.306816 138464051856960 ocdbt_utils.py:56] Param validation support for Zarr3 will be added later (b/362328389). I0421 17:55:45.307823 138464051856960 base_pytree_checkpoint_handler.py:1342] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.611942s. use_zarr3=True, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19/items I0421 17:55:45.308552 138464051856960 atomicity.py:608] Finalizing gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19/items I0421 17:55:45.549129 138464051856960 atomicity.py:608] Finalizing gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19 I0421 17:55:46.239145 138464051856960 atomicity.py:794] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19`. I0421 17:55:46.239920 138464051856960 async_checkpointer.py:420] Finished async_save (blocking + background). Time taken: 39.473265s. directory=gs://wanglance-maxtext/nnx_ckpt_feat_nnx_trainstate_and_training_loop_20260421_161934/nnx_feat_nnx_trainstate_and_training_loop_20260421_161934_05_fp8/checkpoints/19 I0421 17:55:46.239985 138464051856960 async_checkpointer.py:144] [process=0][thread=async_save] Background save thread done. Time taken: 30.844273s. I0421 17:55:46.240151 138455080240704 async_checkpointer.py:273] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0421 17:55:46.240262 138455080240704 async_checkpointer.py:283] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0421 17:55:46.240316 138455080240704 checkpoint_manager.py:2103] [process=0][thread=save_finalize][step=19] CheckpointManager Save Finalize is syncing with other hosts... I0421 17:55:46.240356 138455080240704 checkpoint_manager.py:2112] [process=0][thread=save_finalize][step=19] CheckpointManager Save Finalize is done on all hosts. I0421 17:55:46.240503 138547380145280 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. I0421 17:55:46.240686 138547380145280 checkpoint_manager.py:1983] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0421 17:55:46.241615 138547380145280 metric_logger.py:196] completed step: 19, seconds: 0.111, TFLOP/s/device: 122.873, Tokens/s/device: 18520.696, total_weights: 13717, loss: 10.639, lm_loss: 10.639, perplexity: 41740.996 Per train step: Total TFLOPs: 13.59 split as 93.93% learnable weight flops and 6.07% attention flops