MaxView

XPK Start: Mon Apr 20 06:11:55 UTC 2026
2026-04-20 06:12:42.907706: 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)
I0420 06:12:46.476011 137022069569344 max_utils.py:273] Attempting to initialize the jax distributed system...
INFO:2026-04-20 06:12:55,516:jax._src.distributed:166: Connecting to JAX distributed service on mt-07-distill-smoke-jd1wx-slice-job-0-0.mt-07-distill-smoke-jd1wx:8482
I0420 06:12:55.516344 137022069569344 distributed.py:166] Connecting to JAX distributed service on mt-07-distill-smoke-jd1wx-slice-job-0-0.mt-07-distill-smoke-jd1wx:8482
I0420 06:13:11.152408 137022069569344 max_utils.py:284] Jax distributed system initialized!
I0420 06:13:17.478497 137022069569344 max_utils.py:244] Jax distributed system is already initialized.
I0420 06:13:17.946588 137022069569344 max_utils.py:244] Jax distributed system is already initialized.
I0420 06:13:17.947779 137022069569344 pyconfig.py:432] Config param abort_on_inf_loss: True
I0420 06:13:17.947829 137022069569344 pyconfig.py:432] Config param abort_on_nan_loss: True
I0420 06:13:17.947864 137022069569344 pyconfig.py:432] Config param act_quantization_calibration_method: absmax
I0420 06:13:17.947885 137022069569344 pyconfig.py:432] Config param activation_dropout_for_audio: 0.0
I0420 06:13:17.947906 137022069569344 pyconfig.py:432] Config param activation_function_for_audio: gelu
I0420 06:13:17.947923 137022069569344 pyconfig.py:432] Config param activations_in_float32: False
I0420 06:13:17.947942 137022069569344 pyconfig.py:432] Config param adam_b1: 0.9
I0420 06:13:17.947963 137022069569344 pyconfig.py:432] Config param adam_b2: 0.95
I0420 06:13:17.947981 137022069569344 pyconfig.py:432] Config param adam_eps: 1e-08
I0420 06:13:17.948003 137022069569344 pyconfig.py:432] Config param adam_eps_root: 0.0
I0420 06:13:17.948019 137022069569344 pyconfig.py:432] Config param adam_weight_decay: 0.1
I0420 06:13:17.948036 137022069569344 pyconfig.py:432] Config param adamw_mask: []
I0420 06:13:17.948052 137022069569344 pyconfig.py:432] Config param add_bos: True
I0420 06:13:17.948067 137022069569344 pyconfig.py:432] Config param add_eos: True
I0420 06:13:17.948084 137022069569344 pyconfig.py:432] Config param allow_split_physical_axes: False
I0420 06:13:17.948099 137022069569344 pyconfig.py:432] Config param ar_cache_axis_order: 1,2,0,3
I0420 06:13:17.948115 137022069569344 pyconfig.py:432] Config param async_checkpointing: True
I0420 06:13:17.948130 137022069569344 pyconfig.py:432] Config param async_scheduling: False
I0420 06:13:17.948147 137022069569344 pyconfig.py:432] Config param attention: dot_product
I0420 06:13:17.948162 137022069569344 pyconfig.py:432] Config param attention_bias: False
I0420 06:13:17.948177 137022069569344 pyconfig.py:432] Config param attention_dropout_for_audio: 0.0
I0420 06:13:17.948194 137022069569344 pyconfig.py:432] Config param attention_out: RematLocation.REMAT
I0420 06:13:17.948215 137022069569344 pyconfig.py:432] Config param attention_output_dim: -1
I0420 06:13:17.948230 137022069569344 pyconfig.py:432] Config param attention_sink: False
I0420 06:13:17.948247 137022069569344 pyconfig.py:432] Config param attention_type: global
I0420 06:13:17.948263 137022069569344 pyconfig.py:432] Config param attn_logits_soft_cap: None
I0420 06:13:17.948279 137022069569344 pyconfig.py:432] Config param audio_path: 
I0420 06:13:17.948305 137022069569344 pyconfig.py:432] Config param audio_placeholder: <|audio|>
I0420 06:13:17.948322 137022069569344 pyconfig.py:432] Config param autoregressive_decode_assert: 
I0420 06:13:17.948337 137022069569344 pyconfig.py:432] Config param base_config: base.yml
I0420 06:13:17.948354 137022069569344 pyconfig.py:432] Config param base_emb_dim: 16
I0420 06:13:17.948369 137022069569344 pyconfig.py:432] Config param base_mlp_dim: 64
I0420 06:13:17.948386 137022069569344 pyconfig.py:432] Config param base_moe_mlp_dim: -1
I0420 06:13:17.948400 137022069569344 pyconfig.py:432] Config param base_num_decoder_layers: 1
I0420 06:13:17.948417 137022069569344 pyconfig.py:432] Config param base_num_kv_heads: 2
I0420 06:13:17.948433 137022069569344 pyconfig.py:432] Config param base_num_query_heads: 2
I0420 06:13:17.948449 137022069569344 pyconfig.py:432] Config param base_output_directory: 
I0420 06:13:17.948465 137022069569344 pyconfig.py:432] Config param batch_size: 1
I0420 06:13:17.948480 137022069569344 pyconfig.py:432] Config param batch_split_factor: 1
I0420 06:13:17.948505 137022069569344 pyconfig.py:432] Config param beta_fast: 32
I0420 06:13:17.948520 137022069569344 pyconfig.py:432] Config param beta_slow: 1
I0420 06:13:17.948536 137022069569344 pyconfig.py:432] Config param bwd_quantization_calibration_method: absmax
I0420 06:13:17.948553 137022069569344 pyconfig.py:432] Config param capacity_factor: -1.0
I0420 06:13:17.948568 137022069569344 pyconfig.py:432] Config param cast_logits_to_fp32: True
I0420 06:13:17.948584 137022069569344 pyconfig.py:432] Config param chat_template: 
I0420 06:13:17.948600 137022069569344 pyconfig.py:432] Config param chat_template_path: 
I0420 06:13:17.948617 137022069569344 pyconfig.py:432] Config param checkpoint_conversion_fn: None
I0420 06:13:17.948633 137022069569344 pyconfig.py:432] Config param checkpoint_dir: None
I0420 06:13:17.948651 137022069569344 pyconfig.py:432] Config param checkpoint_is_quantized: False
I0420 06:13:17.948670 137022069569344 pyconfig.py:432] Config param checkpoint_period: 2000
I0420 06:13:17.948688 137022069569344 pyconfig.py:432] Config param checkpoint_storage_concurrent_gb: 96
I0420 06:13:17.948704 137022069569344 pyconfig.py:432] Config param checkpoint_storage_target_data_file_size_bytes: 2147483648
I0420 06:13:17.948720 137022069569344 pyconfig.py:432] Config param checkpoint_storage_use_ocdbt: True
I0420 06:13:17.948737 137022069569344 pyconfig.py:432] Config param checkpoint_storage_use_zarr3: True
I0420 06:13:17.948752 137022069569344 pyconfig.py:432] Config param checkpoint_todelete_full_path: None
I0420 06:13:17.948767 137022069569344 pyconfig.py:432] Config param checkpoint_todelete_subdir: None
I0420 06:13:17.948783 137022069569344 pyconfig.py:432] Config param chips_per_vm: 4
I0420 06:13:17.948798 137022069569344 pyconfig.py:432] Config param chunk_attn_window_size: 0
I0420 06:13:17.948814 137022069569344 pyconfig.py:432] Config param collect_stack_trace: False
I0420 06:13:17.948829 137022069569344 pyconfig.py:432] Config param colocated_python_checkpointing: False
I0420 06:13:17.948848 137022069569344 pyconfig.py:432] Config param colocated_python_data_input: False
I0420 06:13:17.948863 137022069569344 pyconfig.py:432] Config param compile_topology: 
I0420 06:13:17.948878 137022069569344 pyconfig.py:432] Config param compile_topology_num_slices: -1
I0420 06:13:17.948893 137022069569344 pyconfig.py:432] Config param compile_xla_flags: 
I0420 06:13:17.948908 137022069569344 pyconfig.py:432] Config param compiled_trainstep_file: 
I0420 06:13:17.948923 137022069569344 pyconfig.py:432] Config param compute_axis_order: 0,1,2,3
I0420 06:13:17.948938 137022069569344 pyconfig.py:432] Config param constant_bound_config: []
I0420 06:13:17.948954 137022069569344 pyconfig.py:432] Config param context: RematLocation.REMAT
I0420 06:13:17.948971 137022069569344 pyconfig.py:432] Config param context_parallel_load_balance: True
I0420 06:13:17.948985 137022069569344 pyconfig.py:432] Config param context_parallel_size: 1
I0420 06:13:17.949001 137022069569344 pyconfig.py:432] Config param context_parallel_strategy: all_gather
I0420 06:13:17.949015 137022069569344 pyconfig.py:432] Config param context_sharding: context
I0420 06:13:17.949030 137022069569344 pyconfig.py:432] Config param conv_chunksize_for_audio: 500
I0420 06:13:17.949044 137022069569344 pyconfig.py:432] Config param conv_stride_for_vit: 14
I0420 06:13:17.949060 137022069569344 pyconfig.py:432] Config param cost_estimate_flops_bwd: -1
I0420 06:13:17.949074 137022069569344 pyconfig.py:432] Config param cost_estimate_flops_fwd: -1
I0420 06:13:17.949089 137022069569344 pyconfig.py:432] Config param custom_mesh: 
I0420 06:13:17.949103 137022069569344 pyconfig.py:432] Config param custom_mesh_and_rule: 
I0420 06:13:17.949119 137022069569344 pyconfig.py:432] Config param d_model_for_audio: 256
I0420 06:13:17.949133 137022069569344 pyconfig.py:432] Config param data_sharding: (('data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive'),)
I0420 06:13:17.949152 137022069569344 pyconfig.py:432] Config param data_shuffle_seed: 0
I0420 06:13:17.949168 137022069569344 pyconfig.py:432] Config param dataset_name: c4/en:3.0.1
I0420 06:13:17.949183 137022069569344 pyconfig.py:432] Config param dataset_path: 
I0420 06:13:17.949200 137022069569344 pyconfig.py:432] Config param dataset_type: DatasetType.HF
I0420 06:13:17.949217 137022069569344 pyconfig.py:432] Config param dcn_autoregressive_parallelism: 1
I0420 06:13:17.949232 137022069569344 pyconfig.py:432] Config param dcn_context_autoregressive_parallelism: 1
I0420 06:13:17.949246 137022069569344 pyconfig.py:432] Config param dcn_context_parallelism: 1
I0420 06:13:17.949262 137022069569344 pyconfig.py:432] Config param dcn_data_parallelism: -1
I0420 06:13:17.949276 137022069569344 pyconfig.py:432] Config param dcn_diloco_parallelism: 1
I0420 06:13:17.949291 137022069569344 pyconfig.py:432] Config param dcn_expert_parallelism: 1
I0420 06:13:17.949337 137022069569344 pyconfig.py:432] Config param dcn_fsdp_parallelism: 1
I0420 06:13:17.949352 137022069569344 pyconfig.py:432] Config param dcn_fsdp_transpose_parallelism: 1
I0420 06:13:17.949367 137022069569344 pyconfig.py:432] Config param dcn_parallelism: [1, -1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
I0420 06:13:17.949385 137022069569344 pyconfig.py:432] Config param dcn_pipeline_parallelism: 1
I0420 06:13:17.949399 137022069569344 pyconfig.py:432] Config param dcn_sequence_parallelism: 1
I0420 06:13:17.949413 137022069569344 pyconfig.py:432] Config param dcn_tensor_parallelism: 1
I0420 06:13:17.949428 137022069569344 pyconfig.py:432] Config param dcn_tensor_sequence_parallelism: 1
I0420 06:13:17.949442 137022069569344 pyconfig.py:432] Config param dcn_tensor_transpose_parallelism: 1
I0420 06:13:17.949458 137022069569344 pyconfig.py:432] Config param debug: {'rl': False}
I0420 06:13:17.949474 137022069569344 pyconfig.py:432] Config param debug_sharding: False
I0420 06:13:17.949489 137022069569344 pyconfig.py:432] Config param decode_sampling_nucleus_p: -1
I0420 06:13:17.949504 137022069569344 pyconfig.py:432] Config param decode_sampling_strategy: SamplingStrategy.GREEDY
I0420 06:13:17.949522 137022069569344 pyconfig.py:432] Config param decode_sampling_temperature: 1.0
I0420 06:13:17.949537 137022069569344 pyconfig.py:432] Config param decode_sampling_top_k: 0
I0420 06:13:17.949552 137022069569344 pyconfig.py:432] Config param decoder_block: DecoderBlockType.GPT3
I0420 06:13:17.949569 137022069569344 pyconfig.py:432] Config param decoder_layer_input: RematLocation.DEVICE
I0420 06:13:17.949584 137022069569344 pyconfig.py:432] Config param deepstack_visual_indexes_for_vit: []
I0420 06:13:17.949599 137022069569344 pyconfig.py:432] Config param degenerate_group_masking: True
I0420 06:13:17.949615 137022069569344 pyconfig.py:432] Config param dense_init_scale: 1.0
I0420 06:13:17.949629 137022069569344 pyconfig.py:432] Config param diloco_outer_lr: 0.3
I0420 06:13:17.949645 137022069569344 pyconfig.py:432] Config param diloco_outer_momentum: 0.9
I0420 06:13:17.949662 137022069569344 pyconfig.py:432] Config param diloco_sync_period: 36
I0420 06:13:17.949678 137022069569344 pyconfig.py:432] Config param distill_alpha: 0.5
I0420 06:13:17.949693 137022069569344 pyconfig.py:432] Config param distill_alpha_end: None
I0420 06:13:17.949709 137022069569344 pyconfig.py:432] Config param distill_alpha_schedule: constant
I0420 06:13:17.949723 137022069569344 pyconfig.py:432] Config param distill_beta: 0.0
I0420 06:13:17.949739 137022069569344 pyconfig.py:432] Config param distill_beta_end: None
I0420 06:13:17.949753 137022069569344 pyconfig.py:432] Config param distill_beta_schedule: constant
I0420 06:13:17.949769 137022069569344 pyconfig.py:432] Config param distill_feature_loss_type: cosine
I0420 06:13:17.949784 137022069569344 pyconfig.py:432] Config param distill_layer_indices: None
I0420 06:13:17.949798 137022069569344 pyconfig.py:432] Config param distill_temperature: 1.0
I0420 06:13:17.949814 137022069569344 pyconfig.py:432] Config param distill_temperature_end: None
I0420 06:13:17.949829 137022069569344 pyconfig.py:432] Config param distill_temperature_schedule: constant
I0420 06:13:17.949848 137022069569344 pyconfig.py:432] Config param downsample_hidden_size_for_audio: 256
I0420 06:13:17.949864 137022069569344 pyconfig.py:432] Config param dpo_beta: 0.1
I0420 06:13:17.949881 137022069569344 pyconfig.py:432] Config param dpo_label_smoothing: 0.0
I0420 06:13:17.949897 137022069569344 pyconfig.py:432] Config param dq_reduction_steps: 0
I0420 06:13:17.949912 137022069569344 pyconfig.py:432] Config param dropout_rate: 0.0
I0420 06:13:17.949928 137022069569344 pyconfig.py:432] Config param dtype: bfloat16
I0420 06:13:17.949956 137022069569344 pyconfig.py:432] Config param dtype_mm: float32
I0420 06:13:17.949971 137022069569344 pyconfig.py:432] Config param dump_hlo: False
I0420 06:13:17.949986 137022069569344 pyconfig.py:432] Config param dump_hlo_delete_local_after: True
I0420 06:13:17.950000 137022069569344 pyconfig.py:432] Config param dump_hlo_gcs_dir: gpt3-52k_2026-04-20-06-13/xla_dump
I0420 06:13:17.950015 137022069569344 pyconfig.py:432] Config param dump_hlo_local_dir: /tmp/xla_dump/
I0420 06:13:17.950030 137022069569344 pyconfig.py:432] Config param dump_hlo_local_module_name: jit_train_step
I0420 06:13:17.950045 137022069569344 pyconfig.py:432] Config param dump_hlo_module_name: jit_train_step
I0420 06:13:17.950060 137022069569344 pyconfig.py:432] Config param dump_hlo_upload_all: False
I0420 06:13:17.950076 137022069569344 pyconfig.py:432] Config param dump_hlo_xla_flags: 
I0420 06:13:17.950090 137022069569344 pyconfig.py:432] Config param dump_jaxpr: False
I0420 06:13:17.950105 137022069569344 pyconfig.py:432] Config param dump_jaxpr_delete_local_after: True
I0420 06:13:17.950120 137022069569344 pyconfig.py:432] Config param dump_jaxpr_gcs_dir: gpt3-52k_2026-04-20-06-13/jaxpr_dump
I0420 06:13:17.950143 137022069569344 pyconfig.py:432] Config param dump_jaxpr_local_dir: /tmp/jaxpr_dump/
I0420 06:13:17.950166 137022069569344 pyconfig.py:432] Config param dump_step: -1
I0420 06:13:17.950188 137022069569344 pyconfig.py:432] Config param elastic_enabled: False
I0420 06:13:17.950209 137022069569344 pyconfig.py:432] Config param elastic_max_retries: 10
I0420 06:13:17.950231 137022069569344 pyconfig.py:432] Config param elastic_timeout_seconds: 300
I0420 06:13:17.950253 137022069569344 pyconfig.py:432] Config param emb_dim: 16
I0420 06:13:17.950275 137022069569344 pyconfig.py:432] Config param enable_autocheckpoint: False
I0420 06:13:17.950312 137022069569344 pyconfig.py:432] Config param enable_checkpoint_cloud_logger: False
I0420 06:13:17.950339 137022069569344 pyconfig.py:432] Config param enable_checkpointing: True
I0420 06:13:17.950362 137022069569344 pyconfig.py:432] Config param enable_continuous_checkpointing: False
I0420 06:13:17.950383 137022069569344 pyconfig.py:432] Config param enable_data_shuffling: True
I0420 06:13:17.950405 137022069569344 pyconfig.py:432] Config param enable_diloco: False
I0420 06:13:17.950429 137022069569344 pyconfig.py:432] Config param enable_dp_attention: False
I0420 06:13:17.950451 137022069569344 pyconfig.py:432] Config param enable_dropout: False
I0420 06:13:17.950474 137022069569344 pyconfig.py:432] Config param enable_emergency_checkpoint: False
I0420 06:13:17.950499 137022069569344 pyconfig.py:432] Config param enable_expert_parallel: False
I0420 06:13:17.950523 137022069569344 pyconfig.py:432] Config param enable_gcp_goodput_metrics: True
I0420 06:13:17.950548 137022069569344 pyconfig.py:432] Config param enable_gcp_step_deviation_metrics: True
I0420 06:13:17.950573 137022069569344 pyconfig.py:432] Config param enable_goodput_recording: False
I0420 06:13:17.950598 137022069569344 pyconfig.py:432] Config param enable_jax_profiler: False
I0420 06:13:17.950622 137022069569344 pyconfig.py:432] Config param enable_llm_inference_pool: False
I0420 06:13:17.950647 137022069569344 pyconfig.py:432] Config param enable_model_warmup: False
I0420 06:13:17.950673 137022069569344 pyconfig.py:432] Config param enable_multi_tier_checkpointing: False
I0420 06:13:17.950697 137022069569344 pyconfig.py:432] Config param enable_nnx: False
I0420 06:13:17.950722 137022069569344 pyconfig.py:432] Config param enable_orbax_v1: False
I0420 06:13:17.950748 137022069569344 pyconfig.py:432] Config param enable_padding_causal_mask: True
I0420 06:13:17.950774 137022069569344 pyconfig.py:432] Config param enable_pathways_goodput: False
I0420 06:13:17.950799 137022069569344 pyconfig.py:432] Config param enable_prefix_caching: False
I0420 06:13:17.950824 137022069569344 pyconfig.py:432] Config param enable_rampup_batch_size: False
I0420 06:13:17.950856 137022069569344 pyconfig.py:432] Config param enable_single_controller: False
I0420 06:13:17.950883 137022069569344 pyconfig.py:432] Config param enable_single_replica_ckpt_restoring: False
I0420 06:13:17.950909 137022069569344 pyconfig.py:432] Config param enable_tensorboard: True
I0420 06:13:17.950935 137022069569344 pyconfig.py:432] Config param enable_tunix_perf_metrics: False
I0420 06:13:17.950961 137022069569344 pyconfig.py:432] Config param encoder_attention_heads_for_audio: 4
I0420 06:13:17.950986 137022069569344 pyconfig.py:432] Config param encoder_ffn_dim_for_audio: 512
I0420 06:13:17.951011 137022069569344 pyconfig.py:432] Config param encoder_layers_for_audio: 2
I0420 06:13:17.951036 137022069569344 pyconfig.py:432] Config param engram: RematLocation.REMAT
I0420 06:13:17.951063 137022069569344 pyconfig.py:432] Config param engram_head_dim: 1280
I0420 06:13:17.951087 137022069569344 pyconfig.py:432] Config param engram_kernel_size: 4
I0420 06:13:17.951111 137022069569344 pyconfig.py:432] Config param engram_layers: []
I0420 06:13:17.951135 137022069569344 pyconfig.py:432] Config param engram_max_ngram_size: 3
I0420 06:13:17.951160 137022069569344 pyconfig.py:432] Config param engram_num_heads: 8
I0420 06:13:17.951185 137022069569344 pyconfig.py:432] Config param engram_seed: 0
I0420 06:13:17.951210 137022069569344 pyconfig.py:432] Config param engram_vocab_bases: []
I0420 06:13:17.951234 137022069569344 pyconfig.py:432] Config param epsilon_high: None
I0420 06:13:17.951260 137022069569344 pyconfig.py:432] Config param eval_corr_lst: False
I0420 06:13:17.951284 137022069569344 pyconfig.py:432] Config param eval_data_columns: ['text']
I0420 06:13:17.951323 137022069569344 pyconfig.py:432] Config param eval_dataset_name: c4/en:3.0.1
I0420 06:13:17.951350 137022069569344 pyconfig.py:432] Config param eval_image_column: image
I0420 06:13:17.951377 137022069569344 pyconfig.py:432] Config param eval_interval: -1
I0420 06:13:17.951401 137022069569344 pyconfig.py:432] Config param eval_make_lst: False
I0420 06:13:17.951426 137022069569344 pyconfig.py:432] Config param eval_per_device_batch_size: 2
I0420 06:13:17.951450 137022069569344 pyconfig.py:432] Config param eval_sampling_strategy: greedy
I0420 06:13:17.951475 137022069569344 pyconfig.py:432] Config param eval_split: validation
I0420 06:13:17.951499 137022069569344 pyconfig.py:432] Config param eval_steps: -1
I0420 06:13:17.951524 137022069569344 pyconfig.py:432] Config param expansion_factor_real_data: -1.0
I0420 06:13:17.951550 137022069569344 pyconfig.py:432] Config param final_logits_soft_cap: None
I0420 06:13:17.951575 137022069569344 pyconfig.py:432] Config param first_num_dense_layers: 0
I0420 06:13:17.951599 137022069569344 pyconfig.py:432] Config param float32_gate_logits: False
I0420 06:13:17.951623 137022069569344 pyconfig.py:432] Config param float32_logits: False
I0420 06:13:17.951648 137022069569344 pyconfig.py:432] Config param float32_qk_product: False
I0420 06:13:17.951672 137022069569344 pyconfig.py:432] Config param float32_weight_sum: True
I0420 06:13:17.951697 137022069569344 pyconfig.py:432] Config param force_q_layout: False
I0420 06:13:17.951722 137022069569344 pyconfig.py:432] Config param force_unroll: False
I0420 06:13:17.951748 137022069569344 pyconfig.py:432] Config param freeze_audio_encoder_params: True
I0420 06:13:17.951773 137022069569344 pyconfig.py:432] Config param freeze_vision_encoder_params: True
I0420 06:13:17.951797 137022069569344 pyconfig.py:432] Config param fused_mlp: False
I0420 06:13:17.951822 137022069569344 pyconfig.py:432] Config param fused_qkv: True
I0420 06:13:17.951851 137022069569344 pyconfig.py:432] Config param gcs_metrics: False
I0420 06:13:17.951877 137022069569344 pyconfig.py:432] Config param gdn_chunk_size: 64
I0420 06:13:17.951901 137022069569344 pyconfig.py:432] Config param gdn_conv_kernel_dim: 4
I0420 06:13:17.951926 137022069569344 pyconfig.py:432] Config param gdn_key_head_dim: 128
I0420 06:13:17.951951 137022069569344 pyconfig.py:432] Config param gdn_num_key_heads: 16
I0420 06:13:17.951975 137022069569344 pyconfig.py:432] Config param gdn_num_value_heads: 32
I0420 06:13:17.951999 137022069569344 pyconfig.py:432] Config param gdn_value_head_dim: 128
I0420 06:13:17.952023 137022069569344 pyconfig.py:432] Config param generate_padding_batch_eval: False
I0420 06:13:17.952047 137022069569344 pyconfig.py:432] Config param generate_padding_batch_train: False
I0420 06:13:17.952072 137022069569344 pyconfig.py:432] Config param generate_slice: v5e-16
I0420 06:13:17.952096 137022069569344 pyconfig.py:432] Config param generation_configs: {}
I0420 06:13:17.952122 137022069569344 pyconfig.py:432] Config param global_batch_size_to_eval_on: 64
I0420 06:13:17.952146 137022069569344 pyconfig.py:432] Config param global_batch_size_to_load: 512
I0420 06:13:17.952171 137022069569344 pyconfig.py:432] Config param global_batch_size_to_load_eval: 64
I0420 06:13:17.952196 137022069569344 pyconfig.py:432] Config param global_batch_size_to_load_increment: None
I0420 06:13:17.952220 137022069569344 pyconfig.py:432] Config param global_batch_size_to_load_start: None
I0420 06:13:17.952244 137022069569344 pyconfig.py:432] Config param global_batch_size_to_train_on: 512
I0420 06:13:17.952270 137022069569344 pyconfig.py:432] Config param global_head_dim: 0
I0420 06:13:17.952294 137022069569344 pyconfig.py:432] Config param global_num_kv_heads: 0
I0420 06:13:17.952333 137022069569344 pyconfig.py:432] Config param global_parameter_scale: 1
I0420 06:13:17.952358 137022069569344 pyconfig.py:432] Config param global_rampup_samples: 500
I0420 06:13:17.952383 137022069569344 pyconfig.py:432] Config param global_rope_max_timescale: -1
I0420 06:13:17.952408 137022069569344 pyconfig.py:432] Config param global_rope_proportion: 0.25
I0420 06:13:17.952434 137022069569344 pyconfig.py:432] Config param goodput_upload_interval_seconds: 30
I0420 06:13:17.952459 137022069569344 pyconfig.py:432] Config param grad_dtype: float32
I0420 06:13:17.952513 137022069569344 pyconfig.py:432] Config param gradient_accumulation_steps: 8
I0420 06:13:17.952543 137022069569344 pyconfig.py:432] Config param gradient_clipping_threshold: 1.0
I0420 06:13:17.952569 137022069569344 pyconfig.py:432] Config param grain_data_source_max_workers: 16
I0420 06:13:17.952595 137022069569344 pyconfig.py:432] Config param grain_eval_files: 
I0420 06:13:17.952619 137022069569344 pyconfig.py:432] Config param grain_file_type: arrayrecord
I0420 06:13:17.952644 137022069569344 pyconfig.py:432] Config param grain_num_threads: 16
I0420 06:13:17.952668 137022069569344 pyconfig.py:432] Config param grain_num_threads_eval: 16
I0420 06:13:17.952692 137022069569344 pyconfig.py:432] Config param grain_packing_type: first_fit
I0420 06:13:17.952718 137022069569344 pyconfig.py:432] Config param grain_per_worker_buffer_size: 1
I0420 06:13:17.952744 137022069569344 pyconfig.py:432] Config param grain_per_worker_buffer_size_eval: 1
I0420 06:13:17.952769 137022069569344 pyconfig.py:432] Config param grain_prefetch_buffer_size: 500
I0420 06:13:17.952794 137022069569344 pyconfig.py:432] Config param grain_prefetch_buffer_size_eval: 500
I0420 06:13:17.952820 137022069569344 pyconfig.py:432] Config param grain_ram_budget_mb: 1024
I0420 06:13:17.952850 137022069569344 pyconfig.py:432] Config param grain_shuffle_buffer_size: 100
I0420 06:13:17.952876 137022069569344 pyconfig.py:432] Config param grain_train_files: 
I0420 06:13:17.952901 137022069569344 pyconfig.py:432] Config param grain_train_mixture_config_path: 
I0420 06:13:17.952925 137022069569344 pyconfig.py:432] Config param grain_worker_count: 1
I0420 06:13:17.952950 137022069569344 pyconfig.py:432] Config param grain_worker_count_eval: 1
I0420 06:13:17.952975 137022069569344 pyconfig.py:432] Config param grpo_beta: 0.08
I0420 06:13:17.953001 137022069569344 pyconfig.py:432] Config param grpo_epsilon: 0.2
I0420 06:13:17.953026 137022069569344 pyconfig.py:432] Config param hardware: tpu
I0420 06:13:17.953052 137022069569344 pyconfig.py:432] Config param hbm_utilization_vllm: 0.72
I0420 06:13:17.953080 137022069569344 pyconfig.py:432] Config param head_dim: 8
I0420 06:13:17.953105 137022069569344 pyconfig.py:432] Config param heartbeat_reporting_interval_in_seconds: 5
I0420 06:13:17.953130 137022069569344 pyconfig.py:432] Config param hf_data_dir: None
I0420 06:13:17.953155 137022069569344 pyconfig.py:432] Config param hf_eval_files: None
I0420 06:13:17.953180 137022069569344 pyconfig.py:432] Config param hf_eval_split: None
I0420 06:13:17.953205 137022069569344 pyconfig.py:432] Config param hf_name: None
I0420 06:13:17.953230 137022069569344 pyconfig.py:432] Config param hf_path: OptimalScale/ClimbMix
I0420 06:13:17.953254 137022069569344 pyconfig.py:432] Config param hf_train_files: None
I0420 06:13:17.953279 137022069569344 pyconfig.py:432] Config param hidden_size_for_vit: 1408
I0420 06:13:17.953316 137022069569344 pyconfig.py:432] Config param hide_profiler_step_metric: False
I0420 06:13:17.953345 137022069569344 pyconfig.py:432] Config param ici_autoregressive_parallelism: 1
I0420 06:13:17.953371 137022069569344 pyconfig.py:432] Config param ici_context_autoregressive_parallelism: 1
I0420 06:13:17.953396 137022069569344 pyconfig.py:432] Config param ici_context_parallelism: 1
I0420 06:13:17.953421 137022069569344 pyconfig.py:432] Config param ici_data_parallelism: 1
I0420 06:13:17.953445 137022069569344 pyconfig.py:432] Config param ici_diloco_parallelism: 1
I0420 06:13:17.953471 137022069569344 pyconfig.py:432] Config param ici_expert_parallelism: 1
I0420 06:13:17.953495 137022069569344 pyconfig.py:432] Config param ici_fsdp_parallelism: -1
I0420 06:13:17.953520 137022069569344 pyconfig.py:432] Config param ici_fsdp_transpose_parallelism: 1
I0420 06:13:17.953544 137022069569344 pyconfig.py:432] Config param ici_parallelism: [1, 1, 1, -1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
I0420 06:13:17.953571 137022069569344 pyconfig.py:432] Config param ici_pipeline_parallelism: 1
I0420 06:13:17.953596 137022069569344 pyconfig.py:432] Config param ici_sequence_parallelism: 1
I0420 06:13:17.953620 137022069569344 pyconfig.py:432] Config param ici_tensor_parallelism: 1
I0420 06:13:17.953644 137022069569344 pyconfig.py:432] Config param ici_tensor_sequence_parallelism: 1
I0420 06:13:17.953669 137022069569344 pyconfig.py:432] Config param ici_tensor_transpose_parallelism: 1
I0420 06:13:17.953693 137022069569344 pyconfig.py:432] Config param image_path: 
I0420 06:13:17.953717 137022069569344 pyconfig.py:432] Config param image_placeholder: <|image|>
I0420 06:13:17.953742 137022069569344 pyconfig.py:432] Config param image_size_for_vit: 896
I0420 06:13:17.953767 137022069569344 pyconfig.py:432] Config param indexer_head_dim: 128
I0420 06:13:17.953792 137022069569344 pyconfig.py:432] Config param indexer_loss_scaling_factor: 0.0
I0420 06:13:17.953821 137022069569344 pyconfig.py:432] Config param indexer_n_heads: 64
I0420 06:13:17.953855 137022069569344 pyconfig.py:432] Config param indexer_sparse_training: False
I0420 06:13:17.953882 137022069569344 pyconfig.py:432] Config param indexer_topk: 2048
I0420 06:13:17.953907 137022069569344 pyconfig.py:432] Config param inference_benchmark_test: False
I0420 06:13:17.953933 137022069569344 pyconfig.py:432] Config param inference_metadata_file: 
I0420 06:13:17.953958 137022069569344 pyconfig.py:432] Config param inference_microbenchmark_log_file_path: 
I0420 06:13:17.953984 137022069569344 pyconfig.py:432] Config param inference_microbenchmark_loop_iters: 10
I0420 06:13:17.954010 137022069569344 pyconfig.py:432] Config param inference_microbenchmark_num_samples: [1, 2, 3, 4, 5]
I0420 06:13:17.954040 137022069569344 pyconfig.py:432] Config param inference_microbenchmark_prefill_lengths: 64,128,256,512,1024
I0420 06:13:17.954066 137022069569344 pyconfig.py:432] Config param inference_microbenchmark_stages: prefill,generate
I0420 06:13:17.954092 137022069569344 pyconfig.py:432] Config param inference_server: MaxtextInterleavedServer
I0420 06:13:17.954117 137022069569344 pyconfig.py:432] Config param inhomogeneous_layer_cycle_interval: 1
I0420 06:13:17.954142 137022069569344 pyconfig.py:432] Config param init_weights_seed: 0
I0420 06:13:17.954167 137022069569344 pyconfig.py:432] Config param input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activation_norm_length']
I0420 06:13:17.954193 137022069569344 pyconfig.py:432] Config param interleave_moe_layer_step: 1
I0420 06:13:17.954218 137022069569344 pyconfig.py:432] Config param intermediate_size_for_vit: 5632
I0420 06:13:17.954242 137022069569344 pyconfig.py:432] Config param internal_compile: False
I0420 06:13:17.954267 137022069569344 pyconfig.py:432] Config param internal_compile_num_devices: -1
I0420 06:13:17.954292 137022069569344 pyconfig.py:432] Config param jax_cache_dir: ~/jax_cache
I0420 06:13:17.954333 137022069569344 pyconfig.py:432] Config param jax_debug_log_modules: 
I0420 06:13:17.954359 137022069569344 pyconfig.py:432] Config param jax_distributed_initialization_timeout: 300
I0420 06:13:17.954384 137022069569344 pyconfig.py:432] Config param jax_profiler_port: 9999
I0420 06:13:17.954409 137022069569344 pyconfig.py:432] Config param key_proj: RematLocation.REMAT
I0420 06:13:17.954435 137022069569344 pyconfig.py:432] Config param kv_cache_buffer: 256
I0420 06:13:17.954460 137022069569344 pyconfig.py:432] Config param kv_lora_rank: 512
I0420 06:13:17.954485 137022069569344 pyconfig.py:432] Config param kv_quant_axis: KvQuantAxis.HEADS_AND_DKV
I0420 06:13:17.954513 137022069569344 pyconfig.py:432] Config param kv_quant_dtype: int8
I0420 06:13:17.954539 137022069569344 pyconfig.py:432] Config param kv_wa_proj: RematLocation.REMAT
I0420 06:13:17.954564 137022069569344 pyconfig.py:432] Config param learning_rate: 0.0002
I0420 06:13:17.954590 137022069569344 pyconfig.py:432] Config param learning_rate_final_fraction: 0.1
I0420 06:13:17.954616 137022069569344 pyconfig.py:432] Config param learning_rate_schedule_steps: 200000
I0420 06:13:17.954640 137022069569344 pyconfig.py:432] Config param load_balance_loss_weight: 0.0
I0420 06:13:17.954666 137022069569344 pyconfig.py:432] Config param load_checkpoint_only_once: False
I0420 06:13:17.954690 137022069569344 pyconfig.py:432] Config param load_from_prefill_dir: False
I0420 06:13:17.954715 137022069569344 pyconfig.py:432] Config param load_full_state_path: 
I0420 06:13:17.954738 137022069569344 pyconfig.py:432] Config param load_parameters_path: gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items
I0420 06:13:17.954764 137022069569344 pyconfig.py:432] Config param local_checkpoint_directory: 
I0420 06:13:17.954789 137022069569344 pyconfig.py:432] Config param local_checkpoint_period: 0
I0420 06:13:17.954812 137022069569344 pyconfig.py:432] Config param local_rope_max_timescale: -1
I0420 06:13:17.954843 137022069569344 pyconfig.py:432] Config param local_rope_proportion: 1.0
I0420 06:13:17.954869 137022069569344 pyconfig.py:432] Config param log_config: True
I0420 06:13:17.954894 137022069569344 pyconfig.py:432] Config param log_period: 10
I0420 06:13:17.954918 137022069569344 pyconfig.py:432] Config param logical_axis_rules: (('activation_embed_and_logits_batch', ('data', 'stage', 'fsdp', 'fsdp_transpose', 'expert')), ('activation_embed_and_logits_batch_sequence', ('data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('activation_vocab', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('activation_vocab', ('tensor', 'tensor_transpose')), ('activation_vocab', 'tensor_sequence'), ('activation_vocab', ('sequence', 'context')), ('vocab', ('tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive')), ('embed_vocab', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('activation_heads', ('tensor', 'tensor_transpose', 'sequence', 'tensor_sequence', 'autoregressive')), ('activation_kv_heads', ('tensor', 'tensor_transpose', 'sequence', 'tensor_sequence')), ('activation_attn_length', ('sequence', 'context')), ('activation_attn_length', ('context',)), ('activation_q_length', ('context',)), ('activation_kv_length', ()), ('activation_attn_embed', ('tensor', 'tensor_transpose')), ('activation_kv', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('activation_kv_batch', ('data', 'fsdp', 'fsdp_transpose', 'expert')), ('activation_kv_head_dim', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('heads', ('tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive')), ('q_heads', ('tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive')), ('kv_heads', ('tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive')), ('qkv', ()), ('kv', ()), ('kv_head_dim', ()), ('q_lora', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'tensor_transpose', 'expert')), ('q_lora', ('fsdp', 'sequence', 'context', 'tensor_transpose', 'expert')), ('q_lora', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('q_lora', ('fsdp', 'sequence', 'context', 'expert')), ('q_lora_up_proj', ()), ('kv_lora', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'tensor_transpose', 'expert')), ('kv_lora', ('fsdp', 'sequence', 'context', 'tensor_transpose', 'expert')), ('kv_lora', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('kv_lora', ('fsdp', 'sequence', 'context', 'expert')), ('kv_lora_up_proj', ()), ('activation_batch_moe', ('data', 'fsdp', 'fsdp_transpose')), ('activation_length_moe', ('sequence', 'context')), ('activation_length_moe', ('context',)), ('activation_norm_length_moe', ('tensor_sequence', 'context', 'sequence')), ('activation_embed_moe', ('tensor', 'tensor_transpose')), ('activation_mlp_moe', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('activation_exp', ('expert',)), ('exp', 'expert'), ('mlp_moe', ('fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive')), ('embed_moe', ('fsdp', 'fsdp_transpose', 'sequence', 'tensor_transpose', 'context')), ('embed_moe', ('fsdp', 'sequence', 'tensor_transpose', 'context')), ('embed_moe', ('fsdp', 'fsdp_transpose', 'sequence', 'context')), ('embed_moe', ('fsdp', 'sequence', 'context')), ('activation_mlp', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('activation_batch', ('data', 'fsdp', 'fsdp_transpose', 'expert')), ('activation_length', ('sequence', 'context')), ('activation_length', ('context',)), ('activation_norm_length', ('tensor_sequence', 'context', 'sequence')), ('activation_embed', ('tensor', 'tensor_transpose')), ('activation_stage', 'stage'), ('mlp', ('fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive')), ('embed', ('fsdp', 'fsdp_transpose', 'sequence', 'tensor_transpose', 'context', 'expert')), ('embed', ('fsdp', 'sequence', 'tensor_transpose', 'context', 'expert')), ('embed', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('embed', ('fsdp', 'sequence', 'context', 'expert')), ('norm', ('tensor', 'tensor_transpose')), ('layers', 'stage'), ('diloco', 'diloco'), ('engram_dim', ('tensor',)), ('dense_layers', ()), ('moe_layers', ()), ('mhc', ()), ('prefill_activation_length', ('sequence', 'context')), ('prefill_activation_norm_length', ('tensor_sequence', 'context', 'sequence')), ('activation_prefill_kv_batch', ('data', 'fsdp', 'fsdp_transpose', 'expert')), ('decode_batch', ('data', 'fsdp', 'fsdp_transpose', 'expert')), ('decode_length', ('sequence',)), ('cache_heads', ('autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence')), ('cache_heads', ('autoregressive', 'tensor', 'tensor_sequence')), ('paged_kv_heads', ('tensor',)), ('cache_batch_prefill', ()), ('cache_batch', ()), ('cache_heads_none', ()), ('cache_kv', ()), ('cache_sequence', ()), ('num_pages', ()), ('tokens_per_page', ()), ('paged_kv_head_dim_size', ()), ('mlp_no_fsdp', ('tensor', 'tensor_sequence', 'autoregressive')), ('embed_tensor_transpose', ('tensor_transpose',)), ('exp_with_fsdp', 'fsdp'))
I0420 06:13:17.955034 137022069569344 pyconfig.py:432] Config param logits_dot_in_fp32: False
I0420 06:13:17.955061 137022069569344 pyconfig.py:432] Config param logits_via_embedding: True
I0420 06:13:17.955086 137022069569344 pyconfig.py:432] Config param lora_input_adapters_path: 
I0420 06:13:17.955111 137022069569344 pyconfig.py:432] Config param loss_algo: grpo
I0420 06:13:17.955136 137022069569344 pyconfig.py:432] Config param lr_schedule_type: LearningRateScheduleType.COSINE
I0420 06:13:17.955157 137022069569344 pyconfig.py:432] Config param managed_mldiagnostics: False
I0420 06:13:17.955172 137022069569344 pyconfig.py:432] Config param managed_mldiagnostics_dir: None
I0420 06:13:17.955188 137022069569344 pyconfig.py:432] Config param managed_mldiagnostics_run_group: 
I0420 06:13:17.955204 137022069569344 pyconfig.py:432] Config param matmul_precision: MatmulPrecision.DEFAULT
I0420 06:13:17.955222 137022069569344 pyconfig.py:432] Config param max_checkify: False
I0420 06:13:17.955237 137022069569344 pyconfig.py:432] Config param max_concurrency: 256
I0420 06:13:17.955252 137022069569344 pyconfig.py:432] Config param max_corpus_chars: 10000000
I0420 06:13:17.955267 137022069569344 pyconfig.py:432] Config param max_num_batched_tokens: None
I0420 06:13:17.955282 137022069569344 pyconfig.py:432] Config param max_num_checkpoints_to_keep: None
I0420 06:13:17.955306 137022069569344 pyconfig.py:432] Config param max_num_images_per_example: -1
I0420 06:13:17.955322 137022069569344 pyconfig.py:432] Config param max_num_seqs: None
I0420 06:13:17.955337 137022069569344 pyconfig.py:432] Config param max_position_embeddings: 163840
I0420 06:13:17.955353 137022069569344 pyconfig.py:432] Config param max_prefill_predict_length: 64
I0420 06:13:17.955369 137022069569344 pyconfig.py:432] Config param max_sample_len_for_audio: 10000
I0420 06:13:17.955383 137022069569344 pyconfig.py:432] Config param max_segments_per_seq: -1
I0420 06:13:17.955398 137022069569344 pyconfig.py:432] Config param max_source_positions_for_audio: 1500
I0420 06:13:17.955413 137022069569344 pyconfig.py:432] Config param max_target_length: 2048
I0420 06:13:17.955428 137022069569344 pyconfig.py:432] Config param max_timescale_for_audio: 10000.0
I0420 06:13:17.955443 137022069569344 pyconfig.py:432] Config param megablox: True
I0420 06:13:17.955458 137022069569344 pyconfig.py:432] Config param merge_gating_gmm: False
I0420 06:13:17.955472 137022069569344 pyconfig.py:432] Config param mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']
I0420 06:13:17.955491 137022069569344 pyconfig.py:432] Config param metrics_dir: None
I0420 06:13:17.955505 137022069569344 pyconfig.py:432] Config param metrics_file: 
I0420 06:13:17.955519 137022069569344 pyconfig.py:432] Config param mhc_expansion_rate: 1
I0420 06:13:17.955533 137022069569344 pyconfig.py:432] Config param micro_batch_size_to_eval_on: 64
I0420 06:13:17.955549 137022069569344 pyconfig.py:432] Config param micro_batch_size_to_train_on: 64
I0420 06:13:17.955564 137022069569344 pyconfig.py:432] Config param mla_kv: RematLocation.REMAT
I0420 06:13:17.955579 137022069569344 pyconfig.py:432] Config param mla_naive_kvcache: True
I0420 06:13:17.955594 137022069569344 pyconfig.py:432] Config param mla_q: RematLocation.REMAT
I0420 06:13:17.955608 137022069569344 pyconfig.py:432] Config param mlp_activations: ['gelu']
I0420 06:13:17.955624 137022069569344 pyconfig.py:432] Config param mlp_activations_limit: -1.0
I0420 06:13:17.955640 137022069569344 pyconfig.py:432] Config param mlp_bias: False
I0420 06:13:17.955654 137022069569344 pyconfig.py:432] Config param mlp_dim: 64
I0420 06:13:17.955669 137022069569344 pyconfig.py:432] Config param mlpwi: RematLocation.REMAT
I0420 06:13:17.955684 137022069569344 pyconfig.py:432] Config param mlpwi_0: RematLocation.REMAT
I0420 06:13:17.955699 137022069569344 pyconfig.py:432] Config param mlpwi_1: RematLocation.REMAT
I0420 06:13:17.955713 137022069569344 pyconfig.py:432] Config param mlpwo: RematLocation.REMAT
I0420 06:13:17.955729 137022069569344 pyconfig.py:432] Config param moba: False
I0420 06:13:17.955743 137022069569344 pyconfig.py:432] Config param moba_chunk_size: 1024
I0420 06:13:17.955758 137022069569344 pyconfig.py:432] Config param moba_topk: 8
I0420 06:13:17.955774 137022069569344 pyconfig.py:432] Config param model_call_mode: 
I0420 06:13:17.955790 137022069569344 pyconfig.py:432] Config param model_name: gpt3-52k
I0420 06:13:17.955805 137022069569344 pyconfig.py:432] Config param moe_expert_input_dim: -1
I0420 06:13:17.955821 137022069569344 pyconfig.py:432] Config param moe_fsdp_use_two_stage_all_gather: False
I0420 06:13:17.955835 137022069569344 pyconfig.py:432] Config param moe_mlp_dim: -1
I0420 06:13:17.955855 137022069569344 pyconfig.py:432] Config param moe_mlpwi_0: RematLocation.REMAT
I0420 06:13:17.955869 137022069569344 pyconfig.py:432] Config param moe_mlpwi_1: RematLocation.REMAT
I0420 06:13:17.955885 137022069569344 pyconfig.py:432] Config param moe_mlpwo: RematLocation.REMAT
I0420 06:13:17.955900 137022069569344 pyconfig.py:432] Config param monitor_goodput: False
I0420 06:13:17.955915 137022069569344 pyconfig.py:432] Config param monitor_step_time_deviation: True
I0420 06:13:17.955930 137022069569344 pyconfig.py:432] Config param mrope_section: [24, 20, 20]
I0420 06:13:17.955945 137022069569344 pyconfig.py:432] Config param mscale: 1.0
I0420 06:13:17.955960 137022069569344 pyconfig.py:432] Config param mtc_data_parallelism: 0
I0420 06:13:17.955976 137022069569344 pyconfig.py:432] Config param mtp_eval_target_module: 0
I0420 06:13:17.955992 137022069569344 pyconfig.py:432] Config param mtp_loss_scaling_factor: 0.1
I0420 06:13:17.956008 137022069569344 pyconfig.py:432] Config param mtp_num_layers: 0
I0420 06:13:17.956028 137022069569344 pyconfig.py:432] Config param mu_dtype: float32
I0420 06:13:17.956064 137022069569344 pyconfig.py:432] Config param multi_sampling: False
I0420 06:13:17.956081 137022069569344 pyconfig.py:432] Config param multi_tier_checkpointing_backup_interval_minutes: 0
I0420 06:13:17.956095 137022069569344 pyconfig.py:432] Config param muon_beta: 0.95
I0420 06:13:17.956111 137022069569344 pyconfig.py:432] Config param muon_consistent_rms: None
I0420 06:13:17.956126 137022069569344 pyconfig.py:432] Config param muon_weight_decay: 0.0
I0420 06:13:17.956140 137022069569344 pyconfig.py:432] Config param n_routing_groups: -1
I0420 06:13:17.956156 137022069569344 pyconfig.py:432] Config param n_window_for_audio: 50
I0420 06:13:17.956170 137022069569344 pyconfig.py:432] Config param n_window_infer_for_audio: 800
I0420 06:13:17.956185 137022069569344 pyconfig.py:432] Config param nope_layer_interval: -1
I0420 06:13:17.956199 137022069569344 pyconfig.py:432] Config param norm_topk_prob: False
I0420 06:13:17.956214 137022069569344 pyconfig.py:432] Config param normalization_layer_epsilon: 1e-05
I0420 06:13:17.956232 137022069569344 pyconfig.py:432] Config param normalize_embedding_logits: False
I0420 06:13:17.956246 137022069569344 pyconfig.py:432] Config param num_attention_heads_for_vit: 16
I0420 06:13:17.956260 137022069569344 pyconfig.py:432] Config param num_batches: 4
I0420 06:13:17.956276 137022069569344 pyconfig.py:432] Config param num_channels_for_vit: 3
I0420 06:13:17.956290 137022069569344 pyconfig.py:432] Config param num_conv_layers_for_audio: 3
I0420 06:13:17.956322 137022069569344 pyconfig.py:432] Config param num_decoder_layers: 1
I0420 06:13:17.956336 137022069569344 pyconfig.py:432] Config param num_diloco_replicas: 1
I0420 06:13:17.956352 137022069569344 pyconfig.py:432] Config param num_epoch: 1
I0420 06:13:17.956367 137022069569344 pyconfig.py:432] Config param num_eval_passes: 1
I0420 06:13:17.956383 137022069569344 pyconfig.py:432] Config param num_experts: 1
I0420 06:13:17.956398 137022069569344 pyconfig.py:432] Config param num_experts_per_tok: 1
I0420 06:13:17.956413 137022069569344 pyconfig.py:432] Config param num_generations: 2
I0420 06:13:17.956427 137022069569344 pyconfig.py:432] Config param num_hidden_layers_for_vit: 34
I0420 06:13:17.956442 137022069569344 pyconfig.py:432] Config param num_iterations: 1
I0420 06:13:17.956458 137022069569344 pyconfig.py:432] Config param num_kv_heads: 2
I0420 06:13:17.956473 137022069569344 pyconfig.py:432] Config param num_layers_per_pipeline_stage: 1
I0420 06:13:17.956488 137022069569344 pyconfig.py:432] Config param num_mel_bins_for_audio: 128
I0420 06:13:17.956503 137022069569344 pyconfig.py:432] Config param num_pipeline_microbatches: -1
I0420 06:13:17.956517 137022069569344 pyconfig.py:432] Config param num_pipeline_repeats: -1
I0420 06:13:17.956532 137022069569344 pyconfig.py:432] Config param num_position_embeddings_for_vit: 1024
I0420 06:13:17.956547 137022069569344 pyconfig.py:432] Config param num_query_heads: 2
I0420 06:13:17.956561 137022069569344 pyconfig.py:432] Config param num_samplers_slices: -1
I0420 06:13:17.956577 137022069569344 pyconfig.py:432] Config param num_slices: 1
I0420 06:13:17.956591 137022069569344 pyconfig.py:432] Config param num_target_devices: 32
I0420 06:13:17.956605 137022069569344 pyconfig.py:432] Config param num_test_batches: 5
I0420 06:13:17.956620 137022069569344 pyconfig.py:432] Config param num_trainer_slices: -1
I0420 06:13:17.956634 137022069569344 pyconfig.py:432] Config param num_vocab_tiling: 1
I0420 06:13:17.956650 137022069569344 pyconfig.py:432] Config param off_policy_steps: 0
I0420 06:13:17.956665 137022069569344 pyconfig.py:432] Config param offline_data_dir: None
I0420 06:13:17.956679 137022069569344 pyconfig.py:432] Config param opt_type: OptimizerType.ADAM_PAX
I0420 06:13:17.956696 137022069569344 pyconfig.py:432] Config param optimize_mesh_for_tpu_v6e: False
I0420 06:13:17.956710 137022069569344 pyconfig.py:432] Config param optimizer_memory_host_offload: False
I0420 06:13:17.956726 137022069569344 pyconfig.py:432] Config param original_max_position_embeddings: 4096
I0420 06:13:17.956741 137022069569344 pyconfig.py:432] Config param out_hidden_size_for_vit: 512
I0420 06:13:17.956756 137022069569344 pyconfig.py:432] Config param out_proj: RematLocation.REMAT
I0420 06:13:17.956772 137022069569344 pyconfig.py:432] Config param output_dim_for_audio: 512
I0420 06:13:17.956786 137022069569344 pyconfig.py:432] Config param override_logical_axis_rules: False
I0420 06:13:17.956801 137022069569344 pyconfig.py:432] Config param override_model_config: True
I0420 06:13:17.956816 137022069569344 pyconfig.py:432] Config param packing: True
I0420 06:13:17.956831 137022069569344 pyconfig.py:432] Config param pagedattn_head_dim_alignment: 128
I0420 06:13:17.956853 137022069569344 pyconfig.py:432] Config param pagedattn_max_pages_per_group: -1
I0420 06:13:17.956867 137022069569344 pyconfig.py:432] Config param pagedattn_num_pages: 64
I0420 06:13:17.956882 137022069569344 pyconfig.py:432] Config param pagedattn_pages_per_compute_block: 4
I0420 06:13:17.956898 137022069569344 pyconfig.py:432] Config param pagedattn_tokens_per_page: 32
I0420 06:13:17.956913 137022069569344 pyconfig.py:432] Config param param_scan_axis: 1
I0420 06:13:17.956928 137022069569344 pyconfig.py:432] Config param parameter_memory_host_offload: False
I0420 06:13:17.956943 137022069569344 pyconfig.py:432] Config param partial_rotary_factor: 1.0
I0420 06:13:17.956959 137022069569344 pyconfig.py:432] Config param patch_size_for_vit: 14
I0420 06:13:17.956974 137022069569344 pyconfig.py:432] Config param penalty_incorrect_answer: -1.0
I0420 06:13:17.956990 137022069569344 pyconfig.py:432] Config param penalty_incorrect_format: -0.5
I0420 06:13:17.957004 137022069569344 pyconfig.py:432] Config param per_device_batch_size: 2
I0420 06:13:17.957019 137022069569344 pyconfig.py:432] Config param per_device_batch_size_increment: 2.0
I0420 06:13:17.957035 137022069569344 pyconfig.py:432] Config param per_device_batch_size_start: 4.0
I0420 06:13:17.957050 137022069569344 pyconfig.py:432] Config param pipeline_delay_activation_forwarding: False
I0420 06:13:17.957066 137022069569344 pyconfig.py:432] Config param pipeline_fsdp_ag_once: False
I0420 06:13:17.957082 137022069569344 pyconfig.py:432] Config param pipeline_fsdp_ag_per_repeat: False
I0420 06:13:17.957097 137022069569344 pyconfig.py:432] Config param pipeline_parallel_layers: 1
I0420 06:13:17.957112 137022069569344 pyconfig.py:432] Config param pixel_shuffle_ratio_for_vit: 0.5
I0420 06:13:17.957127 137022069569344 pyconfig.py:432] Config param posemb_type_for_vit: learn
I0420 06:13:17.957142 137022069569344 pyconfig.py:432] Config param position_id_per_seconds: 25
I0420 06:13:17.957156 137022069569344 pyconfig.py:432] Config param prefill_cache_axis_order: 1,2,0,3
I0420 06:13:17.957172 137022069569344 pyconfig.py:432] Config param prefill_cache_dir: 
I0420 06:13:17.957186 137022069569344 pyconfig.py:432] Config param prefill_chunk_size: 256
I0420 06:13:17.957201 137022069569344 pyconfig.py:432] Config param prefill_slice: v5e-16
I0420 06:13:17.957217 137022069569344 pyconfig.py:432] Config param prefix_caching_dram_byte: 100000000000
I0420 06:13:17.957233 137022069569344 pyconfig.py:432] Config param prefix_caching_hbm_byte: 10000000000
I0420 06:13:17.957247 137022069569344 pyconfig.py:432] Config param profile_cleanly: True
I0420 06:13:17.957263 137022069569344 pyconfig.py:432] Config param profile_periodically_period: -1
I0420 06:13:17.957277 137022069569344 pyconfig.py:432] Config param profile_power_events: False
I0420 06:13:17.957293 137022069569344 pyconfig.py:432] Config param profiler: ProfilerType.NONE
I0420 06:13:17.957320 137022069569344 pyconfig.py:432] Config param profiler_steps: 5
I0420 06:13:17.957334 137022069569344 pyconfig.py:432] Config param projector_dropout_for_vit: 0.0
I0420 06:13:17.957350 137022069569344 pyconfig.py:432] Config param projector_input_dim_for_vit: 4096
I0420 06:13:17.957364 137022069569344 pyconfig.py:432] Config param projector_output_dim_for_vit: 4096
I0420 06:13:17.957380 137022069569344 pyconfig.py:432] Config param prometheus_port: 0
I0420 06:13:17.957394 137022069569344 pyconfig.py:432] Config param prompt: I love to
I0420 06:13:17.957409 137022069569344 pyconfig.py:432] Config param pure_nnx: False
I0420 06:13:17.957424 137022069569344 pyconfig.py:432] Config param pure_nnx_decoder: False
I0420 06:13:17.957438 137022069569344 pyconfig.py:432] Config param q_lora_rank: 0
I0420 06:13:17.957453 137022069569344 pyconfig.py:432] Config param qk_clip_threshold: 100.0
I0420 06:13:17.957467 137022069569344 pyconfig.py:432] Config param qk_nope_head_dim: 128
I0420 06:13:17.957483 137022069569344 pyconfig.py:432] Config param qk_norm_with_scale: True
I0420 06:13:17.957497 137022069569344 pyconfig.py:432] Config param qk_rope_head_dim: 64
I0420 06:13:17.957511 137022069569344 pyconfig.py:432] Config param qkv_proj: RematLocation.REMAT
I0420 06:13:17.957526 137022069569344 pyconfig.py:432] Config param quant_cfg_path: 
I0420 06:13:17.957540 137022069569344 pyconfig.py:432] Config param quantization: QuantizationType.NONE
I0420 06:13:17.957558 137022069569344 pyconfig.py:432] Config param quantization_local_shard_count: 4
I0420 06:13:17.957574 137022069569344 pyconfig.py:432] Config param quantize_kvcache: False
I0420 06:13:17.957589 137022069569344 pyconfig.py:432] Config param query_proj: RematLocation.REMAT
I0420 06:13:17.957606 137022069569344 pyconfig.py:432] Config param query_wa_proj: RematLocation.REMAT
I0420 06:13:17.957631 137022069569344 pyconfig.py:432] Config param ragged_block_size: 256
I0420 06:13:17.957656 137022069569344 pyconfig.py:432] Config param ragged_buffer_factor: -1.0
I0420 06:13:17.957681 137022069569344 pyconfig.py:432] Config param rampup_end_step: 0
I0420 06:13:17.957706 137022069569344 pyconfig.py:432] Config param rampup_samples_per_increment_to_load: None
I0420 06:13:17.957731 137022069569344 pyconfig.py:432] Config param reasoning_end_token: </reasoning>
I0420 06:13:17.957757 137022069569344 pyconfig.py:432] Config param reasoning_start_token: <reasoning>
I0420 06:13:17.957783 137022069569344 pyconfig.py:432] Config param record_internal_nn_metrics: 0
I0420 06:13:17.957809 137022069569344 pyconfig.py:432] Config param remat_policy: full
I0420 06:13:17.957835 137022069569344 pyconfig.py:432] Config param remat_policy_for_vit: minimal
I0420 06:13:17.957867 137022069569344 pyconfig.py:432] Config param remove_size_one_mesh_axis_from_type: True
I0420 06:13:17.957893 137022069569344 pyconfig.py:432] Config param replicate_quant_scale: False
I0420 06:13:17.957918 137022069569344 pyconfig.py:432] Config param replicator_backup_interval_minutes: 0
I0420 06:13:17.957942 137022069569344 pyconfig.py:432] Config param report_heartbeat_metric_for_gcp_monitoring: False
I0420 06:13:17.957968 137022069569344 pyconfig.py:432] Config param report_performance_metric_for_gcp_monitoring: False
I0420 06:13:17.957992 137022069569344 pyconfig.py:432] Config param reshape_q: False
I0420 06:13:17.958017 137022069569344 pyconfig.py:432] Config param return_log_prob: False
I0420 06:13:17.958042 137022069569344 pyconfig.py:432] Config param reuse_example_batch: 0
I0420 06:13:17.958067 137022069569344 pyconfig.py:432] Config param reward_exact_answer: 5.0
I0420 06:13:17.958092 137022069569344 pyconfig.py:432] Config param reward_exact_format_match: 3.0
I0420 06:13:17.958117 137022069569344 pyconfig.py:432] Config param reward_partial_format_match: 0.5
I0420 06:13:17.958142 137022069569344 pyconfig.py:432] Config param reward_ratio_guess_to_answer_high: 0.5
I0420 06:13:17.958168 137022069569344 pyconfig.py:432] Config param reward_ratio_guess_to_answer_low: 0.25
I0420 06:13:17.958192 137022069569344 pyconfig.py:432] Config param reward_white_space_format_match: 1.5
I0420 06:13:17.958218 137022069569344 pyconfig.py:432] Config param rl: {'num_generations': 2, 'num_iterations': 1, 'grpo_beta': 0.08, 'grpo_epsilon': 0.2, 'loss_algo': 'grpo', 'use_agentic_rollout': False, 'max_concurrency': 256, 'off_policy_steps': 0, 'system_prompt': '', 'degenerate_group_masking': True, 'epsilon_high': None}
I0420 06:13:17.958250 137022069569344 pyconfig.py:432] Config param rollout_data_parallelism: -1
I0420 06:13:17.958276 137022069569344 pyconfig.py:432] Config param rollout_expert_parallelism: 1
I0420 06:13:17.958314 137022069569344 pyconfig.py:432] Config param rollout_micro_batch_size: -1
I0420 06:13:17.958343 137022069569344 pyconfig.py:432] Config param rollout_tensor_parallelism: -1
I0420 06:13:17.958368 137022069569344 pyconfig.py:432] Config param rope_attention_scaling: False
I0420 06:13:17.958393 137022069569344 pyconfig.py:432] Config param rope_factor: 40
I0420 06:13:17.958417 137022069569344 pyconfig.py:432] Config param rope_interleave: True
I0420 06:13:17.958442 137022069569344 pyconfig.py:432] Config param rope_linear_scaling_factor: 1.0
I0420 06:13:17.958468 137022069569344 pyconfig.py:432] Config param rope_max_timescale: 10000
I0420 06:13:17.958492 137022069569344 pyconfig.py:432] Config param rope_min_timescale: 1
I0420 06:13:17.958517 137022069569344 pyconfig.py:432] Config param rope_theta_for_vit: 10000
I0420 06:13:17.958542 137022069569344 pyconfig.py:432] Config param rope_truncate: True
I0420 06:13:17.958567 137022069569344 pyconfig.py:432] Config param rope_type: RopeType.DEFAULT
I0420 06:13:17.958595 137022069569344 pyconfig.py:432] Config param rope_use_scale: True
I0420 06:13:17.958621 137022069569344 pyconfig.py:432] Config param routed_bias: False
I0420 06:13:17.958646 137022069569344 pyconfig.py:432] Config param routed_bias_update_rate: 0.0
I0420 06:13:17.958671 137022069569344 pyconfig.py:432] Config param routed_scaling_factor: 1.0
I0420 06:13:17.958695 137022069569344 pyconfig.py:432] Config param routed_score_func: 
I0420 06:13:17.958720 137022069569344 pyconfig.py:432] Config param run_name: gpt3-52k_2026-04-20-06-13
I0420 06:13:17.958745 137022069569344 pyconfig.py:432] Config param sa_block_kv: 512
I0420 06:13:17.958770 137022069569344 pyconfig.py:432] Config param sa_block_kv_compute: 512
I0420 06:13:17.958795 137022069569344 pyconfig.py:432] Config param sa_block_kv_dkv: 512
I0420 06:13:17.958819 137022069569344 pyconfig.py:432] Config param sa_block_kv_dkv_compute: 512
I0420 06:13:17.958850 137022069569344 pyconfig.py:432] Config param sa_block_kv_dq: 512
I0420 06:13:17.958875 137022069569344 pyconfig.py:432] Config param sa_block_q: 512
I0420 06:13:17.958899 137022069569344 pyconfig.py:432] Config param sa_block_q_dkv: 512
I0420 06:13:17.958923 137022069569344 pyconfig.py:432] Config param sa_block_q_dq: 512
I0420 06:13:17.958948 137022069569344 pyconfig.py:432] Config param sa_k_layout: HEAD_DIM_MINOR
I0420 06:13:17.958972 137022069569344 pyconfig.py:432] Config param sa_q_layout: HEAD_DIM_MINOR
I0420 06:13:17.958997 137022069569344 pyconfig.py:432] Config param sa_use_fused_bwd_kernel: False
I0420 06:13:17.959021 137022069569344 pyconfig.py:432] Config param sa_v_layout: HEAD_DIM_MINOR
I0420 06:13:17.959045 137022069569344 pyconfig.py:432] Config param sampler_devices_fraction: 0.5
I0420 06:13:17.959071 137022069569344 pyconfig.py:432] Config param save_checkpoint_on_completion: True
I0420 06:13:17.959095 137022069569344 pyconfig.py:432] Config param save_config_to_gcs: False
I0420 06:13:17.959120 137022069569344 pyconfig.py:432] Config param save_quantized_params_path: 
I0420 06:13:17.959144 137022069569344 pyconfig.py:432] Config param scale_embedding_for_audio: True
I0420 06:13:17.959168 137022069569344 pyconfig.py:432] Config param scan_layers: True
I0420 06:13:17.959192 137022069569344 pyconfig.py:432] Config param scan_layers_per_stage: False
I0420 06:13:17.959218 137022069569344 pyconfig.py:432] Config param scan_pipeline_iterations: True
I0420 06:13:17.959242 137022069569344 pyconfig.py:432] Config param scan_pipeline_repeats: False
I0420 06:13:17.959267 137022069569344 pyconfig.py:432] Config param set_remat_policy_on_layers_per_stage: False
I0420 06:13:17.959292 137022069569344 pyconfig.py:432] Config param set_remat_policy_on_pipeline_iterations: True
I0420 06:13:17.959330 137022069569344 pyconfig.py:432] Config param sft_train_on_completion_only: False
I0420 06:13:17.959355 137022069569344 pyconfig.py:432] Config param shard_exp_on_fsdp: False
I0420 06:13:17.959381 137022069569344 pyconfig.py:432] Config param shard_mode: ShardMode.AUTO
I0420 06:13:17.959408 137022069569344 pyconfig.py:432] Config param shard_optimizer_over_data: False
I0420 06:13:17.959433 137022069569344 pyconfig.py:432] Config param sharding_strategy: None
I0420 06:13:17.959458 137022069569344 pyconfig.py:432] Config param sharding_tolerance: 0.02
I0420 06:13:17.959483 137022069569344 pyconfig.py:432] Config param shardy: True
I0420 06:13:17.959508 137022069569344 pyconfig.py:432] Config param share_kv_projections: False
I0420 06:13:17.959533 137022069569344 pyconfig.py:432] Config param shared_experts: 0
I0420 06:13:17.959557 137022069569344 pyconfig.py:432] Config param sinkhorn_iterations: 20
I0420 06:13:17.959582 137022069569344 pyconfig.py:432] Config param skip_first_n_steps_for_profiler: 1
I0420 06:13:17.959608 137022069569344 pyconfig.py:432] Config param skip_jax_distributed_system: False
I0420 06:13:17.959633 137022069569344 pyconfig.py:432] Config param skip_step_interval: 128
I0420 06:13:17.959657 137022069569344 pyconfig.py:432] Config param skip_step_on_spikes: False
I0420 06:13:17.959681 137022069569344 pyconfig.py:432] Config param skip_step_scaling_factor: 6.0
I0420 06:13:17.959706 137022069569344 pyconfig.py:432] Config param sliding_window_size: 0
I0420 06:13:17.959731 137022069569344 pyconfig.py:432] Config param solution_end_token: </answer>
I0420 06:13:17.959757 137022069569344 pyconfig.py:432] Config param solution_start_token: <answer>
I0420 06:13:17.959784 137022069569344 pyconfig.py:432] Config param source_checkpoint_layout: orbax
I0420 06:13:17.959808 137022069569344 pyconfig.py:432] Config param sparse_matmul: True
I0420 06:13:17.959833 137022069569344 pyconfig.py:432] Config param spatial_merge_size_for_vit: 2
I0420 06:13:17.959863 137022069569344 pyconfig.py:432] Config param stack_prefill_result_cache: False
I0420 06:13:17.959888 137022069569344 pyconfig.py:432] Config param stack_trace_interval_seconds: 600
I0420 06:13:17.959913 137022069569344 pyconfig.py:432] Config param stack_trace_to_cloud: False
I0420 06:13:17.959938 137022069569344 pyconfig.py:432] Config param step_deviation_interval_seconds: 30
I0420 06:13:17.959962 137022069569344 pyconfig.py:432] Config param steps: 200000
I0420 06:13:17.959987 137022069569344 pyconfig.py:432] Config param stop_strings: None
I0420 06:13:17.960012 137022069569344 pyconfig.py:432] Config param student_overrides: {'model_name': 'llama3.1-8b'}
I0420 06:13:17.960040 137022069569344 pyconfig.py:432] Config param student_params_to_update: None
I0420 06:13:17.960066 137022069569344 pyconfig.py:432] Config param subslice_shape: 
I0420 06:13:17.960090 137022069569344 pyconfig.py:432] Config param swap_space_vllm_gb: 2
I0420 06:13:17.960115 137022069569344 pyconfig.py:432] Config param system_prompt: 
I0420 06:13:17.960141 137022069569344 pyconfig.py:432] Config param target_eval_loss: 0.0
I0420 06:13:17.960166 137022069569344 pyconfig.py:432] Config param teacher_overrides: {'model_name': 'llama3.1-8b'}
I0420 06:13:17.960192 137022069569344 pyconfig.py:432] Config param temperature_tuning: False
I0420 06:13:17.960217 137022069569344 pyconfig.py:432] Config param temporal_patch_size_for_vit: 2
I0420 06:13:17.960242 137022069569344 pyconfig.py:432] Config param tensorboard_dir: None
I0420 06:13:17.960267 137022069569344 pyconfig.py:432] Config param tensors_on_device: None
I0420 06:13:17.960292 137022069569344 pyconfig.py:432] Config param tensors_to_offload: None
I0420 06:13:17.960330 137022069569344 pyconfig.py:432] Config param test_batch_start_index: 0
I0420 06:13:17.960357 137022069569344 pyconfig.py:432] Config param tile_size_for_vit: 336
I0420 06:13:17.960382 137022069569344 pyconfig.py:432] Config param tokenize_eval_data: True
I0420 06:13:17.960406 137022069569344 pyconfig.py:432] Config param tokenize_train_data: True
I0420 06:13:17.960430 137022069569344 pyconfig.py:432] Config param tokenizer_path: meta-llama/Llama-3.1-8B
I0420 06:13:17.960456 137022069569344 pyconfig.py:432] Config param tokenizer_type: TokenizerType.HUGGINGFACE
I0420 06:13:17.960484 137022069569344 pyconfig.py:432] Config param topk_routing_group: -1
I0420 06:13:17.960508 137022069569344 pyconfig.py:432] Config param train_data_columns: ['text']
I0420 06:13:17.960534 137022069569344 pyconfig.py:432] Config param train_fraction: 1.0
I0420 06:13:17.960559 137022069569344 pyconfig.py:432] Config param train_image_column: image
I0420 06:13:17.960583 137022069569344 pyconfig.py:432] Config param train_micro_batch_size: -1
I0420 06:13:17.960608 137022069569344 pyconfig.py:432] Config param train_split: train
I0420 06:13:17.960633 137022069569344 pyconfig.py:432] Config param trainable_parameters_mask: []
I0420 06:13:17.960658 137022069569344 pyconfig.py:432] Config param trainable_position_size: 2048
I0420 06:13:17.960682 137022069569344 pyconfig.py:432] Config param trainer_devices_fraction: 0.5
I0420 06:13:17.960707 137022069569344 pyconfig.py:432] Config param upload_all_profiler_results: False
I0420 06:13:17.960732 137022069569344 pyconfig.py:432] Config param use_2d_fsdp_sharding: False
I0420 06:13:17.960757 137022069569344 pyconfig.py:432] Config param use_agentic_rollout: False
I0420 06:13:17.960781 137022069569344 pyconfig.py:432] Config param use_audio: False
I0420 06:13:17.960805 137022069569344 pyconfig.py:432] Config param use_audio_in_video: False
I0420 06:13:17.960829 137022069569344 pyconfig.py:432] Config param use_batch_split_schedule: False
I0420 06:13:17.960860 137022069569344 pyconfig.py:432] Config param use_chat_template: False
I0420 06:13:17.960885 137022069569344 pyconfig.py:432] Config param use_chunked_prefill: False
I0420 06:13:17.960910 137022069569344 pyconfig.py:432] Config param use_custom_sort_vjp: True
I0420 06:13:17.960934 137022069569344 pyconfig.py:432] Config param use_dpo: False
I0420 06:13:17.960959 137022069569344 pyconfig.py:432] Config param use_gather_mosaic_kernel: False
I0420 06:13:17.960983 137022069569344 pyconfig.py:432] Config param use_grpo: True
I0420 06:13:17.961007 137022069569344 pyconfig.py:432] Config param use_indexer: False
I0420 06:13:17.961032 137022069569344 pyconfig.py:432] Config param use_iota_embed: True
I0420 06:13:17.961057 137022069569344 pyconfig.py:432] Config param use_jax_splash: False
I0420 06:13:17.961081 137022069569344 pyconfig.py:432] Config param use_max_logit_estimate: -1
I0420 06:13:17.961105 137022069569344 pyconfig.py:432] Config param use_mrope: False
I0420 06:13:17.961130 137022069569344 pyconfig.py:432] Config param use_multimodal: False
I0420 06:13:17.961154 137022069569344 pyconfig.py:432] Config param use_pathways: True
I0420 06:13:17.961178 137022069569344 pyconfig.py:432] Config param use_post_attn_norm: False
I0420 06:13:17.961203 137022069569344 pyconfig.py:432] Config param use_post_ffw_norm: False
I0420 06:13:17.961228 137022069569344 pyconfig.py:432] Config param use_qk_clip: False
I0420 06:13:17.961252 137022069569344 pyconfig.py:432] Config param use_qk_norm: False
I0420 06:13:17.961277 137022069569344 pyconfig.py:432] Config param use_qk_norm_in_gdn: True
I0420 06:13:17.961314 137022069569344 pyconfig.py:432] Config param use_qwix_quantization: False
I0420 06:13:17.961342 137022069569344 pyconfig.py:432] Config param use_ragged_attention: False
I0420 06:13:17.961367 137022069569344 pyconfig.py:432] Config param use_random_routing: False
I0420 06:13:17.961392 137022069569344 pyconfig.py:432] Config param use_replicator_service: False
I0420 06:13:17.961416 137022069569344 pyconfig.py:432] Config param use_ring_of_experts: False
I0420 06:13:17.961440 137022069569344 pyconfig.py:432] Config param use_sft: False
I0420 06:13:17.961464 137022069569344 pyconfig.py:432] Config param use_splash_scheduler: False
I0420 06:13:17.961489 137022069569344 pyconfig.py:432] Config param use_tokamax_gmm: False
I0420 06:13:17.961513 137022069569344 pyconfig.py:432] Config param use_tokamax_splash: False
I0420 06:13:17.961537 137022069569344 pyconfig.py:432] Config param use_truncation: True
I0420 06:13:17.961561 137022069569344 pyconfig.py:432] Config param use_tunix_gradient_accumulation: False
I0420 06:13:17.961584 137022069569344 pyconfig.py:432] Config param use_untrainable_positional_embedding: False
I0420 06:13:17.961609 137022069569344 pyconfig.py:432] Config param use_vertex_tensorboard: False
I0420 06:13:17.961633 137022069569344 pyconfig.py:432] Config param using_pipeline_parallelism: False
I0420 06:13:17.961657 137022069569344 pyconfig.py:432] Config param v_head_dim: 128
I0420 06:13:17.961681 137022069569344 pyconfig.py:432] Config param v_norm_with_scale: True
I0420 06:13:17.961705 137022069569344 pyconfig.py:432] Config param value_proj: RematLocation.REMAT
I0420 06:13:17.961731 137022069569344 pyconfig.py:432] Config param vertex_tensorboard_project: 
I0420 06:13:17.961756 137022069569344 pyconfig.py:432] Config param vertex_tensorboard_region: 
I0420 06:13:17.961780 137022069569344 pyconfig.py:432] Config param video_path: 
I0420 06:13:17.961805 137022069569344 pyconfig.py:432] Config param video_placeholder: <|video|>
I0420 06:13:17.961830 137022069569344 pyconfig.py:432] Config param vision_output_dim_for_vit: 4096
I0420 06:13:17.961862 137022069569344 pyconfig.py:432] Config param vision_output_length: -1
I0420 06:13:17.961888 137022069569344 pyconfig.py:432] Config param vllm_additional_config: {}
I0420 06:13:17.961913 137022069569344 pyconfig.py:432] Config param vllm_hf_config_path: 
I0420 06:13:17.961937 137022069569344 pyconfig.py:432] Config param vllm_hf_overrides: {}
I0420 06:13:17.961962 137022069569344 pyconfig.py:432] Config param vocab_size: 32000
I0420 06:13:17.961988 137022069569344 pyconfig.py:432] Config param warmup_steps_fraction: 0.1
I0420 06:13:17.962014 137022069569344 pyconfig.py:432] Config param weight_dtype: float32
I0420 06:13:17.962054 137022069569344 pyconfig.py:432] Config param weight_quantization_calibration_method: absmax
I0420 06:13:17.962081 137022069569344 pyconfig.py:432] Config param wi_tile_dlhs_batch_seq: 512
I0420 06:13:17.962106 137022069569344 pyconfig.py:432] Config param wi_tile_dlhs_embed_dim: 1024
I0420 06:13:17.962130 137022069569344 pyconfig.py:432] Config param wi_tile_dlhs_mlp_dim: 1024
I0420 06:13:17.962155 137022069569344 pyconfig.py:432] Config param wi_tile_drhs_batch_seq: 512
I0420 06:13:17.962179 137022069569344 pyconfig.py:432] Config param wi_tile_drhs_embed_dim: 1024
I0420 06:13:17.962204 137022069569344 pyconfig.py:432] Config param wi_tile_drhs_mlp_dim: 1024
I0420 06:13:17.962228 137022069569344 pyconfig.py:432] Config param wi_tile_fwd_batch_seq: 512
I0420 06:13:17.962253 137022069569344 pyconfig.py:432] Config param wi_tile_fwd_embed_dim: 1024
I0420 06:13:17.962278 137022069569344 pyconfig.py:432] Config param wi_tile_fwd_mlp_dim: 1024
I0420 06:13:17.962312 137022069569344 pyconfig.py:432] Config param wo_tile_dlhs_batch_seq: 512
I0420 06:13:17.962332 137022069569344 pyconfig.py:432] Config param wo_tile_dlhs_embed_dim: 1024
I0420 06:13:17.962356 137022069569344 pyconfig.py:432] Config param wo_tile_dlhs_mlp_dim: 1024
I0420 06:13:17.962381 137022069569344 pyconfig.py:432] Config param wo_tile_drhs_batch_seq: 512
I0420 06:13:17.962405 137022069569344 pyconfig.py:432] Config param wo_tile_drhs_embed_dim: 1024
I0420 06:13:17.962429 137022069569344 pyconfig.py:432] Config param wo_tile_drhs_mlp_dim: 1024
I0420 06:13:17.962453 137022069569344 pyconfig.py:432] Config param wo_tile_fwd_batch_seq: 512
I0420 06:13:17.962479 137022069569344 pyconfig.py:432] Config param wo_tile_fwd_embed_dim: 1024
I0420 06:13:17.962504 137022069569344 pyconfig.py:432] Config param wo_tile_fwd_mlp_dim: 1024
I0420 06:13:17.962522 137022069569344 pyconfig.py:432] Config param wsd_decay_steps_fraction: 0.1
I0420 06:13:17.962538 137022069569344 pyconfig.py:432] Config param wsd_decay_style: WsdDecayStyle.LINEAR
I0420 06:13:17.962556 137022069569344 pyconfig.py:432] Config param xprof_e2e_enable_fw_power_level_event: False
I0420 06:13:17.962581 137022069569344 pyconfig.py:432] Config param xprof_e2e_enable_fw_thermal_event: False
I0420 06:13:17.962606 137022069569344 pyconfig.py:432] Config param xprof_e2e_enable_fw_throttle_event: False
I0420 06:13:17.962630 137022069569344 pyconfig.py:432] Config param xprof_tpu_power_trace_level: 0
I0420 06:13:17.962657 137022069569344 pyconfig.py:432] Config param z_loss_multiplier: 0.0
I0420 06:13:17.963130 137022069569344 tokenizer.py:245] Tokenizer path: meta-llama/Llama-2-7b-chat-hf
I0420 06:13:17.963175 137022069569344 tokenizer.py:224] Loading HF tokenizer: meta-llama/Llama-2-7b-chat-hf
I0420 06:13:21.983571 137022069569344 _schedule.py:129] A polynomial schedule was set with a non-positive `transition_steps` value; this results in a constant schedule with value `init_value`.
I0420 06:13:21.986649 137022069569344 maxtext_utils.py:1551] Num_devices: 32, shape (1, 4, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1)
I0420 06:13:21.986760 137022069569344 train_distill.py:596] Applying logical axis rules for model initialization and training...
I0420 06:13:21.986834 137022069569344 train_distill.py:600] Loading Student from ...
I0420 06:13:21.986862 137022069569344 train_distill.py:169] --- Student Configuration ---
I0420 06:13:21.986884 137022069569344 train_distill.py:170]   Model Name:      gpt3-52k
I0420 06:13:21.986906 137022069569344 train_distill.py:171]   Dimensions:      1 Layers, 16 Emb Dim, 8 Head Dim
I0420 06:13:21.986925 137022069569344 train_distill.py:174]   Attention Heads: 2 Query, 2 KV
I0420 06:13:21.986941 137022069569344 train_distill.py:175]   Vocab Size:      32000
I0420 06:13:21.986959 137022069569344 train_distill.py:176]   Checkpoint:      
I0420 06:13:21.986975 137022069569344 train_distill.py:465] Initializing model: gpt3-52k...
I0420 06:13:23.400933 137022069569344 train_distill.py:614] Loading Teacher from gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items...
I0420 06:13:23.401047 137022069569344 train_distill.py:169] --- Teacher Configuration ---
I0420 06:13:23.401077 137022069569344 train_distill.py:170]   Model Name:      gpt3-52k
I0420 06:13:23.401106 137022069569344 train_distill.py:171]   Dimensions:      1 Layers, 16 Emb Dim, 8 Head Dim
I0420 06:13:23.401129 137022069569344 train_distill.py:174]   Attention Heads: 2 Query, 2 KV
I0420 06:13:23.401148 137022069569344 train_distill.py:175]   Vocab Size:      32000
I0420 06:13:23.401167 137022069569344 train_distill.py:176]   Checkpoint:      gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items
I0420 06:13:23.401185 137022069569344 train_distill.py:465] Initializing model: gpt3-52k...
I0420 06:13:24.477232 137022069569344 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0420 06:13:24.477694 137022069569344 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 0x7c9e38243b00>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0420 06:13:24.477753 137022069569344 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.28
W0420 06:13:25.003862 137022069569344 checkpoint.py:202] Metadata file does not exist: gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items/_CHECKPOINT_METADATA
I0420 06:13:25.543234    1864 google_auth_provider.cc:181] Running on GCE, using service account 562977990677-compute@developer.gserviceaccount.com
I0420 06:13:26.809472 137022069569344 checkpointer.py:304] Restoring checkpoint from gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items.
W0420 06:13:29.054955 137022069569344 transform_utils.py:230] The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.
I0420 06:13:29.055356 137022069569344 transform_utils.py:288] The following keys are not loaded from the original tree after applying specified transforms: params/params/decoder/to_nnx__rngs/aqt/count, params/params/decoder/to_nnx__rngs/aqt/key, params/params/decoder/to_nnx__rngs/dropout/count, params/params/decoder/to_nnx__rngs/dropout/key, params/params/decoder/to_nnx__rngs/params/count, params/params/decoder/to_nnx__rngs/params/key
I0420 06:13:29.443317 137022069569344 checkpointer.py:318] Finished restoring checkpoint in 3.00 seconds from gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items.
I0420 06:13:30.139032 137022069569344 train_distill.py:640] Initializing Data Iterators via MaxText pipeline...
I0420 06:13:30.202599 137022069569344 config.py:112] TensorFlow version 2.20.0 available.
I0420 06:13:30.203083 137022069569344 config.py:125] JAX version 0.8.3 available.
E0420 06:13:32.488108 137022069569344 packing.py:209] PackAndBatchOperation is deprecated. Please use lazy_dataset.FirstFitPackIterDataset instead.
I0420 06:13:32.488345 137022069569344 data_loader.py:408] Adding CopyNumPyArrayToSharedMemory MapTransform.
I0420 06:13:32.491403 137022069569344 train_distill.py:410] Input Pipeline Checkpointing: DISABLED
I0420 06:13:32.491471 137022069569344 train_distill.py:414] Reason: Iterator 'MultiHostDataLoadIterator' is not recognized as Grain (dataset_type='DatasetType.HF', has_save=False)
I0420 06:13:32.491533 137022069569344 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0420 06:13:32.491611 137022069569344 base_pytree_checkpoint_handler.py:411] 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 0x7c9e38243b00>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0420 06:13:32.491653 137022069569344 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0420 06:13:32.491686 137022069569344 base_pytree_checkpoint_handler.py:411] 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 0x7c9e38243b00>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0420 06:13:32.491730 137022069569344 checkpoint_manager.py:702] [process=3][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers={'model_params': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c87e8da8830>, 'optimizer_state': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418063290>, 'custom_metadata': <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c84180631d0>}, handler_registry=None
I0420 06:13:32.491935 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "model_params". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c87e8da8830>` 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`.
I0420 06:13:32.491978 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "optimizer_state". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418063290>` 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`.
I0420 06:13:32.492005 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "custom_metadata". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c84180631d0>` 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`.
I0420 06:13:32.492029 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c8418062e40>` 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`.
I0420 06:13:32.492056 137022069569344 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 0x7c87e8da8830>, ('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c87e8da8830>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418063290>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418063290>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c84180631d0>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c84180631d0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c8418062e40>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c8418062e40>}).
I0420 06:13:32.492488 137022069569344 async_checkpointer.py:177] [process=3][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7c7f940e7920> timeout: 600 secs and primary_host=0 for async checkpoint writes
I0420 06:13:35.771604 137022069569344 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://lance-maxtext/pt_ckpt_xpk_main_20260420_060038/pt_distill_linen_xpk_main_20260420_060038_07_distill_smoke/checkpoints
I0420 06:13:35.773882 137022069569344 checkpoint_manager.py:921] [process=3][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=2000, 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=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=gs://lance-maxtext/pt_ckpt_xpk_main_20260420_060038/pt_distill_linen_xpk_main_20260420_060038_07_distill_smoke/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7c84180631a0>
I0420 06:13:35.773991 137022069569344 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0420 06:13:35.774054 137022069569344 base_pytree_checkpoint_handler.py:411] 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 0x7c9e38243b00>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0420 06:13:35.774089 137022069569344 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0420 06:13:35.774118 137022069569344 base_pytree_checkpoint_handler.py:411] 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 0x7c9e38243b00>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0420 06:13:35.774154 137022069569344 checkpoint_manager.py:1983] [process=3][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
I0420 06:13:35.774204 137022069569344 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=137022069569344 count=1 at 0x7c856820ec00>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7c8418062fc0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7c8418062f90>, _write_futures=[])
I0420 06:13:35.774572 137022069569344 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=137022069569344 count=1 at 0x7c856820ec00>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7c8418062fc0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7c8418062f90>, _write_futures=[])
I0420 06:13:35.774601 137022069569344 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=137022069569344 count=1 at 0x7c856820ec00>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7c8418062fc0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7c8418062f90>, _write_futures=[])
I0420 06:13:35.774632 137022069569344 checkpoint_manager.py:702] [process=3][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers={'model_params': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418063170>, 'optimizer_state': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418062390>, 'custom_metadata': <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c7f94055910>, 'iter': <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x7c7f94055880>}, handler_registry=None
I0420 06:13:35.774728 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "model_params". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418063170>` 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`.
I0420 06:13:35.774762 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "optimizer_state". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418062390>` 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`.
I0420 06:13:35.774785 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "custom_metadata". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c7f94055910>` 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`.
I0420 06:13:35.774812 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "iter". Adding handler `<maxtext.common.checkpointing.GrainCheckpointHandler object at 0x7c7f94055880>` for item "iter" and save args `<class 'maxtext.common.checkpointing.GrainCheckpointSave'>` and restore args `<class 'maxtext.common.checkpointing.GrainCheckpointRestore'>` to `_handler_registry`.
I0420 06:13:35.774834 137022069569344 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c7f94055d00>` 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`.
I0420 06:13:35.774864 137022069569344 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 0x7c8418063170>, ('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418063170>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418062390>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7c8418062390>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c7f94055910>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c7f94055910>, ('iter', <class 'maxtext.common.checkpointing.GrainCheckpointSave'>): <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x7c7f94055880>, ('iter', <class 'maxtext.common.checkpointing.GrainCheckpointRestore'>): <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x7c7f94055880>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c7f94055d00>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7c7f94055d00>}).
I0420 06:13:35.774933 137022069569344 async_checkpointer.py:177] [process=3][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7c7f940e7a60> timeout: 600 secs and primary_host=0 for async checkpoint writes
I0420 06:13:36.157115 137022069569344 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://lance-maxtext/pt_ckpt_xpk_main_20260420_060038/pt_distill_linen_xpk_main_20260420_060038_07_distill_smoke/checkpoints
I0420 06:13:36.441055 137022069569344 checkpoint_manager.py:921] [process=3][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=2000, 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=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=gs://lance-maxtext/pt_ckpt_xpk_main_20260420_060038/pt_distill_linen_xpk_main_20260420_060038_07_distill_smoke/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7c94685398e0>
I0420 06:13:36.441682 137022069569344 train_distill.py:691] Starting Distillation Training...
I0420 06:13:36.441788 137022069569344 peft_trainer.py:590] Training with mesh: Mesh('diloco': 1, 'data': 4, '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))
I0420 06:13:36.789440 137022069569344 peft_trainer.py:600] Compiled train_step cache size: 0

