2026-04-16 21:11:20.865701: 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) I0416 21:11:21.338263 125630730390656 max_utils.py:238] Skipping jax distributed system due to skip_jax_distributed_system=True flag. I0416 21:12:18.422026 125630730390656 max_utils.py:800] System Information: Jax Version: 0.8.3 I0416 21:12:18.422156 125630730390656 max_utils.py:801] System Information: Jaxlib Version: 0.8.3 I0416 21:12:18.422191 125630730390656 max_utils.py:802] System Information: Jax Backend: PJRT C API TFRT TPU v6 lite Built on Dec 15 2025 14:03:46 (1765836226) cl/844590465 I0416 21:12:18.424954 125630730390656 maxtext_utils.py:1687] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0416 21:12:18.506169 125630730390656 maxtext_utils.py:1687] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0416 21:12:18.589121 125630730390656 maxtext_utils.py:1687] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0416 21:12:19.550992 125630730390656 max_utils.py:194] tensorboardX not available; using no-op SummaryWriter. I0416 21:12:19.572965 125630730390656 config.py:112] TensorFlow version 2.20.0 available. I0416 21:12:19.573403 125630730390656 config.py:125] JAX version 0.8.3 available. E0416 21:12:21.213800 125630730390656 packing.py:209] PackAndBatchOperation is deprecated. Please use lazy_dataset.FirstFitPackIterDataset instead. I0416 21:12:21.557279 125630730390656 pytree_checkpoint_handler.py:592] save_device_host_concurrent_bytes=None I0416 21:12:21.557686 125630730390656 base_pytree_checkpoint_handler.py:441] Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x72420cdf0170>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0416 21:12:21.557728 125630730390656 pytree_checkpoint_handler.py:592] save_device_host_concurrent_bytes=None I0416 21:12:21.557760 125630730390656 base_pytree_checkpoint_handler.py:441] Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x72420cdf0170>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0416 21:12:21.557800 125630730390656 checkpoint_manager.py:708] [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers={'model_params': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x723b70d01820>, 'optimizer_state': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x723b75992300>, 'custom_metadata': <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x723b72502540>}, handler_registry=None I0416 21:12:21.558089 125630730390656 composite_checkpoint_handler.py:237] Deferred registration for item: "model_params". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x723b70d01820>` for item "model_params" 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`. I0416 21:12:21.558130 125630730390656 composite_checkpoint_handler.py:237] Deferred registration for item: "optimizer_state". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x723b75992300>` for item "optimizer_state" 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`. I0416 21:12:21.558153 125630730390656 composite_checkpoint_handler.py:237] Deferred registration for item: "custom_metadata". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x723b72502540>` for item "custom_metadata" 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`. I0416 21:12:21.558172 125630730390656 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x723b70d78b90>` 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`. I0416 21:12:21.558194 125630730390656 composite_checkpoint_handler.py:505] Initialized registry DefaultCheckpointHandlerRegistry({('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x723b70d01820>, ('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x723b70d01820>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x723b75992300>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x723b75992300>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x723b72502540>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x723b72502540>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x723b70d78b90>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x723b70d78b90>}). I0416 21:12:21.558350 125630730390656 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.36 I0416 21:12:21.558386 125630730390656 async_checkpointer.py:192] [process=0][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>.<lambda> at 0x723b70d9a980> timeout: 1200 secs and primary_host=0 for async checkpoint writes I0416 21:12:23.530811 125630730390656 checkpoint_manager.py:564] Created directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints I0416 21:12:25.884749 125630730390656 checkpoint_manager.py:1812] Found 0 checkpoint steps in gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints I0416 21:12:25.885014 125630730390656 checkpoint_manager.py:929] [process=0][thread=MainThread] CheckpointManager created, primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=10000, max_to_keep=None, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=False, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_background_delete=False, read_only=False, enable_async_checkpointing=True, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False, lightweight_initialize=False), root_directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x723b70d34500> I0416 21:12:27.919018 125630730390656 metrics_logger.py:64] WandbBackend skipped: 'wandb' library not installed. I0416 21:12:27.919302 125630730390656 peft_trainer.py:590] Training with mesh: Mesh('diloco': 1, 'data': 1, 'stage': 1, 'fsdp': 8, 'fsdp_transpose': 1, 'sequence': 1, 'context': 1, 'context_autoregressive': 1, 'tensor': 1, 'tensor_transpose': 1, 'tensor_sequence': 1, 'expert': 1, 'autoregressive': 1, axis_types=(Auto, Auto, Auto, Auto, Auto, Auto, Auto, Auto, Auto, Auto, Auto, Auto, Auto)) I0416 21:12:28.343440 125630730390656 peft_trainer.py:600] Compiled train_step cache size: 0 [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. Training: 0%| | 0/5 [00:00<?, ?step/s]I0416 21:12:28.346850 125630730390656 metric_logger.py:289] number parameters: 0.000 billion Per train step: Total TFLOPs: 0.00 split as 54.29% learnable weight flops and 45.71% attention flops 2026-04-16 21:12:31.523597: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. 2026-04-16 21:12:31.559899: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. 2026-04-16 21:12:32.586056: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. 2026-04-16 21:12:35.038875: 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) I0416 21:12:43.622141 125630730390656 checkpoint_manager.py:2009] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0416 21:12:43.622338 125630730390656 checkpoint_manager.py:1512] [process=0] Saving checkpoint at step 1 I0416 21:12:43.622401 125630730390656 event_tracking.py:70] [process=0] [async] Started save checkpoint @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1. I0416 21:12:43.721835 125630730390656 signaling_client.py:373] Using ThreadSafeKeyValueSignalingClient I0416 21:12:43.745375 125630730390656 jax_array_handlers.py:360] Scheduling D2H of 46 prioritized jax.Array. I0416 21:12:43.745437 125630730390656 replica_slices.py:424] Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=False I0416 21:12:43.820411 125520742188608 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1 I0416 21:12:44.522585 125520731702848 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/model_params I0416 21:12:44.530431 125520731702848 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/optimizer_state I0416 21:12:44.601477 125630730390656 base_pytree_checkpoint_handler.py:154] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.857357s I0416 21:12:44.601936 125630730390656 jax_array_handlers.py:360] Scheduling D2H of 52 prioritized jax.Array. I0416 21:12:44.601984 125630730390656 replica_slices.py:424] Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=False I0416 21:12:44.632354 125630730390656 base_pytree_checkpoint_handler.py:154] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.030776s I0416 21:12:44.632738 125630730390656 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/blocking_gbytes_per_sec: 225.446 KiB/s (total gbytes: 205.1 KiB) (time elapsed: 0.9098124504089355 s) (per-host) I0416 21:12:44.632850 125630730390656 base_pytree_checkpoint_handler.py:768] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.909937s (batch_requests_ready=0.002929s, total_serialization_initiated=0.906625s, others=0.000383s) I0416 21:12:44.633179 125630730390656 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/blocking_gbytes_per_sec: 685.026 KiB/s (total gbytes: 614.8 KiB) (time elapsed: 0.8975248336791992 s) (per-host) I0416 21:12:44.633291 125630730390656 base_pytree_checkpoint_handler.py:768] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.897647s (batch_requests_ready=0.002311s, total_serialization_initiated=0.894947s, others=0.000388s) I0416 21:12:44.633352 125630730390656 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.910980s (all_items=0.000020s, per_item={'model_params': '0.00001621', 'optimizer_state': '0.00000381'}, temp_paths=0.910960) I0416 21:12:44.633954 125630730390656 event_tracking.py:125] [process=0] [async] Finished blocking save in 1.01 seconds. Continuing save @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1. I0416 21:12:44.634176 125520622650944 async_checkpointer.py:76] [process=0][thread=async_save] Background save thread started. Deadline for this save operation is 2026-04-16 21:32:44.634141 I0416 21:12:44.634407 125630730390656 checkpoint_manager.py:1560] [process=0][thread=MainThread][step=1] Starting CheckpointManager Save Finalize thread=save_finalize I0416 21:12:44.634724 125630730390656 standard_logger.py:34] {'step': 1, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776373963.622122, 'wait_for_prev_duration_secs': 9.250640869140625e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776373963.6223605, 'checkpointer_blocking_duration_secs': 1.0119056701660156, 'get_old_steps_start_time': 1776373964.6342828, 'get_old_steps_duration_secs': 8.678436279296875e-05, 'checkpoint_manager_blocking_start_time': 1776373963.6220398, 'checkpoint_manager_blocking_duration_secs': 1.0126631259918213} I0416 21:12:44.634864 125630730390656 profiler.py:85] Starting JAX profiler at step 1. I0416 21:12:44.889782 125520696051264 checkpoint.py:188] Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776373964428003108, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1776373964428003108, "commit_timestamp_nsecs": null, "custom_metadata": {}} to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/_CHECKPOINT_METADATA I0416 21:12:44.890212 125520752674368 async_checkpointer.py:280] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0416 21:12:44.973328 125630730390656 peft_trainer.py:485] Train step 1 training loss: 6.004404 - training perplexity: 405.209259 Training: 0%| | 0/5 [00:16<?, ?step/s, _train_loss=6, _train_perplexity=405, _train_steps_per_sec=0.065] Training: 20%|██ | 1/5 [00:16<01:06, 16.63s/step, _train_loss=6, _train_perplexity=405, _train_steps_per_sec=0.065]I0416 21:12:44.974439 125630730390656 max_utils.py:750] Memstats: After params initialized: I0416 21:12:44.974524 125630730390656 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_0(process=0,(0,0,0,0)) I0416 21:12:44.974582 125630730390656 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_1(process=0,(1,0,0,0)) I0416 21:12:44.974632 125630730390656 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_2(process=0,(0,1,0,0)) I0416 21:12:44.974676 125630730390656 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_3(process=0,(1,1,0,0)) I0416 21:12:44.974719 125630730390656 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_4(process=0,(0,2,0,0)) I0416 21:12:44.974761 125630730390656 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_5(process=0,(1,2,0,0)) I0416 21:12:44.974807 125630730390656 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_6(process=0,(0,3,0,0)) I0416 21:12:44.974851 125630730390656 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_7(process=0,(1,3,0,0)) I0416 21:12:45.100788 125630730390656 metric_logger.py:185] completed step: 1, seconds: 16.628, TFLOP/s/device: 0.000, Tokens/s/device: 61.585, total_weights: 6826, loss: 6.004 I0416 21:12:45.113594 125630730390656 peft_trainer.py:485] Train step 2 training loss: 6.137612 - training perplexity: 462.946655 Training: 20%|██ | 1/5 [00:16<01:06, 16.63s/step, _train_loss=6.07, _train_perplexity=433, _train_steps_per_sec=0.403] Training: 40%|████ | 2/5 [00:16<00:20, 6.93s/step, _train_loss=6.07, _train_perplexity=433, _train_steps_per_sec=0.403]I0416 21:12:45.115060 125630730390656 metric_logger.py:185] completed step: 2, seconds: 0.140, TFLOP/s/device: 0.002, Tokens/s/device: 7336.279, total_weights: 4636, loss: 6.138 I0416 21:12:45.135096 125630730390656 peft_trainer.py:485] Train step 3 training loss: 5.638189 - training perplexity: 280.953430 Training: 40%|████ | 2/5 [00:16<00:20, 6.93s/step, _train_loss=5.93, _train_perplexity=375, _train_steps_per_sec=2.59] I0416 21:12:45.136373 125630730390656 metric_logger.py:185] completed step: 3, seconds: 0.021, TFLOP/s/device: 0.010, Tokens/s/device: 48230.951, total_weights: 5886, loss: 5.638 I0416 21:12:45.147703 125630730390656 peft_trainer.py:485] Train step 4 training loss: 5.768358 - training perplexity: 320.011780 Training: 60%|██████ | 3/5 [00:16<00:13, 6.93s/step, _train_loss=5.89, _train_perplexity=360, _train_steps_per_sec=13.6]I0416 21:12:45.148957 125630730390656 metric_logger.py:185] completed step: 4, seconds: 0.013, TFLOP/s/device: 0.017, Tokens/s/device: 80794.500, total_weights: 4990, loss: 5.768 I0416 21:12:45.149272 125630730390656 profiler.py:113] Stopping JAX profiler at step 5. I0416 21:12:45.220222 4168901 google_auth_provider.cc:149] Using credentials at ~/.config/gcloud/application_default_credentials.json I0416 21:12:45.220320 4168901 google_auth_provider.cc:156] Using OAuth2 AuthProvider I0416 21:12:45.998951 125520654108224 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 52 array_metadata.ArrayMetadata to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/optimizer_state/array_metadatas/process_0 I0416 21:12:45.999851 125520675079744 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 46 array_metadata.ArrayMetadata to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/model_params/array_metadatas/process_0 I0416 21:12:47.702682 125520633136704 base_pytree_checkpoint_handler.py:1282] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 2.480485s (commit=2.027834s, array_metadata_write=0.452651s) I0416 21:12:47.715854 125520643622464 base_pytree_checkpoint_handler.py:1282] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 2.487516s (commit=2.049904s, array_metadata_write=0.437612s) I0416 21:12:47.716807 125520622650944 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/gbytes_per_sec: 51.357 KiB/s (total gbytes: 205.1 KiB) (time elapsed: 3.993842601776123 s) (per-host) I0416 21:12:47.717123 125520622650944 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/gbytes_per_sec: 154.422 KiB/s (total gbytes: 614.8 KiB) (time elapsed: 3.9814741611480713 s) (per-host) I0416 21:12:47.717203 125520622650944 async_checkpointer.py:90] [process=0][thread=async_save] 4 Handler Commit operations completed. Time taken: 3.082754s. I0416 21:12:47.899681 125520622650944 checkpoint.py:228] Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776373964428003108, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/_CHECKPOINT_METADATA I0416 21:12:48.063940 125520622650944 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0416 21:12:48.277314 125520696051264 checkpoint.py:247] Updated Metadata={'item_handlers': {'model_params': 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler', 'optimizer_state': 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776373964428003108, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/_CHECKPOINT_METADATA I0416 21:12:48.492764 125520622650944 base_pytree_checkpoint_handler.py:1406] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.556041s. use_zarr3=False, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/model_params I0416 21:12:48.493509 125520622650944 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/model_params I0416 21:12:48.904649 125520622650944 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0416 21:12:49.361347 125520622650944 base_pytree_checkpoint_handler.py:1406] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.584717s. use_zarr3=False, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/optimizer_state I0416 21:12:49.362140 125520622650944 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/optimizer_state I0416 21:12:49.630800 125520622650944 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1 I0416 21:12:50.011156 125630730390656 utils.py:86] Train loop finished in: 21.6636 seconds I0416 21:12:50.011788 125630730390656 peft_trainer.py:485] Train step 5 training loss: 5.830140 - training perplexity: 340.406403 Training: 80%|████████ | 4/5 [00:21<00:06, 6.93s/step, _train_loss=5.88, _train_perplexity=356, _train_steps_per_sec=26.6] Training: 100%|██████████| 5/5 [00:21<00:00, 3.14s/step, _train_loss=5.88, _train_perplexity=356, _train_steps_per_sec=26.6]I0416 21:12:50.013059 125630730390656 metric_logger.py:185] completed step: 5, seconds: 4.864, TFLOP/s/device: 0.000, Tokens/s/device: 210.534, total_weights: 4264, loss: 5.830 I0416 21:12:50.016418 125630730390656 checkpoint_manager.py:2020] [process=0][thread=MainThread][step=1][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0416 21:12:50.396630 125520622650944 atomicity.py:847] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1`. I0416 21:12:50.397328 125520622650944 event_tracking.py:138] [process=0] [async] Finished save (blocking + background) in 6.77 seconds @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1 I0416 21:12:50.397404 125520622650944 async_checkpointer.py:160] [process=0][thread=async_save] Background save thread done. Time taken: 5.762956s. I0416 21:12:50.397569 125520752674368 async_checkpointer.py:288] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0416 21:12:50.397688 125520752674368 async_checkpointer.py:298] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0416 21:12:50.397747 125520752674368 checkpoint_manager.py:2137] [process=0][thread=save_finalize][step=1] CheckpointManager Save Finalize is syncing with other hosts... I0416 21:12:50.397792 125520752674368 checkpoint_manager.py:2146] [process=0][thread=save_finalize][step=1] CheckpointManager Save Finalize is done on all hosts. I0416 21:12:50.397879 125630730390656 checkpoint_manager.py:2032] [process=0][thread=MainThread][step=1][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=1. I0416 21:12:50.398112 125630730390656 checkpoint_manager.py:1512] [process=0] Saving checkpoint at step 5 I0416 21:12:50.398177 125630730390656 event_tracking.py:70] [process=0] [async] Started save checkpoint @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5. I0416 21:12:50.495798 125630730390656 jax_array_handlers.py:360] Scheduling D2H of 46 prioritized jax.Array. I0416 21:12:50.495931 125630730390656 replica_slices.py:424] Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=False I0416 21:12:50.509776 125630730390656 base_pytree_checkpoint_handler.py:154] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.015279s I0416 21:12:50.510094 125630730390656 jax_array_handlers.py:360] Scheduling D2H of 52 prioritized jax.Array. I0416 21:12:50.510136 125630730390656 replica_slices.py:424] Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=False I0416 21:12:50.540337 125630730390656 base_pytree_checkpoint_handler.py:154] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.030495s I0416 21:12:50.540679 125630730390656 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/blocking_gbytes_per_sec: 3.024 MiB/s (total gbytes: 205.1 KiB) (time elapsed: 0.06624436378479004 s) (per-host) I0416 21:12:50.540791 125630730390656 base_pytree_checkpoint_handler.py:768] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.066370s (batch_requests_ready=0.002296s, total_serialization_initiated=0.063724s, others=0.000350s) I0416 21:12:50.541223 125630730390656 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/blocking_gbytes_per_sec: 10.814 MiB/s (total gbytes: 614.8 KiB) (time elapsed: 0.05552482604980469 s) (per-host) I0416 21:12:50.541297 125630730390656 base_pytree_checkpoint_handler.py:768] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.055608s (batch_requests_ready=0.002334s, total_serialization_initiated=0.052819s, others=0.000455s) I0416 21:12:50.541356 125630730390656 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.067413s (all_items=0.000013s, per_item={'model_params': '0.00001025', 'optimizer_state': '0.00000262'}, temp_paths=0.067400) I0416 21:12:50.541937 125630730390656 event_tracking.py:125] [process=0] [async] Finished blocking save in 0.14 seconds. Continuing save @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5. I0416 21:12:50.542116 125520486336064 async_checkpointer.py:76] [process=0][thread=async_save] Background save thread started. Deadline for this save operation is 2026-04-16 21:32:50.542083 I0416 21:12:50.542349 125630730390656 checkpoint_manager.py:1560] [process=0][thread=MainThread][step=5] Starting CheckpointManager Save Finalize thread=save_finalize I0416 21:12:50.542663 125520475850304 async_checkpointer.py:280] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0416 21:12:50.542773 125630730390656 standard_logger.py:34] {'step': 5, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': False, 'wait_for_prev_start_time': 1776373970.0163968, 'wait_for_prev_duration_secs': 0.3815345764160156, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776373970.3981366, 'checkpointer_blocking_duration_secs': 0.14407014846801758, 'get_old_steps_start_time': 1776373970.5422237, 'get_old_steps_duration_secs': 8.7738037109375e-05, 'checkpoint_manager_blocking_start_time': 1776373970.0163558, 'checkpoint_manager_blocking_duration_secs': 0.5263934135437012} I0416 21:12:50.542941 125630730390656 checkpoint_manager.py:2020] [process=0][thread=MainThread][step=5][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0416 21:12:50.557421 125520622650944 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5 I0416 21:12:51.208363 125520643622464 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/optimizer_state I0416 21:12:51.214314 125520643622464 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/model_params I0416 21:12:52.446114 125520675079744 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 46 array_metadata.ArrayMetadata to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/model_params/array_metadatas/process_0 I0416 21:12:52.460892 125520633136704 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 52 array_metadata.ArrayMetadata to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/optimizer_state/array_metadatas/process_0 I0416 21:12:54.014780 125520591193664 base_pytree_checkpoint_handler.py:1282] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 2.347946s (commit=1.925215s, array_metadata_write=0.422731s) I0416 21:12:54.015827 125520559736384 base_pytree_checkpoint_handler.py:1282] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 2.349145s (commit=1.928037s, array_metadata_write=0.421108s) I0416 21:12:54.016997 125520486336064 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/gbytes_per_sec: 57.900 KiB/s (total gbytes: 205.1 KiB) (time elapsed: 3.542527198791504 s) (per-host) I0416 21:12:54.017487 125520486336064 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/gbytes_per_sec: 174.084 KiB/s (total gbytes: 614.8 KiB) (time elapsed: 3.5317907333374023 s) (per-host) I0416 21:12:54.017590 125520486336064 async_checkpointer.py:90] [process=0][thread=async_save] 4 Handler Commit operations completed. Time taken: 3.475188s. I0416 21:12:54.371819 125520486336064 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0416 21:12:54.770210 125520486336064 base_pytree_checkpoint_handler.py:1406] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.534747s. use_zarr3=False, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/model_params I0416 21:12:54.770857 125520486336064 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/model_params I0416 21:12:55.187703 125520486336064 array_metadata_store.py:367] [process=0][thread=async_save] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0. I0416 21:12:55.610383 125520486336064 base_pytree_checkpoint_handler.py:1406] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.558571s. use_zarr3=False, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/optimizer_state I0416 21:12:55.611114 125520486336064 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/optimizer_state I0416 21:12:55.873681 125520486336064 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5 I0416 21:12:56.615331 125520486336064 atomicity.py:847] [process=0][thread=async_save] Finished saving checkpoint (finalized tmp dir) to `gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5`. I0416 21:12:56.615975 125520486336064 event_tracking.py:138] [process=0] [async] Finished save (blocking + background) in 6.22 seconds @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_nnx_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5 I0416 21:12:56.616064 125520486336064 async_checkpointer.py:160] [process=0][thread=async_save] Background save thread done. Time taken: 6.073657s. I0416 21:12:56.616255 125520475850304 async_checkpointer.py:288] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0416 21:12:56.616367 125520475850304 async_checkpointer.py:298] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0416 21:12:56.616421 125520475850304 checkpoint_manager.py:2137] [process=0][thread=save_finalize][step=5] CheckpointManager Save Finalize is syncing with other hosts... I0416 21:12:56.616463 125520475850304 checkpoint_manager.py:2146] [process=0][thread=save_finalize][step=5] CheckpointManager Save Finalize is done on all hosts. I0416 21:12:56.616599 125630730390656 checkpoint_manager.py:2032] [process=0][thread=MainThread][step=5][wait_until_finished] Done waiting for Save Finalize thread (save_finalize) running at step=5. I0416 21:12:56.616852 125630730390656 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=125630730390656 count=1 at 0x723b75984900>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x723b6f163950>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x723b6f163980>, _write_futures=[]) I0416 21:12:56.617353 125630730390656 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=125630730390656 count=1 at 0x723b75984900>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x723b6f163950>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x723b6f163980>, _write_futures=[]) I0416 21:12:56.617381 125630730390656 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=125630730390656 count=1 at 0x723b75984900>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x723b6f163950>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x723b6f163980>, _write_futures=[]) Training: 100%|██████████| 5/5 [00:29<00:00, 5.89s/step, _train_loss=5.88, _train_perplexity=356, _train_steps_per_sec=26.6] [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. ~/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/multiprocessing/resource_tracker.py:279: UserWarning: resource_tracker: There appear to be 15 leaked shared_memory objects to clean up at shutdown warnings.warn('resource_tracker: There appear to be %d '