2026-04-16 21:05:58.072298: 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:05:58.607383 131141424139392 max_utils.py:238] Skipping jax distributed system due to skip_jax_distributed_system=True flag. I0416 21:06:27.700041 131141424139392 max_utils.py:800] System Information: Jax Version: 0.8.3 I0416 21:06:27.700183 131141424139392 max_utils.py:801] System Information: Jaxlib Version: 0.8.3 I0416 21:06:27.700224 131141424139392 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:06:27.704583 131141424139392 maxtext_utils.py:1687] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0416 21:06:27.896583 131141424139392 maxtext_utils.py:1687] Num_devices: 8, shape (1, 1, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0416 21:06:28.928112 131141424139392 max_utils.py:194] tensorboardX not available; using no-op SummaryWriter. I0416 21:06:28.953595 131141424139392 config.py:112] TensorFlow version 2.20.0 available. I0416 21:06:28.953993 131141424139392 config.py:125] JAX version 0.8.3 available. E0416 21:06:31.870134 131141424139392 packing.py:209] PackAndBatchOperation is deprecated. Please use lazy_dataset.FirstFitPackIterDataset instead. I0416 21:06:32.236766 131141424139392 pytree_checkpoint_handler.py:592] save_device_host_concurrent_bytes=None I0416 21:06:32.237290 131141424139392 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 0x77451c39c920>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0416 21:06:32.237338 131141424139392 pytree_checkpoint_handler.py:592] save_device_host_concurrent_bytes=None I0416 21:06:32.237371 131141424139392 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 0x77451c39c920>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0416 21:06:32.237428 131141424139392 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 0x773e7f3d99d0>, 'optimizer_state': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x773e81b025d0>, 'custom_metadata': <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x773e7d7ff8c0>}, handler_registry=None I0416 21:06:32.237707 131141424139392 composite_checkpoint_handler.py:237] Deferred registration for item: "model_params". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x773e7f3d99d0>` 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:06:32.237746 131141424139392 composite_checkpoint_handler.py:237] Deferred registration for item: "optimizer_state". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x773e81b025d0>` 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:06:32.237767 131141424139392 composite_checkpoint_handler.py:237] Deferred registration for item: "custom_metadata". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x773e7d7ff8c0>` 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:06:32.237785 131141424139392 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x773e7d7fc230>` 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:06:32.237809 131141424139392 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 0x773e7f3d99d0>, ('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x773e7f3d99d0>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x773e81b025d0>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x773e81b025d0>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x773e7d7ff8c0>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x773e7d7ff8c0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x773e7d7fc230>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x773e7d7fc230>}). I0416 21:06:32.238992 131141424139392 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.36 I0416 21:06:32.239040 131141424139392 async_checkpointer.py:192] [process=0][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>.<lambda> at 0x773e7d743d80> timeout: 1200 secs and primary_host=0 for async checkpoint writes I0416 21:06:34.580789 131141424139392 checkpoint_manager.py:564] Created directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints I0416 21:06:36.940611 131141424139392 checkpoint_manager.py:1812] Found 0 checkpoint steps in gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints I0416 21:06:36.940864 131141424139392 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_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x773e7d7fc080> I0416 21:06:38.989873 131141424139392 metrics_logger.py:64] WandbBackend skipped: 'wandb' library not installed. I0416 21:06:38.990192 131141424139392 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:06:39.431640 131141424139392 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:06:39.433747 131141424139392 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:06:42.610804: 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:06:42.649431: 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:06:43.658852: 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:06:46.061876: 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:06:56.713685 131141424139392 checkpoint_manager.py:2009] [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0416 21:06:56.713886 131141424139392 checkpoint_manager.py:1512] [process=0] Saving checkpoint at step 1 I0416 21:06:56.713952 131141424139392 event_tracking.py:70] [process=0] [async] Started save checkpoint @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1. I0416 21:06:56.816035 131141424139392 signaling_client.py:373] Using ThreadSafeKeyValueSignalingClient I0416 21:06:56.834959 131141424139392 jax_array_handlers.py:360] Scheduling D2H of 22 prioritized jax.Array. I0416 21:06:56.835082 131141424139392 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:06:56.906855 131031418013248 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1 I0416 21:06:57.625192 131031407527488 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/model_params I0416 21:06:57.631219 131031407527488 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/optimizer_state I0416 21:06:57.694640 131141424139392 base_pytree_checkpoint_handler.py:154] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.860302s I0416 21:06:57.695139 131141424139392 jax_array_handlers.py:360] Scheduling D2H of 52 prioritized jax.Array. I0416 21:06:57.695184 131141424139392 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:06:57.727409 131141424139392 base_pytree_checkpoint_handler.py:154] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.032664s I0416 21:06:57.727786 131141424139392 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/blocking_gbytes_per_sec: 225.179 KiB/s (total gbytes: 205.0 KiB) (time elapsed: 0.9102644920349121 s) (per-host) I0416 21:06:57.727900 131141424139392 base_pytree_checkpoint_handler.py:768] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.910393s (batch_requests_ready=0.002069s, total_serialization_initiated=0.907947s, others=0.000377s) I0416 21:06:57.728331 131141424139392 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/blocking_gbytes_per_sec: 680.282 KiB/s (total gbytes: 614.8 KiB) (time elapsed: 0.9037835597991943 s) (per-host) I0416 21:06:57.728412 131141424139392 base_pytree_checkpoint_handler.py:768] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.903876s (batch_requests_ready=0.002553s, total_serialization_initiated=0.900864s, others=0.000459s) I0416 21:06:57.728477 131141424139392 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.911857s (all_items=0.000020s, per_item={'model_params': '0.00001621', 'optimizer_state': '0.00000405'}, temp_paths=0.911836) I0416 21:06:57.729208 131141424139392 event_tracking.py:125] [process=0] [async] Finished blocking save in 1.02 seconds. Continuing save @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1. I0416 21:06:57.729418 131031313155648 async_checkpointer.py:76] [process=0][thread=async_save] Background save thread started. Deadline for this save operation is 2026-04-16 21:26:57.729383 I0416 21:06:57.729664 131141424139392 checkpoint_manager.py:1560] [process=0][thread=MainThread][step=1] Starting CheckpointManager Save Finalize thread=save_finalize I0416 21:06:57.730023 131141424139392 standard_logger.py:34] {'step': 1, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_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': 1776373616.713666, 'wait_for_prev_duration_secs': 9.775161743164062e-05, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776373616.7139091, 'checkpointer_blocking_duration_secs': 1.0156042575836182, 'get_old_steps_start_time': 1776373617.7295296, 'get_old_steps_duration_secs': 9.703636169433594e-05, 'checkpoint_manager_blocking_start_time': 1776373616.7135963, 'checkpoint_manager_blocking_duration_secs': 1.0164036750793457} I0416 21:06:57.730171 131141424139392 profiler.py:85] Starting JAX profiler at step 1. I0416 21:06:57.858526 131031376070208 checkpoint.py:188] Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776373617531409892, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1776373617531409892, "commit_timestamp_nsecs": null, "custom_metadata": {}} to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/_CHECKPOINT_METADATA I0416 21:06:57.859720 131031428499008 async_checkpointer.py:280] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0416 21:06:58.049573 131141424139392 peft_trainer.py:485] Train step 1 training loss: 6.124308 - training perplexity: 456.828308 Training: 0%| | 0/5 [00:18<?, ?step/s, _train_loss=6.12, _train_perplexity=457, _train_steps_per_sec=0.058] Training: 20%|██ | 1/5 [00:18<01:14, 18.62s/step, _train_loss=6.12, _train_perplexity=457, _train_steps_per_sec=0.058]I0416 21:06:58.050524 131141424139392 max_utils.py:750] Memstats: After params initialized: I0416 21:06:58.050608 131141424139392 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_0(process=0,(0,0,0,0)) I0416 21:06:58.050659 131141424139392 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_1(process=0,(1,0,0,0)) I0416 21:06:58.050703 131141424139392 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_2(process=0,(0,1,0,0)) I0416 21:06:58.050742 131141424139392 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_3(process=0,(1,1,0,0)) I0416 21:06:58.050779 131141424139392 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_4(process=0,(0,2,0,0)) I0416 21:06:58.050815 131141424139392 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_5(process=0,(1,2,0,0)) I0416 21:06:58.050850 131141424139392 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_6(process=0,(0,3,0,0)) I0416 21:06:58.050885 131141424139392 max_utils.py:756] Using (GB) 0.01 / 31.25 (0.032000%) on TPU_7(process=0,(1,3,0,0)) I0416 21:06:58.184929 131141424139392 metric_logger.py:185] completed step: 1, seconds: 18.617, TFLOP/s/device: 0.000, Tokens/s/device: 55.005, total_weights: 6826, loss: 6.124 I0416 21:06:58.195772 131141424139392 peft_trainer.py:485] Train step 2 training loss: 6.129846 - training perplexity: 459.365448 Training: 20%|██ | 1/5 [00:18<01:14, 18.62s/step, _train_loss=6.13, _train_perplexity=458, _train_steps_per_sec=0.404] Training: 40%|████ | 2/5 [00:18<00:23, 7.75s/step, _train_loss=6.13, _train_perplexity=458, _train_steps_per_sec=0.404]I0416 21:06:58.197339 131141424139392 metric_logger.py:185] completed step: 2, seconds: 0.146, TFLOP/s/device: 0.002, Tokens/s/device: 7035.691, total_weights: 4636, loss: 6.130 I0416 21:06:58.211188 131141424139392 peft_trainer.py:485] Train step 3 training loss: 5.572245 - training perplexity: 263.023956 Training: 40%|████ | 2/5 [00:18<00:23, 7.75s/step, _train_loss=5.94, _train_perplexity=381, _train_steps_per_sec=2.5] I0416 21:06:58.212568 131141424139392 metric_logger.py:185] completed step: 3, seconds: 0.015, TFLOP/s/device: 0.014, Tokens/s/device: 66797.783, total_weights: 5886, loss: 5.572 I0416 21:06:58.222007 131141424139392 peft_trainer.py:485] Train step 4 training loss: 5.824473 - training perplexity: 338.482697 Training: 60%|██████ | 3/5 [00:18<00:15, 7.75s/step, _train_loss=5.91, _train_perplexity=370, _train_steps_per_sec=18.1]I0416 21:06:58.223372 131141424139392 metric_logger.py:185] completed step: 4, seconds: 0.011, TFLOP/s/device: 0.020, Tokens/s/device: 94734.379, total_weights: 4990, loss: 5.824 I0416 21:06:58.223663 131141424139392 profiler.py:113] Stopping JAX profiler at step 5. I0416 21:06:58.316122 4156963 google_auth_provider.cc:149] Using credentials at ~/.config/gcloud/application_default_credentials.json I0416 21:06:58.316198 4156963 google_auth_provider.cc:156] Using OAuth2 AuthProvider I0416 21:06:59.120499 131031365584448 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 22 array_metadata.ArrayMetadata to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/model_params/array_metadatas/process_0 I0416 21:06:59.126204 131031344612928 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_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/optimizer_state/array_metadatas/process_0 I0416 21:07:00.805595 131031334127168 base_pytree_checkpoint_handler.py:1282] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 2.493249s (commit=2.042352s, array_metadata_write=0.450896s) I0416 21:07:00.806669 131031313155648 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/gbytes_per_sec: 51.383 KiB/s (total gbytes: 205.0 KiB) (time elapsed: 3.9891133308410645 s) (per-host) I0416 21:07:00.875975 131031323641408 base_pytree_checkpoint_handler.py:1282] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 2.555255s (commit=2.130029s, array_metadata_write=0.425225s) I0416 21:07:00.877107 131031313155648 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/gbytes_per_sec: 151.715 KiB/s (total gbytes: 614.8 KiB) (time elapsed: 4.052513837814331 s) (per-host) I0416 21:07:00.877363 131031313155648 async_checkpointer.py:90] [process=0][thread=async_save] 4 Handler Commit operations completed. Time taken: 3.147656s. I0416 21:07:01.076204 131031313155648 checkpoint.py:228] Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1776373617531409892, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/_CHECKPOINT_METADATA I0416 21:07:01.259173 131031313155648 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:07:01.482818 131031376070208 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': 1776373617531409892, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/_CHECKPOINT_METADATA I0416 21:07:01.723141 131031313155648 base_pytree_checkpoint_handler.py:1406] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.606190s. use_zarr3=False, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/model_params I0416 21:07:01.723968 131031313155648 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/model_params I0416 21:07:02.147941 131031313155648 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:07:02.599188 131031313155648 base_pytree_checkpoint_handler.py:1406] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.596655s. use_zarr3=False, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/optimizer_state I0416 21:07:02.600133 131031313155648 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1/optimizer_state I0416 21:07:02.893655 131031313155648 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1 I0416 21:07:03.035638 131141424139392 utils.py:86] Train loop finished in: 23.6012 seconds I0416 21:07:03.036417 131141424139392 peft_trainer.py:485] Train step 5 training loss: 5.844853 - training perplexity: 345.451904 Training: 80%|████████ | 4/5 [00:23<00:07, 7.75s/step, _train_loss=5.9, _train_perplexity=365, _train_steps_per_sec=32.8] Training: 100%|██████████| 5/5 [00:23<00:00, 3.36s/step, _train_loss=5.9, _train_perplexity=365, _train_steps_per_sec=32.8]I0416 21:07:03.038115 131141424139392 metric_logger.py:185] completed step: 5, seconds: 4.814, TFLOP/s/device: 0.000, Tokens/s/device: 212.701, total_weights: 4264, loss: 5.845 I0416 21:07:03.040626 131141424139392 checkpoint_manager.py:2020] [process=0][thread=MainThread][step=1][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0416 21:07:03.635948 131031313155648 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_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1`. I0416 21:07:03.636744 131031313155648 event_tracking.py:138] [process=0] [async] Finished save (blocking + background) in 6.92 seconds @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/1 I0416 21:07:03.636819 131031313155648 async_checkpointer.py:160] [process=0][thread=async_save] Background save thread done. Time taken: 5.907113s. I0416 21:07:03.636979 131031428499008 async_checkpointer.py:288] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0416 21:07:03.637105 131031428499008 async_checkpointer.py:298] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0416 21:07:03.637160 131031428499008 checkpoint_manager.py:2137] [process=0][thread=save_finalize][step=1] CheckpointManager Save Finalize is syncing with other hosts... I0416 21:07:03.637205 131031428499008 checkpoint_manager.py:2146] [process=0][thread=save_finalize][step=1] CheckpointManager Save Finalize is done on all hosts. I0416 21:07:03.637351 131141424139392 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:07:03.637683 131141424139392 checkpoint_manager.py:1512] [process=0] Saving checkpoint at step 5 I0416 21:07:03.637762 131141424139392 event_tracking.py:70] [process=0] [async] Started save checkpoint @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5. I0416 21:07:03.740402 131141424139392 jax_array_handlers.py:360] Scheduling D2H of 22 prioritized jax.Array. I0416 21:07:03.740522 131141424139392 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:07:03.752720 131141424139392 base_pytree_checkpoint_handler.py:154] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.013043s I0416 21:07:03.753026 131141424139392 jax_array_handlers.py:360] Scheduling D2H of 52 prioritized jax.Array. I0416 21:07:03.753084 131141424139392 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:07:03.784501 131141424139392 base_pytree_checkpoint_handler.py:154] [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.jax_array_handlers.ArrayHandler".serialize. Time taken: 0.031715s I0416 21:07:03.784797 131141424139392 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/blocking_gbytes_per_sec: 3.317 MiB/s (total gbytes: 205.0 KiB) (time elapsed: 0.06035447120666504 s) (per-host) I0416 21:07:03.784916 131141424139392 base_pytree_checkpoint_handler.py:768] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.060488s (batch_requests_ready=0.001481s, total_serialization_initiated=0.058696s, others=0.000311s) I0416 21:07:03.785264 131141424139392 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/blocking_gbytes_per_sec: 10.898 MiB/s (total gbytes: 614.8 KiB) (time elapsed: 0.05509686470031738 s) (per-host) I0416 21:07:03.785342 131141424139392 base_pytree_checkpoint_handler.py:768] [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.055185s (batch_requests_ready=0.002453s, total_serialization_initiated=0.052355s, others=0.000377s) I0416 21:07:03.785403 131141424139392 composite_checkpoint_handler.py:715] [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.061501s (all_items=0.000013s, per_item={'model_params': '0.00001025', 'optimizer_state': '0.00000262'}, temp_paths=0.061488) I0416 21:07:03.786124 131141424139392 event_tracking.py:125] [process=0] [async] Finished blocking save in 0.15 seconds. Continuing save @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5. I0416 21:07:03.786350 131031227172416 async_checkpointer.py:76] [process=0][thread=async_save] Background save thread started. Deadline for this save operation is 2026-04-16 21:27:03.786315 I0416 21:07:03.786589 131141424139392 checkpoint_manager.py:1560] [process=0][thread=MainThread][step=5] Starting CheckpointManager Save Finalize thread=save_finalize I0416 21:07:03.786915 131031153772096 async_checkpointer.py:280] [process=0][thread=save_finalize] Waiting for background save thread=async_save. I0416 21:07:03.787040 131141424139392 standard_logger.py:34] {'step': 5, 'event_type': 'save', 'directory': 'gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_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': 1776373623.0406003, 'wait_for_prev_duration_secs': 0.59686279296875, 'time_between_consecutive_saves_sec': None, 'checkpointer_blocking_start_time': 1776373623.6377094, 'checkpointer_blocking_duration_secs': 0.14873147010803223, 'get_old_steps_start_time': 1776373623.7864592, 'get_old_steps_duration_secs': 9.417533874511719e-05, 'checkpoint_manager_blocking_start_time': 1776373623.0405579, 'checkpoint_manager_blocking_duration_secs': 0.7464585304260254} I0416 21:07:03.787205 131141424139392 checkpoint_manager.py:2020] [process=0][thread=MainThread][step=5][wait_until_finished] Waiting for Save Finalize thread (save_finalize) to complete. I0416 21:07:03.816087 131031428499008 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5 I0416 21:07:04.514771 131031323641408 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/model_params I0416 21:07:04.518631 131031323641408 atomicity.py:140] Creating tmp directory gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/optimizer_state I0416 21:07:05.790271 131031365584448 array_metadata_store.py:203] [process=0][thread=array_type_handler] Wrote 22 array_metadata.ArrayMetadata to gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/model_params/array_metadatas/process_0 I0416 21:07:05.792644 131031334127168 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_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/optimizer_state/array_metadatas/process_0 I0416 21:07:07.334915 131031281698368 base_pytree_checkpoint_handler.py:1282] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 2.326957s (commit=1.907889s, array_metadata_write=0.419068s) I0416 21:07:07.336055 131031227172416 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/gbytes_per_sec: 56.755 KiB/s (total gbytes: 205.0 KiB) (time elapsed: 3.611562490463257 s) (per-host) I0416 21:07:07.369688 131031237658176 base_pytree_checkpoint_handler.py:1282] [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 2.362172s (commit=1.932085s, array_metadata_write=0.430088s) I0416 21:07:07.676827 131031227172416 base_pytree_checkpoint_handler.py:130] [process=0] /jax/orbax/write/gbytes_per_sec: 155.786 KiB/s (total gbytes: 614.8 KiB) (time elapsed: 3.9466214179992676 s) (per-host) I0416 21:07:07.677260 131031227172416 async_checkpointer.py:90] [process=0][thread=async_save] 4 Handler Commit operations completed. Time taken: 3.890627s. I0416 21:07:08.063488 131031227172416 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:07:08.507838 131031227172416 base_pytree_checkpoint_handler.py:1406] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.591783s. use_zarr3=False, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/model_params I0416 21:07:08.508700 131031227172416 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/model_params I0416 21:07:08.929417 131031227172416 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:07:09.405541 131031227172416 base_pytree_checkpoint_handler.py:1406] [process=0][thread=async_save] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.621231s. use_zarr3=False, enable_post_merge_validation=True, directory=gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/optimizer_state I0416 21:07:09.406440 131031227172416 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5/optimizer_state I0416 21:07:09.673276 131031227172416 atomicity.py:666] Finalizing gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5 I0416 21:07:10.382983 131031227172416 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_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5`. I0416 21:07:10.383759 131031227172416 event_tracking.py:138] [process=0] [async] Finished save (blocking + background) in 6.75 seconds @ gs://wanglance-maxtext/pt_ckpt_feat_nnx_post_train_fixes_20260416_210550/pt_sft_linen_feat_nnx_post_train_fixes_20260416_210550_01_sft_smoke/checkpoints/5 I0416 21:07:10.383828 131031227172416 async_checkpointer.py:160] [process=0][thread=async_save] Background save thread done. Time taken: 6.597198s. I0416 21:07:10.383990 131031153772096 async_checkpointer.py:288] [process=0][thread=save_finalize] Done with waiting for background save thread=async_save. I0416 21:07:10.384119 131031153772096 async_checkpointer.py:298] [process=0][thread=save_finalize] No errors found in background save thread=async_save. I0416 21:07:10.384172 131031153772096 checkpoint_manager.py:2137] [process=0][thread=save_finalize][step=5] CheckpointManager Save Finalize is syncing with other hosts... I0416 21:07:10.384212 131031153772096 checkpoint_manager.py:2146] [process=0][thread=save_finalize][step=5] CheckpointManager Save Finalize is done on all hosts. I0416 21:07:10.384423 131141424139392 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:07:10.384559 131141424139392 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=131141424139392 count=1 at 0x773e82bb7440>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x773e7d779010>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x773e7d779a60>, _write_futures=[]) I0416 21:07:10.384980 131141424139392 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=131141424139392 count=1 at 0x773e82bb7440>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x773e7d779010>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x773e7d779a60>, _write_futures=[]) I0416 21:07:10.385017 131141424139392 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=131141424139392 count=1 at 0x773e82bb7440>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x773e7d779010>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x773e7d779a60>, _write_futures=[]) Training: 100%|██████████| 5/5 [00:32<00:00, 6.41s/step, _train_loss=5.9, _train_perplexity=365, _train_steps_per_sec=32.8] [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 '