Training:   0%|          | 0/5 [00:00<?, ?step/s]I0420 06:13:36.791257 136883510155008 grain_pool.py:367] Grain pool will use 1 processes.
I0420 06:13:36.817793 136883510155008 grain_pool.py:440] Grain pool will start child processes.
I0420 06:13:36.822834 136883510155008 grain_pool.py:448] Grain pool started all child processes.
2026-04-20 06:13:42.870008: 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)
I0420 06:13:46.459178 137022069569344 utils.py:86] Train loop finished in: 9.6692 seconds
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/deps/src/maxtext/trainers/post_train/distillation/train_distill.py", line 765, in <module>
    app.run(main)
  File "/usr/local/lib/python3.12/site-packages/absl/app.py", line 316, in run
    _run_main(main, args)
  File "/usr/local/lib/python3.12/site-packages/absl/app.py", line 261, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "/deps/src/maxtext/trainers/post_train/distillation/train_distill.py", line 761, in main
    train_distill(student_config, teacher_config, is_offline, global_config.offline_data_dir)
  File "/deps/src/maxtext/trainers/post_train/distillation/train_distill.py", line 693, in train_distill
    trainer.train(train_iter, eval_iter)
  File "/usr/local/lib/python3.12/site-packages/tunix/sft/peft_trainer.py", line 659, in train
    train_example = sharding_utils.shard_input(
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/tunix/sft/sharding_utils.py", line 58, in shard_input
    return jax.tree.map(
           ^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/tree.py", line 155, in map
    return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/tree_util.py", line 362, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/tree_util.py", line 362, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
                             ^^^^^^
  File "/usr/local/lib/python3.12/site-packages/tunix/sft/sharding_utils.py", line 59, in <lambda>
    lambda x: jax.make_array_from_process_local_data(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/array.py", line 986, in make_array_from_process_local_data
    out = [_array_from_process_local_data(data, s, shape)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/array.py", line 1048, in _array_from_process_local_data
    return make_array_from_callback(global_shape, sharding, cb)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/array.py", line 845, in make_array_from_callback
    per_device_values = api.device_put(per_device_values, devices)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/api.py", line 2729, in device_put
    out_flat = dispatch._batched_device_put_impl(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/dispatch.py", line 558, in _batched_device_put_impl
    y = _device_put_impl(x, device=device, src=src, copy=cp, aval=aval)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/dispatch.py", line 545, in _device_put_impl
    return _device_put_sharding_impl(x, aval, device, copy)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jax/_src/dispatch.py", line 487, in _device_put_sharding_impl
    raise ValueError(
ValueError: device_put's first argument must be a fully addressable array, but got value with devices {TpuDevice(id=30, process_index=7, coords=(2,7,0), core_on_chip=0), TpuDevice(id=23, process_index=5, coords=(3,5,0), core_on_chip=0), TpuDevice(id=8, process_index=2, coords=(0,2,0), core_on_chip=0), TpuDevice(id=16, process_index=4, coords=(0,4,0), core_on_chip=0), TpuDevice(id=2, process_index=1, coords=(2,0,0), core_on_chip=0), TpuDevice(id=29, process_index=6, coords=(1,7,0), core_on_chip=0), TpuDevice(id=17, process_index=4, coords=(1,4,0), core_on_chip=0), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=7, process_index=1, coords=(3,1,0), core_on_chip=0), TpuDevice(id=6, process_index=1, coords=(2,1,0), core_on_chip=0), TpuDevice(id=3, process_index=1, coords=(3,0,0), core_on_chip=0), TpuDevice(id=9, process_index=2, coords=(1,2,0), core_on_chip=0), TpuDevice(id=10, process_index=3, coords=(2,2,0), core_on_chip=0), TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=22, process_index=5, coords=(2,5,0), core_on_chip=0), TpuDevice(id=27, process_index=7, coords=(3,6,0), core_on_chip=0), TpuDevice(id=28, process_index=6, coords=(0,7,0), core_on_chip=0), TpuDevice(id=24, process_index=6, coords=(0,6,0), core_on_chip=0), TpuDevice(id=14, process_index=3, coords=(2,3,0), core_on_chip=0), TpuDevice(id=26, process_index=7, coords=(2,6,0), core_on_chip=0), TpuDevice(id=13, process_index=2, coords=(1,3,0), core_on_chip=0), TpuDevice(id=11, process_index=3, coords=(3,2,0), core_on_chip=0), TpuDevice(id=20, process_index=4, coords=(0,5,0), core_on_chip=0), TpuDevice(id=31, process_index=7, coords=(3,7,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=15, process_index=3, coords=(3,3,0), core_on_chip=0), TpuDevice(id=12, process_index=2, coords=(0,3,0), core_on_chip=0), TpuDevice(id=18, process_index=5, coords=(2,4,0), core_on_chip=0), TpuDevice(id=19, process_index=5, coords=(3,4,0), core_on_chip=0), TpuDevice(id=25, process_index=6, coords=(1,6,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=21, process_index=4, coords=(1,5,0), core_on_chip=0)}
I0420 06:13:46.806343 136883510155008 grain_pool.py:542] Grain pool is exiting.
I0420 06:13:46.806442 136883510155008 grain_pool.py:547] Shutting down multiprocessing system.
I0420 06:13:48.251580 136883510155008 grain_pool.py:547] Shutting down multiprocessing system.

Training:   0%|          | 0/5 [00:13<?, ?step/s]
/usr/local/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 '
XPK End: Mon Apr 20 06:13:55 UTC 2026
EXIT_CODE=1