MaxView

XPK Start: Wed Apr 22 07:35:22 UTC 2026
2026-04-22 07:35:39.857371: 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)
Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'}
`rope_scaling`'s factor field must be a float >= 1, got 40
`rope_scaling`'s beta_fast field must be a float, got 32
`rope_scaling`'s beta_slow field must be a float, got 1
Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'}
Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'}
Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'}
I0422 07:35:43.873320 134641274300224 max_utils.py:273] Attempting to initialize the jax distributed system...
INFO:2026-04-22 07:35:52,914:jax._src.distributed:149: Starting JAX distributed service on [::]:8482
I0422 07:35:52.914499 134641274300224 distributed.py:149] Starting JAX distributed service on [::]:8482
INFO:2026-04-22 07:35:52,916:jax._src.distributed:166: Connecting to JAX distributed service on mt-07-distill-smoke-1gr1v-slice-job-0-0.mt-07-distill-smoke-1gr1v:8482
I0422 07:35:52.916772 134641274300224 distributed.py:166] Connecting to JAX distributed service on mt-07-distill-smoke-1gr1v-slice-job-0-0.mt-07-distill-smoke-1gr1v:8482
I0422 07:35:54.103542 134641274300224 max_utils.py:284] Jax distributed system initialized!
I0422 07:35:59.737997 134641274300224 max_utils.py:244] Jax distributed system is already initialized.
W0422 07:35:59.870467 134641274300224 pyconfig.py:111] base_output_directory is not provided; Using local directory called maxtext_output
I0422 07:35:59.930157 134641274300224 max_utils.py:244] Jax distributed system is already initialized.
I0422 07:35:59.931339 134641274300224 pyconfig.py:471] Config param abort_on_inf_loss: True
I0422 07:35:59.931389 134641274300224 pyconfig.py:471] Config param abort_on_nan_loss: True
I0422 07:35:59.931417 134641274300224 pyconfig.py:471] Config param act_quantization_calibration_method: absmax
I0422 07:35:59.931443 134641274300224 pyconfig.py:471] Config param activation_dropout_for_audio: 0.0
I0422 07:35:59.931463 134641274300224 pyconfig.py:471] Config param activation_function_for_audio: gelu
I0422 07:35:59.931481 134641274300224 pyconfig.py:471] Config param activations_in_float32: False
I0422 07:35:59.931498 134641274300224 pyconfig.py:471] Config param adam_b1: 0.9
I0422 07:35:59.931517 134641274300224 pyconfig.py:471] Config param adam_b2: 0.95
I0422 07:35:59.931535 134641274300224 pyconfig.py:471] Config param adam_eps: 1e-08
I0422 07:35:59.931558 134641274300224 pyconfig.py:471] Config param adam_eps_root: 0.0
I0422 07:35:59.931575 134641274300224 pyconfig.py:471] Config param adam_weight_decay: 0.1
I0422 07:35:59.931592 134641274300224 pyconfig.py:471] Config param adamw_mask: []
I0422 07:35:59.931608 134641274300224 pyconfig.py:471] Config param add_bos: True
I0422 07:35:59.931625 134641274300224 pyconfig.py:471] Config param add_eos: True
I0422 07:35:59.931642 134641274300224 pyconfig.py:471] Config param allow_split_physical_axes: False
I0422 07:35:59.931658 134641274300224 pyconfig.py:471] Config param ar_cache_axis_order: 1,2,0,3
I0422 07:35:59.931674 134641274300224 pyconfig.py:471] Config param async_checkpointing: True
I0422 07:35:59.931690 134641274300224 pyconfig.py:471] Config param async_scheduling: False
I0422 07:35:59.931706 134641274300224 pyconfig.py:471] Config param attention: dot_product
I0422 07:35:59.931734 134641274300224 pyconfig.py:471] Config param attention_bias: False
I0422 07:35:59.931750 134641274300224 pyconfig.py:471] Config param attention_dropout_for_audio: 0.0
I0422 07:35:59.931768 134641274300224 pyconfig.py:471] Config param attention_out: RematLocation.REMAT
I0422 07:35:59.931788 134641274300224 pyconfig.py:471] Config param attention_output_dim: -1
I0422 07:35:59.931805 134641274300224 pyconfig.py:471] Config param attention_sink: False
I0422 07:35:59.931823 134641274300224 pyconfig.py:471] Config param attention_type: global
I0422 07:35:59.931839 134641274300224 pyconfig.py:471] Config param attn_logits_soft_cap: None
I0422 07:35:59.931855 134641274300224 pyconfig.py:471] Config param audio_path: 
I0422 07:35:59.931871 134641274300224 pyconfig.py:471] Config param audio_placeholder: <|audio|>
I0422 07:35:59.931887 134641274300224 pyconfig.py:471] Config param autoregressive_decode_assert: 
I0422 07:35:59.931902 134641274300224 pyconfig.py:471] Config param base_config: base.yml
I0422 07:35:59.931918 134641274300224 pyconfig.py:471] Config param base_emb_dim: 16
I0422 07:35:59.931934 134641274300224 pyconfig.py:471] Config param base_mlp_dim: 64
I0422 07:35:59.931950 134641274300224 pyconfig.py:471] Config param base_moe_mlp_dim: -1
I0422 07:35:59.931967 134641274300224 pyconfig.py:471] Config param base_num_decoder_layers: 1
I0422 07:35:59.931983 134641274300224 pyconfig.py:471] Config param base_num_kv_heads: 2
I0422 07:35:59.931999 134641274300224 pyconfig.py:471] Config param base_num_query_heads: 2
I0422 07:35:59.932018 134641274300224 pyconfig.py:471] Config param base_output_directory: /deps/maxtext_output
I0422 07:35:59.932033 134641274300224 pyconfig.py:471] Config param batch_size: 1
I0422 07:35:59.932050 134641274300224 pyconfig.py:471] Config param batch_split_factor: 1
I0422 07:35:59.932065 134641274300224 pyconfig.py:471] Config param beta_fast: 32
I0422 07:35:59.932081 134641274300224 pyconfig.py:471] Config param beta_slow: 1
I0422 07:35:59.932097 134641274300224 pyconfig.py:471] Config param bwd_quantization_calibration_method: absmax
I0422 07:35:59.932115 134641274300224 pyconfig.py:471] Config param capacity_factor: -1.0
I0422 07:35:59.932131 134641274300224 pyconfig.py:471] Config param cast_logits_to_fp32: True
I0422 07:35:59.932148 134641274300224 pyconfig.py:471] Config param chat_template: 
I0422 07:35:59.932164 134641274300224 pyconfig.py:471] Config param chat_template_path: 
I0422 07:35:59.932182 134641274300224 pyconfig.py:471] Config param checkpoint_conversion_fn: None
I0422 07:35:59.932199 134641274300224 pyconfig.py:471] Config param checkpoint_dir: /deps/maxtext_output/gpt3-52k_2026-04-22-07-35/checkpoints/
I0422 07:35:59.932216 134641274300224 pyconfig.py:471] Config param checkpoint_is_quantized: False
I0422 07:35:59.932232 134641274300224 pyconfig.py:471] Config param checkpoint_period: 2000
I0422 07:35:59.932249 134641274300224 pyconfig.py:471] Config param checkpoint_storage_concurrent_gb: 96
I0422 07:35:59.932265 134641274300224 pyconfig.py:471] Config param checkpoint_storage_target_data_file_size_bytes: 2147483648
I0422 07:35:59.932281 134641274300224 pyconfig.py:471] Config param checkpoint_storage_use_ocdbt: True
I0422 07:35:59.932296 134641274300224 pyconfig.py:471] Config param checkpoint_storage_use_zarr3: True
I0422 07:35:59.932312 134641274300224 pyconfig.py:471] Config param checkpoint_todelete_full_path: None
I0422 07:35:59.932327 134641274300224 pyconfig.py:471] Config param checkpoint_todelete_subdir: None
I0422 07:35:59.932342 134641274300224 pyconfig.py:471] Config param chips_per_vm: 4
I0422 07:35:59.932358 134641274300224 pyconfig.py:471] Config param chunk_attn_window_size: 0
I0422 07:35:59.932372 134641274300224 pyconfig.py:471] Config param collect_stack_trace: False
I0422 07:35:59.932388 134641274300224 pyconfig.py:471] Config param colocated_python_checkpointing: False
I0422 07:35:59.932404 134641274300224 pyconfig.py:471] Config param colocated_python_data_input: False
I0422 07:35:59.932419 134641274300224 pyconfig.py:471] Config param compile_topology: 
I0422 07:35:59.932440 134641274300224 pyconfig.py:471] Config param compile_topology_num_slices: -1
I0422 07:35:59.932456 134641274300224 pyconfig.py:471] Config param compile_xla_flags: 
I0422 07:35:59.932471 134641274300224 pyconfig.py:471] Config param compiled_trainstep_file: 
I0422 07:35:59.932486 134641274300224 pyconfig.py:471] Config param compute_axis_order: 0,1,2,3
I0422 07:35:59.932502 134641274300224 pyconfig.py:471] Config param constant_bound_config: []
I0422 07:35:59.932518 134641274300224 pyconfig.py:471] Config param context: RematLocation.REMAT
I0422 07:35:59.932534 134641274300224 pyconfig.py:471] Config param context_parallel_load_balance: True
I0422 07:35:59.932549 134641274300224 pyconfig.py:471] Config param context_parallel_reorder_strategy: ReorderStrategy.AUTO
I0422 07:35:59.932566 134641274300224 pyconfig.py:471] Config param context_parallel_size: 1
I0422 07:35:59.932581 134641274300224 pyconfig.py:471] Config param context_parallel_strategy: all_gather
I0422 07:35:59.932597 134641274300224 pyconfig.py:471] Config param context_sharding: context
I0422 07:35:59.932613 134641274300224 pyconfig.py:471] Config param conv_chunksize_for_audio: 500
I0422 07:35:59.932627 134641274300224 pyconfig.py:471] Config param conv_stride_for_vit: 14
I0422 07:35:59.932651 134641274300224 pyconfig.py:471] Config param convert_checkpoint_if_possible: False
I0422 07:35:59.932667 134641274300224 pyconfig.py:471] Config param cost_estimate_flops_bwd: -1
I0422 07:35:59.932682 134641274300224 pyconfig.py:471] Config param cost_estimate_flops_fwd: -1
I0422 07:35:59.932698 134641274300224 pyconfig.py:471] Config param custom_mesh: 
I0422 07:35:59.932736 134641274300224 pyconfig.py:471] Config param custom_mesh_and_rule: 
I0422 07:35:59.932757 134641274300224 pyconfig.py:471] Config param d_model_for_audio: 256
I0422 07:35:59.932778 134641274300224 pyconfig.py:471] Config param data_sharding: (('data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive'),)
I0422 07:35:59.932809 134641274300224 pyconfig.py:471] Config param data_shuffle_seed: 0
I0422 07:35:59.932835 134641274300224 pyconfig.py:471] Config param dataset_name: c4/en:3.0.1
I0422 07:35:59.932850 134641274300224 pyconfig.py:471] Config param dataset_path: 
I0422 07:35:59.932866 134641274300224 pyconfig.py:471] Config param dataset_type: DatasetType.HF
I0422 07:35:59.932882 134641274300224 pyconfig.py:471] Config param dcn_autoregressive_parallelism: 1
I0422 07:35:59.932898 134641274300224 pyconfig.py:471] Config param dcn_context_autoregressive_parallelism: 1
I0422 07:35:59.932913 134641274300224 pyconfig.py:471] Config param dcn_context_parallelism: 1
I0422 07:35:59.932928 134641274300224 pyconfig.py:471] Config param dcn_data_parallelism: -1
I0422 07:35:59.932943 134641274300224 pyconfig.py:471] Config param dcn_diloco_parallelism: 1
I0422 07:35:59.932958 134641274300224 pyconfig.py:471] Config param dcn_expert_parallelism: 1
I0422 07:35:59.932974 134641274300224 pyconfig.py:471] Config param dcn_fsdp_parallelism: 1
I0422 07:35:59.932988 134641274300224 pyconfig.py:471] Config param dcn_fsdp_transpose_parallelism: 1
I0422 07:35:59.933004 134641274300224 pyconfig.py:471] Config param dcn_parallelism: [1, -1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
I0422 07:35:59.933021 134641274300224 pyconfig.py:471] Config param dcn_pipeline_parallelism: 1
I0422 07:35:59.933036 134641274300224 pyconfig.py:471] Config param dcn_sequence_parallelism: 1
I0422 07:35:59.933052 134641274300224 pyconfig.py:471] Config param dcn_tensor_parallelism: 1
I0422 07:35:59.933066 134641274300224 pyconfig.py:471] Config param dcn_tensor_sequence_parallelism: 1
I0422 07:35:59.933083 134641274300224 pyconfig.py:471] Config param dcn_tensor_transpose_parallelism: 1
I0422 07:35:59.933099 134641274300224 pyconfig.py:471] Config param debug: {'rl': False}
I0422 07:35:59.933115 134641274300224 pyconfig.py:471] Config param debug_sharding: False
I0422 07:35:59.933130 134641274300224 pyconfig.py:471] Config param decode_sampling_nucleus_p: -1
I0422 07:35:59.933146 134641274300224 pyconfig.py:471] Config param decode_sampling_strategy: SamplingStrategy.GREEDY
I0422 07:35:59.933165 134641274300224 pyconfig.py:471] Config param decode_sampling_temperature: 1.0
I0422 07:35:59.933181 134641274300224 pyconfig.py:471] Config param decode_sampling_top_k: 0
I0422 07:35:59.933197 134641274300224 pyconfig.py:471] Config param decoder_block: DecoderBlockType.GPT3
I0422 07:35:59.933216 134641274300224 pyconfig.py:471] Config param decoder_layer_input: RematLocation.DEVICE
I0422 07:35:59.933233 134641274300224 pyconfig.py:471] Config param deepstack_visual_indexes_for_vit: []
I0422 07:35:59.933248 134641274300224 pyconfig.py:471] Config param degenerate_group_masking: True
I0422 07:35:59.933263 134641274300224 pyconfig.py:471] Config param dense_init_scale: 1.0
I0422 07:35:59.933279 134641274300224 pyconfig.py:471] Config param diloco_outer_lr: 0.3
I0422 07:35:59.933295 134641274300224 pyconfig.py:471] Config param diloco_outer_momentum: 0.9
I0422 07:35:59.933310 134641274300224 pyconfig.py:471] Config param diloco_sync_period: 36
I0422 07:35:59.933327 134641274300224 pyconfig.py:471] Config param distill_alpha: 0.5
I0422 07:35:59.933343 134641274300224 pyconfig.py:471] Config param distill_alpha_end: None
I0422 07:35:59.933358 134641274300224 pyconfig.py:471] Config param distill_alpha_schedule: constant
I0422 07:35:59.933373 134641274300224 pyconfig.py:471] Config param distill_beta: 0.0
I0422 07:35:59.933388 134641274300224 pyconfig.py:471] Config param distill_beta_end: None
I0422 07:35:59.933403 134641274300224 pyconfig.py:471] Config param distill_beta_schedule: constant
I0422 07:35:59.933418 134641274300224 pyconfig.py:471] Config param distill_feature_loss_type: cosine
I0422 07:35:59.933434 134641274300224 pyconfig.py:471] Config param distill_layer_indices: None
I0422 07:35:59.933453 134641274300224 pyconfig.py:471] Config param distill_temperature: 1.0
I0422 07:35:59.933469 134641274300224 pyconfig.py:471] Config param distill_temperature_end: None
I0422 07:35:59.933483 134641274300224 pyconfig.py:471] Config param distill_temperature_schedule: constant
I0422 07:35:59.933499 134641274300224 pyconfig.py:471] Config param downsample_hidden_size_for_audio: 256
I0422 07:35:59.933513 134641274300224 pyconfig.py:471] Config param dpo_beta: 0.1
I0422 07:35:59.933529 134641274300224 pyconfig.py:471] Config param dpo_label_smoothing: 0.0
I0422 07:35:59.933544 134641274300224 pyconfig.py:471] Config param dq_reduction_steps: 0
I0422 07:35:59.933560 134641274300224 pyconfig.py:471] Config param dropout_rate: 0.0
I0422 07:35:59.933575 134641274300224 pyconfig.py:471] Config param dtype: bfloat16
I0422 07:35:59.933606 134641274300224 pyconfig.py:471] Config param dtype_mm: float32
I0422 07:35:59.933622 134641274300224 pyconfig.py:471] Config param dump_hlo: False
I0422 07:35:59.933638 134641274300224 pyconfig.py:471] Config param dump_hlo_delete_local_after: True
I0422 07:35:59.933652 134641274300224 pyconfig.py:471] Config param dump_hlo_gcs_dir: /deps/maxtext_output/gpt3-52k_2026-04-22-07-35/xla_dump
I0422 07:35:59.933668 134641274300224 pyconfig.py:471] Config param dump_hlo_local_dir: /tmp/xla_dump/
I0422 07:35:59.933682 134641274300224 pyconfig.py:471] Config param dump_hlo_local_module_name: jit_train_step
I0422 07:35:59.933698 134641274300224 pyconfig.py:471] Config param dump_hlo_module_name: jit_train_step
I0422 07:35:59.933727 134641274300224 pyconfig.py:471] Config param dump_hlo_upload_all: False
I0422 07:35:59.933752 134641274300224 pyconfig.py:471] Config param dump_hlo_xla_flags: 
I0422 07:35:59.933777 134641274300224 pyconfig.py:471] Config param dump_jaxpr: False
I0422 07:35:59.933793 134641274300224 pyconfig.py:471] Config param dump_jaxpr_delete_local_after: True
I0422 07:35:59.933809 134641274300224 pyconfig.py:471] Config param dump_jaxpr_gcs_dir: /deps/maxtext_output/gpt3-52k_2026-04-22-07-35/jaxpr_dump
I0422 07:35:59.933824 134641274300224 pyconfig.py:471] Config param dump_jaxpr_local_dir: /tmp/jaxpr_dump/
I0422 07:35:59.933842 134641274300224 pyconfig.py:471] Config param dump_step: -1
I0422 07:35:59.933868 134641274300224 pyconfig.py:471] Config param elastic_enabled: False
I0422 07:35:59.933893 134641274300224 pyconfig.py:471] Config param elastic_max_retries: 10
I0422 07:35:59.933908 134641274300224 pyconfig.py:471] Config param elastic_timeout_seconds: 300
I0422 07:35:59.933925 134641274300224 pyconfig.py:471] Config param emb_dim: 16
I0422 07:35:59.933951 134641274300224 pyconfig.py:471] Config param enable_autocheckpoint: False
I0422 07:35:59.933978 134641274300224 pyconfig.py:471] Config param enable_checkpoint_cloud_logger: False
I0422 07:35:59.934003 134641274300224 pyconfig.py:471] Config param enable_checkpointing: True
I0422 07:35:59.934030 134641274300224 pyconfig.py:471] Config param enable_continuous_checkpointing: False
I0422 07:35:59.934055 134641274300224 pyconfig.py:471] Config param enable_data_shuffling: True
I0422 07:35:59.934081 134641274300224 pyconfig.py:471] Config param enable_diloco: False
I0422 07:35:59.934107 134641274300224 pyconfig.py:471] Config param enable_dp_attention: False
I0422 07:35:59.934134 134641274300224 pyconfig.py:471] Config param enable_dropout: False
I0422 07:35:59.934158 134641274300224 pyconfig.py:471] Config param enable_emergency_checkpoint: False
I0422 07:35:59.934182 134641274300224 pyconfig.py:471] Config param enable_expert_parallel: False
I0422 07:35:59.934207 134641274300224 pyconfig.py:471] Config param enable_gcp_goodput_metrics: True
I0422 07:35:59.934233 134641274300224 pyconfig.py:471] Config param enable_gcp_step_deviation_metrics: True
I0422 07:35:59.934258 134641274300224 pyconfig.py:471] Config param enable_goodput_recording: False
I0422 07:35:59.934280 134641274300224 pyconfig.py:471] Config param enable_jax_profiler: False
I0422 07:35:59.934298 134641274300224 pyconfig.py:471] Config param enable_llm_inference_pool: False
I0422 07:35:59.934322 134641274300224 pyconfig.py:471] Config param enable_model_warmup: False
I0422 07:35:59.934348 134641274300224 pyconfig.py:471] Config param enable_multi_tier_checkpointing: False
I0422 07:35:59.934372 134641274300224 pyconfig.py:471] Config param enable_nnx: False
I0422 07:35:59.934398 134641274300224 pyconfig.py:471] Config param enable_orbax_v1: False
I0422 07:35:59.934423 134641274300224 pyconfig.py:471] Config param enable_padding_causal_mask: True
I0422 07:35:59.934456 134641274300224 pyconfig.py:471] Config param enable_pathways_goodput: False
I0422 07:35:59.934483 134641274300224 pyconfig.py:471] Config param enable_prefix_caching: False
I0422 07:35:59.934508 134641274300224 pyconfig.py:471] Config param enable_rampup_batch_size: False
I0422 07:35:59.934532 134641274300224 pyconfig.py:471] Config param enable_single_controller: False
I0422 07:35:59.934554 134641274300224 pyconfig.py:471] Config param enable_single_replica_ckpt_restoring: False
I0422 07:35:59.934578 134641274300224 pyconfig.py:471] Config param enable_tensorboard: True
I0422 07:35:59.934603 134641274300224 pyconfig.py:471] Config param enable_tunix_perf_metrics: False
I0422 07:35:59.934630 134641274300224 pyconfig.py:471] Config param encoder_attention_heads_for_audio: 4
I0422 07:35:59.934654 134641274300224 pyconfig.py:471] Config param encoder_ffn_dim_for_audio: 512
I0422 07:35:59.934677 134641274300224 pyconfig.py:471] Config param encoder_layers_for_audio: 2
I0422 07:35:59.934699 134641274300224 pyconfig.py:471] Config param engram: RematLocation.REMAT
I0422 07:35:59.934738 134641274300224 pyconfig.py:471] Config param engram_head_dim: 1280
I0422 07:35:59.934766 134641274300224 pyconfig.py:471] Config param engram_kernel_size: 4
I0422 07:35:59.934792 134641274300224 pyconfig.py:471] Config param engram_layers: []
I0422 07:35:59.934815 134641274300224 pyconfig.py:471] Config param engram_max_ngram_size: 3
I0422 07:35:59.934842 134641274300224 pyconfig.py:471] Config param engram_num_heads: 8
I0422 07:35:59.934867 134641274300224 pyconfig.py:471] Config param engram_seed: 0
I0422 07:35:59.934890 134641274300224 pyconfig.py:471] Config param engram_vocab_bases: []
I0422 07:35:59.934919 134641274300224 pyconfig.py:471] Config param epsilon_high: None
I0422 07:35:59.934945 134641274300224 pyconfig.py:471] Config param eval_corr_lst: False
I0422 07:35:59.934968 134641274300224 pyconfig.py:471] Config param eval_data_columns: ['text']
I0422 07:35:59.934991 134641274300224 pyconfig.py:471] Config param eval_dataset_name: c4/en:3.0.1
I0422 07:35:59.935013 134641274300224 pyconfig.py:471] Config param eval_image_column: image
I0422 07:35:59.935039 134641274300224 pyconfig.py:471] Config param eval_interval: -1
I0422 07:35:59.935064 134641274300224 pyconfig.py:471] Config param eval_make_lst: False
I0422 07:35:59.935088 134641274300224 pyconfig.py:471] Config param eval_per_device_batch_size: 2
I0422 07:35:59.935114 134641274300224 pyconfig.py:471] Config param eval_sampling_strategy: greedy
I0422 07:35:59.935139 134641274300224 pyconfig.py:471] Config param eval_split: validation
I0422 07:35:59.935164 134641274300224 pyconfig.py:471] Config param eval_steps: -1
I0422 07:35:59.935186 134641274300224 pyconfig.py:471] Config param expansion_factor_real_data: -1.0
I0422 07:35:59.935210 134641274300224 pyconfig.py:471] Config param final_logits_soft_cap: None
I0422 07:35:59.935227 134641274300224 pyconfig.py:471] Config param first_num_dense_layers: 0
I0422 07:35:59.935243 134641274300224 pyconfig.py:471] Config param float32_gate_logits: False
I0422 07:35:59.935257 134641274300224 pyconfig.py:471] Config param float32_logits: False
I0422 07:35:59.935272 134641274300224 pyconfig.py:471] Config param float32_qk_product: False
I0422 07:35:59.935287 134641274300224 pyconfig.py:471] Config param float32_weight_sum: True
I0422 07:35:59.935302 134641274300224 pyconfig.py:471] Config param force_q_layout: False
I0422 07:35:59.935317 134641274300224 pyconfig.py:471] Config param force_unroll: False
I0422 07:35:59.935333 134641274300224 pyconfig.py:471] Config param freeze_audio_encoder_params: True
I0422 07:35:59.935348 134641274300224 pyconfig.py:471] Config param freeze_vision_encoder_params: True
I0422 07:35:59.935364 134641274300224 pyconfig.py:471] Config param fused_mlp: False
I0422 07:35:59.935378 134641274300224 pyconfig.py:471] Config param fused_qkv: True
I0422 07:35:59.935394 134641274300224 pyconfig.py:471] Config param gcs_metrics: False
I0422 07:35:59.935410 134641274300224 pyconfig.py:471] Config param gdn_chunk_size: 64
I0422 07:35:59.935424 134641274300224 pyconfig.py:471] Config param gdn_conv_kernel_dim: 4
I0422 07:35:59.935443 134641274300224 pyconfig.py:471] Config param gdn_key_head_dim: 128
I0422 07:35:59.935459 134641274300224 pyconfig.py:471] Config param gdn_num_key_heads: 16
I0422 07:35:59.935475 134641274300224 pyconfig.py:471] Config param gdn_num_value_heads: 32
I0422 07:35:59.935489 134641274300224 pyconfig.py:471] Config param gdn_value_head_dim: 128
I0422 07:35:59.935505 134641274300224 pyconfig.py:471] Config param generate_padding_batch_eval: False
I0422 07:35:59.935519 134641274300224 pyconfig.py:471] Config param generate_padding_batch_train: False
I0422 07:35:59.935534 134641274300224 pyconfig.py:471] Config param generate_slice: v5e-16
I0422 07:35:59.935549 134641274300224 pyconfig.py:471] Config param generation_configs: {}
I0422 07:35:59.935565 134641274300224 pyconfig.py:471] Config param global_batch_size_to_eval_on: 64
I0422 07:35:59.935580 134641274300224 pyconfig.py:471] Config param global_batch_size_to_load: 512
I0422 07:35:59.935595 134641274300224 pyconfig.py:471] Config param global_batch_size_to_load_eval: 64
I0422 07:35:59.935611 134641274300224 pyconfig.py:471] Config param global_batch_size_to_load_increment: None
I0422 07:35:59.935625 134641274300224 pyconfig.py:471] Config param global_batch_size_to_load_start: None
I0422 07:35:59.935642 134641274300224 pyconfig.py:471] Config param global_batch_size_to_train_on: 512
I0422 07:35:59.935657 134641274300224 pyconfig.py:471] Config param global_head_dim: 0
I0422 07:35:59.935672 134641274300224 pyconfig.py:471] Config param global_num_kv_heads: 0
I0422 07:35:59.935687 134641274300224 pyconfig.py:471] Config param global_parameter_scale: 1
I0422 07:35:59.935703 134641274300224 pyconfig.py:471] Config param global_rampup_samples: 500
I0422 07:35:59.935730 134641274300224 pyconfig.py:471] Config param global_rope_max_timescale: -1
I0422 07:35:59.935744 134641274300224 pyconfig.py:471] Config param global_rope_proportion: 0.25
I0422 07:35:59.935761 134641274300224 pyconfig.py:471] Config param goodput_upload_interval_seconds: 30
I0422 07:35:59.935776 134641274300224 pyconfig.py:471] Config param grad_dtype: float32
I0422 07:35:59.935815 134641274300224 pyconfig.py:471] Config param gradient_accumulation_steps: 8
I0422 07:35:59.935831 134641274300224 pyconfig.py:471] Config param gradient_clipping_threshold: 1.0
I0422 07:35:59.935847 134641274300224 pyconfig.py:471] Config param grain_data_source_max_workers: 16
I0422 07:35:59.935863 134641274300224 pyconfig.py:471] Config param grain_eval_files: 
I0422 07:35:59.935878 134641274300224 pyconfig.py:471] Config param grain_file_type: arrayrecord
I0422 07:35:59.935894 134641274300224 pyconfig.py:471] Config param grain_num_threads: 16
I0422 07:35:59.935909 134641274300224 pyconfig.py:471] Config param grain_num_threads_eval: 16
I0422 07:35:59.935925 134641274300224 pyconfig.py:471] Config param grain_packing_type: first_fit
I0422 07:35:59.935941 134641274300224 pyconfig.py:471] Config param grain_per_worker_buffer_size: 1
I0422 07:35:59.935956 134641274300224 pyconfig.py:471] Config param grain_per_worker_buffer_size_eval: 1
I0422 07:35:59.935972 134641274300224 pyconfig.py:471] Config param grain_prefetch_buffer_size: 500
I0422 07:35:59.935988 134641274300224 pyconfig.py:471] Config param grain_prefetch_buffer_size_eval: 500
I0422 07:35:59.936004 134641274300224 pyconfig.py:471] Config param grain_ram_budget_mb: 1024
I0422 07:35:59.936018 134641274300224 pyconfig.py:471] Config param grain_shuffle_buffer_size: 100
I0422 07:35:59.936033 134641274300224 pyconfig.py:471] Config param grain_train_files: 
I0422 07:35:59.936048 134641274300224 pyconfig.py:471] Config param grain_train_mixture_config_path: 
I0422 07:35:59.936064 134641274300224 pyconfig.py:471] Config param grain_worker_count: 1
I0422 07:35:59.936080 134641274300224 pyconfig.py:471] Config param grain_worker_count_eval: 1
I0422 07:35:59.936094 134641274300224 pyconfig.py:471] Config param grpo_beta: 0.08
I0422 07:35:59.936110 134641274300224 pyconfig.py:471] Config param grpo_epsilon: 0.2
I0422 07:35:59.936126 134641274300224 pyconfig.py:471] Config param hardware: tpu
I0422 07:35:59.936141 134641274300224 pyconfig.py:471] Config param hbm_utilization_vllm: 0.72
I0422 07:35:59.936157 134641274300224 pyconfig.py:471] Config param head_dim: 8
I0422 07:35:59.936182 134641274300224 pyconfig.py:471] Config param heartbeat_reporting_interval_in_seconds: 5
I0422 07:35:59.936207 134641274300224 pyconfig.py:471] Config param hf_data_dir: None
I0422 07:35:59.936234 134641274300224 pyconfig.py:471] Config param hf_eval_files: None
I0422 07:35:59.936261 134641274300224 pyconfig.py:471] Config param hf_eval_split: None
I0422 07:35:59.936286 134641274300224 pyconfig.py:471] Config param hf_name: None
I0422 07:35:59.936311 134641274300224 pyconfig.py:471] Config param hf_path: OptimalScale/ClimbMix
I0422 07:35:59.936335 134641274300224 pyconfig.py:471] Config param hf_train_files: None
I0422 07:35:59.936361 134641274300224 pyconfig.py:471] Config param hidden_size_for_vit: 1408
I0422 07:35:59.936387 134641274300224 pyconfig.py:471] Config param hide_profiler_step_metric: False
I0422 07:35:59.936411 134641274300224 pyconfig.py:471] Config param ici_autoregressive_parallelism: 1
I0422 07:35:59.936439 134641274300224 pyconfig.py:471] Config param ici_context_autoregressive_parallelism: 1
I0422 07:35:59.936460 134641274300224 pyconfig.py:471] Config param ici_context_parallelism: 1
I0422 07:35:59.936477 134641274300224 pyconfig.py:471] Config param ici_data_parallelism: 1
I0422 07:35:59.936500 134641274300224 pyconfig.py:471] Config param ici_diloco_parallelism: 1
I0422 07:35:59.936526 134641274300224 pyconfig.py:471] Config param ici_expert_parallelism: 1
I0422 07:35:59.936551 134641274300224 pyconfig.py:471] Config param ici_fsdp_parallelism: -1
I0422 07:35:59.936575 134641274300224 pyconfig.py:471] Config param ici_fsdp_transpose_parallelism: 1
I0422 07:35:59.936599 134641274300224 pyconfig.py:471] Config param ici_parallelism: [1, 1, 1, -1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
I0422 07:35:59.936626 134641274300224 pyconfig.py:471] Config param ici_pipeline_parallelism: 1
I0422 07:35:59.936651 134641274300224 pyconfig.py:471] Config param ici_sequence_parallelism: 1
I0422 07:35:59.936675 134641274300224 pyconfig.py:471] Config param ici_tensor_parallelism: 1
I0422 07:35:59.936694 134641274300224 pyconfig.py:471] Config param ici_tensor_sequence_parallelism: 1
I0422 07:35:59.936724 134641274300224 pyconfig.py:471] Config param ici_tensor_transpose_parallelism: 1
I0422 07:35:59.936752 134641274300224 pyconfig.py:471] Config param image_path: 
I0422 07:35:59.936779 134641274300224 pyconfig.py:471] Config param image_placeholder: <|image|>
I0422 07:35:59.936802 134641274300224 pyconfig.py:471] Config param image_size_for_vit: 896
I0422 07:35:59.936824 134641274300224 pyconfig.py:471] Config param indexer_head_dim: 128
I0422 07:35:59.936847 134641274300224 pyconfig.py:471] Config param indexer_loss_scaling_factor: 0.0
I0422 07:35:59.936874 134641274300224 pyconfig.py:471] Config param indexer_n_heads: 64
I0422 07:35:59.936897 134641274300224 pyconfig.py:471] Config param indexer_sparse_training: False
I0422 07:35:59.936923 134641274300224 pyconfig.py:471] Config param indexer_topk: 2048
I0422 07:35:59.936949 134641274300224 pyconfig.py:471] Config param inference_benchmark_test: False
I0422 07:35:59.936974 134641274300224 pyconfig.py:471] Config param inference_metadata_file: 
I0422 07:35:59.937000 134641274300224 pyconfig.py:471] Config param inference_microbenchmark_log_file_path: 
I0422 07:35:59.937025 134641274300224 pyconfig.py:471] Config param inference_microbenchmark_loop_iters: 10
I0422 07:35:59.937049 134641274300224 pyconfig.py:471] Config param inference_microbenchmark_num_samples: [1, 2, 3, 4, 5]
I0422 07:35:59.937075 134641274300224 pyconfig.py:471] Config param inference_microbenchmark_prefill_lengths: 64,128,256,512,1024
I0422 07:35:59.937097 134641274300224 pyconfig.py:471] Config param inference_microbenchmark_stages: prefill,generate
I0422 07:35:59.937119 134641274300224 pyconfig.py:471] Config param inference_server: MaxtextInterleavedServer
I0422 07:35:59.937134 134641274300224 pyconfig.py:471] Config param inhomogeneous_layer_cycle_interval: 1
I0422 07:35:59.937149 134641274300224 pyconfig.py:471] Config param init_weights_seed: 0
I0422 07:35:59.937165 134641274300224 pyconfig.py:471] Config param input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activation_norm_length']
I0422 07:35:59.937180 134641274300224 pyconfig.py:471] Config param interleave_moe_layer_step: 1
I0422 07:35:59.937198 134641274300224 pyconfig.py:471] Config param intermediate_size_for_vit: 5632
I0422 07:35:59.937213 134641274300224 pyconfig.py:471] Config param internal_compile: False
I0422 07:35:59.937230 134641274300224 pyconfig.py:471] Config param internal_compile_num_devices: -1
I0422 07:35:59.937245 134641274300224 pyconfig.py:471] Config param jax_cache_dir: ~/jax_cache
I0422 07:35:59.937261 134641274300224 pyconfig.py:471] Config param jax_debug_log_modules: 
I0422 07:35:59.937276 134641274300224 pyconfig.py:471] Config param jax_distributed_initialization_timeout: 300
I0422 07:35:59.937292 134641274300224 pyconfig.py:471] Config param jax_profiler_port: 9999
I0422 07:35:59.937308 134641274300224 pyconfig.py:471] Config param key_proj: RematLocation.REMAT
I0422 07:35:59.937325 134641274300224 pyconfig.py:471] Config param kv_cache_buffer: 256
I0422 07:35:59.937340 134641274300224 pyconfig.py:471] Config param kv_lora_rank: 512
I0422 07:35:59.937355 134641274300224 pyconfig.py:471] Config param kv_quant_axis: KvQuantAxis.HEADS_AND_DKV
I0422 07:35:59.937372 134641274300224 pyconfig.py:471] Config param kv_quant_dtype: int8
I0422 07:35:59.937388 134641274300224 pyconfig.py:471] Config param kv_wa_proj: RematLocation.REMAT
I0422 07:35:59.937403 134641274300224 pyconfig.py:471] Config param learning_rate: 0.0002
I0422 07:35:59.937420 134641274300224 pyconfig.py:471] Config param learning_rate_final_fraction: 0.1
I0422 07:35:59.937441 134641274300224 pyconfig.py:471] Config param learning_rate_schedule_steps: 200000
I0422 07:35:59.937457 134641274300224 pyconfig.py:471] Config param load_balance_loss_weight: 0.0
I0422 07:35:59.937471 134641274300224 pyconfig.py:471] Config param load_checkpoint_only_once: False
I0422 07:35:59.937487 134641274300224 pyconfig.py:471] Config param load_from_prefill_dir: False
I0422 07:35:59.937502 134641274300224 pyconfig.py:471] Config param load_full_state_path: 
I0422 07:35:59.937518 134641274300224 pyconfig.py:471] Config param load_parameters_path: gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items
I0422 07:35:59.937534 134641274300224 pyconfig.py:471] Config param local_checkpoint_directory: 
I0422 07:35:59.937550 134641274300224 pyconfig.py:471] Config param local_checkpoint_period: 0
I0422 07:35:59.937564 134641274300224 pyconfig.py:471] Config param local_rope_max_timescale: -1
I0422 07:35:59.937580 134641274300224 pyconfig.py:471] Config param local_rope_proportion: 1.0
I0422 07:35:59.937596 134641274300224 pyconfig.py:471] Config param log_config: True
I0422 07:35:59.937611 134641274300224 pyconfig.py:471] Config param log_period: 10
I0422 07:35:59.937627 134641274300224 pyconfig.py:471] 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'))
I0422 07:35:59.937705 134641274300224 pyconfig.py:471] Config param logits_dot_in_fp32: False
I0422 07:35:59.937741 134641274300224 pyconfig.py:471] Config param logits_via_embedding: True
I0422 07:35:59.937755 134641274300224 pyconfig.py:471] Config param lora_input_adapters_path: 
I0422 07:35:59.937771 134641274300224 pyconfig.py:471] Config param loss_algo: grpo
I0422 07:35:59.937787 134641274300224 pyconfig.py:471] Config param lr_schedule_type: LearningRateScheduleType.COSINE
I0422 07:35:59.937806 134641274300224 pyconfig.py:471] Config param managed_mldiagnostics: False
I0422 07:35:59.937821 134641274300224 pyconfig.py:471] Config param managed_mldiagnostics_dir: /deps/maxtext_output/gpt3-52k_2026-04-22-07-35/managed-mldiagnostics
I0422 07:35:59.937837 134641274300224 pyconfig.py:471] Config param managed_mldiagnostics_run_group: 
I0422 07:35:59.937853 134641274300224 pyconfig.py:471] Config param matmul_precision: MatmulPrecision.DEFAULT
I0422 07:35:59.937871 134641274300224 pyconfig.py:471] Config param max_checkify: False
I0422 07:35:59.937887 134641274300224 pyconfig.py:471] Config param max_concurrency: 256
I0422 07:35:59.937901 134641274300224 pyconfig.py:471] Config param max_corpus_chars: 10000000
I0422 07:35:59.937918 134641274300224 pyconfig.py:471] Config param max_num_batched_tokens: None
I0422 07:35:59.937933 134641274300224 pyconfig.py:471] Config param max_num_checkpoints_to_keep: None
I0422 07:35:59.937948 134641274300224 pyconfig.py:471] Config param max_num_images_per_example: -1
I0422 07:35:59.937963 134641274300224 pyconfig.py:471] Config param max_num_seqs: None
I0422 07:35:59.937979 134641274300224 pyconfig.py:471] Config param max_position_embeddings: 163840
I0422 07:35:59.937994 134641274300224 pyconfig.py:471] Config param max_prefill_predict_length: 64
I0422 07:35:59.938009 134641274300224 pyconfig.py:471] Config param max_sample_len_for_audio: 10000
I0422 07:35:59.938025 134641274300224 pyconfig.py:471] Config param max_segments_per_seq: -1
I0422 07:35:59.938048 134641274300224 pyconfig.py:471] Config param max_source_positions_for_audio: 1500
I0422 07:35:59.938074 134641274300224 pyconfig.py:471] Config param max_target_length: 2048
I0422 07:35:59.938100 134641274300224 pyconfig.py:471] Config param max_timescale_for_audio: 10000.0
I0422 07:35:59.938127 134641274300224 pyconfig.py:471] Config param megablox: True
I0422 07:35:59.938146 134641274300224 pyconfig.py:471] Config param merge_gating_gmm: False
I0422 07:35:59.938161 134641274300224 pyconfig.py:471] Config param mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']
I0422 07:35:59.938179 134641274300224 pyconfig.py:471] Config param metrics_dir: /deps/maxtext_output/gpt3-52k_2026-04-22-07-35/metrics/
I0422 07:35:59.938193 134641274300224 pyconfig.py:471] Config param metrics_file: 
I0422 07:35:59.938208 134641274300224 pyconfig.py:471] Config param mhc_expansion_rate: 1
I0422 07:35:59.938222 134641274300224 pyconfig.py:471] Config param micro_batch_size_to_eval_on: 64
I0422 07:35:59.938237 134641274300224 pyconfig.py:471] Config param micro_batch_size_to_train_on: 64
I0422 07:35:59.938251 134641274300224 pyconfig.py:471] Config param mla_kv: RematLocation.REMAT
I0422 07:35:59.938268 134641274300224 pyconfig.py:471] Config param mla_naive_kvcache: True
I0422 07:35:59.938282 134641274300224 pyconfig.py:471] Config param mla_q: RematLocation.REMAT
I0422 07:35:59.938298 134641274300224 pyconfig.py:471] Config param mlp_activations: ['gelu']
I0422 07:35:59.938313 134641274300224 pyconfig.py:471] Config param mlp_activations_limit: -1.0
I0422 07:35:59.938328 134641274300224 pyconfig.py:471] Config param mlp_bias: False
I0422 07:35:59.938343 134641274300224 pyconfig.py:471] Config param mlp_dim: 64
I0422 07:35:59.938359 134641274300224 pyconfig.py:471] Config param mlpwi: RematLocation.REMAT
I0422 07:35:59.938373 134641274300224 pyconfig.py:471] Config param mlpwi_0: RematLocation.REMAT
I0422 07:35:59.938389 134641274300224 pyconfig.py:471] Config param mlpwi_1: RematLocation.REMAT
I0422 07:35:59.938406 134641274300224 pyconfig.py:471] Config param mlpwo: RematLocation.REMAT
I0422 07:35:59.938422 134641274300224 pyconfig.py:471] Config param moba: False
I0422 07:35:59.938442 134641274300224 pyconfig.py:471] Config param moba_chunk_size: 1024
I0422 07:35:59.938458 134641274300224 pyconfig.py:471] Config param moba_topk: 8
I0422 07:35:59.938472 134641274300224 pyconfig.py:471] Config param model_call_mode: 
I0422 07:35:59.938487 134641274300224 pyconfig.py:471] Config param model_name: gpt3-52k
I0422 07:35:59.938501 134641274300224 pyconfig.py:471] Config param moe_expert_input_dim: -1
I0422 07:35:59.938516 134641274300224 pyconfig.py:471] Config param moe_fsdp_use_two_stage_all_gather: False
I0422 07:35:59.938531 134641274300224 pyconfig.py:471] Config param moe_mlp_dim: -1
I0422 07:35:59.938547 134641274300224 pyconfig.py:471] Config param moe_mlpwi_0: RematLocation.REMAT
I0422 07:35:59.938561 134641274300224 pyconfig.py:471] Config param moe_mlpwi_1: RematLocation.REMAT
I0422 07:35:59.938577 134641274300224 pyconfig.py:471] Config param moe_mlpwo: RematLocation.REMAT
I0422 07:35:59.938593 134641274300224 pyconfig.py:471] Config param monitor_goodput: False
I0422 07:35:59.938607 134641274300224 pyconfig.py:471] Config param monitor_step_time_deviation: True
I0422 07:35:59.938623 134641274300224 pyconfig.py:471] Config param mrope_section: [24, 20, 20]
I0422 07:35:59.938640 134641274300224 pyconfig.py:471] Config param mscale: 1.0
I0422 07:35:59.938654 134641274300224 pyconfig.py:471] Config param mtc_data_parallelism: 0
I0422 07:35:59.938670 134641274300224 pyconfig.py:471] Config param mtp_eval_target_module: 0
I0422 07:35:59.938687 134641274300224 pyconfig.py:471] Config param mtp_loss_scaling_factor: 0.1
I0422 07:35:59.938703 134641274300224 pyconfig.py:471] Config param mtp_num_layers: 0
I0422 07:35:59.938732 134641274300224 pyconfig.py:471] Config param mu_dtype: float32
I0422 07:35:59.938759 134641274300224 pyconfig.py:471] Config param multi_sampling: False
I0422 07:35:59.938774 134641274300224 pyconfig.py:471] Config param multi_tier_checkpointing_backup_interval_minutes: 0
I0422 07:35:59.938788 134641274300224 pyconfig.py:471] Config param muon_beta: 0.95
I0422 07:35:59.938805 134641274300224 pyconfig.py:471] Config param muon_consistent_rms: None
I0422 07:35:59.938819 134641274300224 pyconfig.py:471] Config param muon_weight_decay: 0.0
I0422 07:35:59.938835 134641274300224 pyconfig.py:471] Config param n_routing_groups: -1
I0422 07:35:59.938849 134641274300224 pyconfig.py:471] Config param n_window_for_audio: 50
I0422 07:35:59.938864 134641274300224 pyconfig.py:471] Config param n_window_infer_for_audio: 800
I0422 07:35:59.938879 134641274300224 pyconfig.py:471] Config param nope_layer_interval: -1
I0422 07:35:59.938894 134641274300224 pyconfig.py:471] Config param norm_topk_prob: False
I0422 07:35:59.938910 134641274300224 pyconfig.py:471] Config param normalization_layer_epsilon: 1e-05
I0422 07:35:59.938928 134641274300224 pyconfig.py:471] Config param normalize_embedding_logits: False
I0422 07:35:59.938942 134641274300224 pyconfig.py:471] Config param num_attention_heads_for_vit: 16
I0422 07:35:59.938958 134641274300224 pyconfig.py:471] Config param num_batches: 4
I0422 07:35:59.938972 134641274300224 pyconfig.py:471] Config param num_channels_for_vit: 3
I0422 07:35:59.938988 134641274300224 pyconfig.py:471] Config param num_conv_layers_for_audio: 3
I0422 07:35:59.939002 134641274300224 pyconfig.py:471] Config param num_decoder_layers: 1
I0422 07:35:59.939018 134641274300224 pyconfig.py:471] Config param num_diloco_replicas: 1
I0422 07:35:59.939032 134641274300224 pyconfig.py:471] Config param num_epoch: 1
I0422 07:35:59.939048 134641274300224 pyconfig.py:471] Config param num_eval_passes: 1
I0422 07:35:59.939064 134641274300224 pyconfig.py:471] Config param num_experts: 1
I0422 07:35:59.939079 134641274300224 pyconfig.py:471] Config param num_experts_per_tok: 1
I0422 07:35:59.939095 134641274300224 pyconfig.py:471] Config param num_generations: 2
I0422 07:35:59.939108 134641274300224 pyconfig.py:471] Config param num_hidden_layers_for_vit: 34
I0422 07:35:59.939124 134641274300224 pyconfig.py:471] Config param num_iterations: 1
I0422 07:35:59.939140 134641274300224 pyconfig.py:471] Config param num_kv_heads: 2
I0422 07:35:59.939155 134641274300224 pyconfig.py:471] Config param num_layers_per_pipeline_stage: 1
I0422 07:35:59.939169 134641274300224 pyconfig.py:471] Config param num_mel_bins_for_audio: 128
I0422 07:35:59.939185 134641274300224 pyconfig.py:471] Config param num_pipeline_microbatches: -1
I0422 07:35:59.939199 134641274300224 pyconfig.py:471] Config param num_pipeline_repeats: -1
I0422 07:35:59.939216 134641274300224 pyconfig.py:471] Config param num_position_embeddings_for_vit: 1024
I0422 07:35:59.939232 134641274300224 pyconfig.py:471] Config param num_query_heads: 2
I0422 07:35:59.939246 134641274300224 pyconfig.py:471] Config param num_samplers_slices: -1
I0422 07:35:59.939267 134641274300224 pyconfig.py:471] Config param num_slices: 1
I0422 07:35:59.939292 134641274300224 pyconfig.py:471] Config param num_target_devices: 32
I0422 07:35:59.939318 134641274300224 pyconfig.py:471] Config param num_test_batches: 5
I0422 07:35:59.939335 134641274300224 pyconfig.py:471] Config param num_trainer_slices: -1
I0422 07:35:59.939352 134641274300224 pyconfig.py:471] Config param num_vocab_tiling: 1
I0422 07:35:59.939366 134641274300224 pyconfig.py:471] Config param off_policy_steps: 0
I0422 07:35:59.939381 134641274300224 pyconfig.py:471] Config param offline_data_dir: None
I0422 07:35:59.939395 134641274300224 pyconfig.py:471] Config param opt_type: OptimizerType.ADAM_PAX
I0422 07:35:59.939414 134641274300224 pyconfig.py:471] Config param optimize_mesh_for_tpu_v6e: False
I0422 07:35:59.939429 134641274300224 pyconfig.py:471] Config param optimizer_memory_host_offload: False
I0422 07:35:59.939449 134641274300224 pyconfig.py:471] Config param original_max_position_embeddings: 4096
I0422 07:35:59.939464 134641274300224 pyconfig.py:471] Config param out_hidden_size_for_vit: 512
I0422 07:35:59.939480 134641274300224 pyconfig.py:471] Config param out_proj: RematLocation.REMAT
I0422 07:35:59.939494 134641274300224 pyconfig.py:471] Config param output_dim_for_audio: 512
I0422 07:35:59.939510 134641274300224 pyconfig.py:471] Config param override_logical_axis_rules: False
I0422 07:35:59.939524 134641274300224 pyconfig.py:471] Config param override_model_config: True
I0422 07:35:59.939540 134641274300224 pyconfig.py:471] Config param packing: True
I0422 07:35:59.939556 134641274300224 pyconfig.py:471] Config param pagedattn_head_dim_alignment: 128
I0422 07:35:59.939570 134641274300224 pyconfig.py:471] Config param pagedattn_max_pages_per_group: -1
I0422 07:35:59.939586 134641274300224 pyconfig.py:471] Config param pagedattn_num_pages: 64
I0422 07:35:59.939600 134641274300224 pyconfig.py:471] Config param pagedattn_pages_per_compute_block: 4
I0422 07:35:59.939616 134641274300224 pyconfig.py:471] Config param pagedattn_tokens_per_page: 32
I0422 07:35:59.939630 134641274300224 pyconfig.py:471] Config param param_scan_axis: 1
I0422 07:35:59.939646 134641274300224 pyconfig.py:471] Config param parameter_memory_host_offload: False
I0422 07:35:59.939661 134641274300224 pyconfig.py:471] Config param partial_rotary_factor: 1.0
I0422 07:35:59.939676 134641274300224 pyconfig.py:471] Config param patch_size_for_vit: 14
I0422 07:35:59.939691 134641274300224 pyconfig.py:471] Config param penalty_incorrect_answer: -1.0
I0422 07:35:59.939706 134641274300224 pyconfig.py:471] Config param penalty_incorrect_format: -0.5
I0422 07:35:59.939734 134641274300224 pyconfig.py:471] Config param per_device_batch_size: 2
I0422 07:35:59.939749 134641274300224 pyconfig.py:471] Config param per_device_batch_size_increment: 2.0
I0422 07:35:59.939765 134641274300224 pyconfig.py:471] Config param per_device_batch_size_start: 4.0
I0422 07:35:59.939780 134641274300224 pyconfig.py:471] Config param pipeline_delay_activation_forwarding: False
I0422 07:35:59.939795 134641274300224 pyconfig.py:471] Config param pipeline_fsdp_ag_once: False
I0422 07:35:59.939811 134641274300224 pyconfig.py:471] Config param pipeline_fsdp_ag_per_repeat: False
I0422 07:35:59.939826 134641274300224 pyconfig.py:471] Config param pipeline_parallel_layers: 1
I0422 07:35:59.939841 134641274300224 pyconfig.py:471] Config param pixel_shuffle_ratio_for_vit: 0.5
I0422 07:35:59.939857 134641274300224 pyconfig.py:471] Config param posemb_type_for_vit: learn
I0422 07:35:59.939872 134641274300224 pyconfig.py:471] Config param position_id_per_seconds: 25
I0422 07:35:59.939887 134641274300224 pyconfig.py:471] Config param prefill_cache_axis_order: 1,2,0,3
I0422 07:35:59.939903 134641274300224 pyconfig.py:471] Config param prefill_cache_dir: 
I0422 07:35:59.939917 134641274300224 pyconfig.py:471] Config param prefill_chunk_size: 256
I0422 07:35:59.939932 134641274300224 pyconfig.py:471] Config param prefill_slice: v5e-16
I0422 07:35:59.939947 134641274300224 pyconfig.py:471] Config param prefix_caching_dram_byte: 100000000000
I0422 07:35:59.939962 134641274300224 pyconfig.py:471] Config param prefix_caching_hbm_byte: 10000000000
I0422 07:35:59.939977 134641274300224 pyconfig.py:471] Config param profile_cleanly: True
I0422 07:35:59.939993 134641274300224 pyconfig.py:471] Config param profile_periodically_period: -1
I0422 07:35:59.940007 134641274300224 pyconfig.py:471] Config param profile_power_events: False
I0422 07:35:59.940023 134641274300224 pyconfig.py:471] Config param profiler: ProfilerType.NONE
I0422 07:35:59.940041 134641274300224 pyconfig.py:471] Config param profiler_steps: 5
I0422 07:35:59.940056 134641274300224 pyconfig.py:471] Config param projector_dropout_for_vit: 0.0
I0422 07:35:59.940071 134641274300224 pyconfig.py:471] Config param projector_input_dim_for_vit: 4096
I0422 07:35:59.940086 134641274300224 pyconfig.py:471] Config param projector_output_dim_for_vit: 4096
I0422 07:35:59.940100 134641274300224 pyconfig.py:471] Config param prometheus_port: 0
I0422 07:35:59.940116 134641274300224 pyconfig.py:471] Config param prompt: I love to
I0422 07:35:59.940130 134641274300224 pyconfig.py:471] Config param pure_nnx: False
I0422 07:35:59.940146 134641274300224 pyconfig.py:471] Config param pure_nnx_decoder: False
I0422 07:35:59.940161 134641274300224 pyconfig.py:471] Config param q_lora_rank: 0
I0422 07:35:59.940177 134641274300224 pyconfig.py:471] Config param qk_clip_threshold: 100.0
I0422 07:35:59.940193 134641274300224 pyconfig.py:471] Config param qk_nope_head_dim: 128
I0422 07:35:59.940207 134641274300224 pyconfig.py:471] Config param qk_norm_with_scale: True
I0422 07:35:59.940223 134641274300224 pyconfig.py:471] Config param qk_rope_head_dim: 64
I0422 07:35:59.940238 134641274300224 pyconfig.py:471] Config param qkv_proj: RematLocation.REMAT
I0422 07:35:59.940254 134641274300224 pyconfig.py:471] Config param quant_cfg_path: 
I0422 07:35:59.940270 134641274300224 pyconfig.py:471] Config param quantization: QuantizationType.NONE
I0422 07:35:59.940289 134641274300224 pyconfig.py:471] Config param quantization_local_shard_count: 4
I0422 07:35:59.940304 134641274300224 pyconfig.py:471] Config param quantize_kvcache: False
I0422 07:35:59.940319 134641274300224 pyconfig.py:471] Config param query_proj: RematLocation.REMAT
I0422 07:35:59.940336 134641274300224 pyconfig.py:471] Config param query_wa_proj: RematLocation.REMAT
I0422 07:35:59.940353 134641274300224 pyconfig.py:471] Config param ragged_block_size: 256
I0422 07:35:59.940369 134641274300224 pyconfig.py:471] Config param ragged_buffer_factor: -1.0
I0422 07:35:59.940383 134641274300224 pyconfig.py:471] Config param rampup_end_step: 0
I0422 07:35:59.940400 134641274300224 pyconfig.py:471] Config param rampup_samples_per_increment_to_load: None
I0422 07:35:59.940414 134641274300224 pyconfig.py:471] Config param reasoning_end_token: </reasoning>
I0422 07:35:59.940430 134641274300224 pyconfig.py:471] Config param reasoning_start_token: <reasoning>
I0422 07:35:59.940450 134641274300224 pyconfig.py:471] Config param record_internal_nn_metrics: 0
I0422 07:35:59.940465 134641274300224 pyconfig.py:471] Config param remat_policy: full
I0422 07:35:59.940481 134641274300224 pyconfig.py:471] Config param remat_policy_for_vit: minimal
I0422 07:35:59.940496 134641274300224 pyconfig.py:471] Config param remove_size_one_mesh_axis_from_type: True
I0422 07:35:59.940511 134641274300224 pyconfig.py:471] Config param replicate_quant_scale: False
I0422 07:35:59.940526 134641274300224 pyconfig.py:471] Config param replicator_backup_interval_minutes: 0
I0422 07:35:59.940541 134641274300224 pyconfig.py:471] Config param report_heartbeat_metric_for_gcp_monitoring: False
I0422 07:35:59.940557 134641274300224 pyconfig.py:471] Config param report_performance_metric_for_gcp_monitoring: False
I0422 07:35:59.940571 134641274300224 pyconfig.py:471] Config param reshape_q: False
I0422 07:35:59.940587 134641274300224 pyconfig.py:471] Config param return_log_prob: False
I0422 07:35:59.940601 134641274300224 pyconfig.py:471] Config param reuse_example_batch: 0
I0422 07:35:59.940617 134641274300224 pyconfig.py:471] Config param reward_exact_answer: 5.0
I0422 07:35:59.940633 134641274300224 pyconfig.py:471] Config param reward_exact_format_match: 3.0
I0422 07:35:59.940648 134641274300224 pyconfig.py:471] Config param reward_partial_format_match: 0.5
I0422 07:35:59.940664 134641274300224 pyconfig.py:471] Config param reward_ratio_guess_to_answer_high: 0.5
I0422 07:35:59.940680 134641274300224 pyconfig.py:471] Config param reward_ratio_guess_to_answer_low: 0.25
I0422 07:35:59.940696 134641274300224 pyconfig.py:471] Config param reward_white_space_format_match: 1.5
I0422 07:35:59.940732 134641274300224 pyconfig.py:471] 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}
I0422 07:35:59.940768 134641274300224 pyconfig.py:471] Config param rollout_data_parallelism: -1
I0422 07:35:59.940794 134641274300224 pyconfig.py:471] Config param rollout_expert_parallelism: 1
I0422 07:35:59.940816 134641274300224 pyconfig.py:471] Config param rollout_micro_batch_size: -1
I0422 07:35:59.940831 134641274300224 pyconfig.py:471] Config param rollout_tensor_parallelism: -1
I0422 07:35:59.940847 134641274300224 pyconfig.py:471] Config param rope_attention_scaling: False
I0422 07:35:59.940861 134641274300224 pyconfig.py:471] Config param rope_factor: 40
I0422 07:35:59.940878 134641274300224 pyconfig.py:471] Config param rope_interleave: True
I0422 07:35:59.940894 134641274300224 pyconfig.py:471] Config param rope_linear_scaling_factor: 1.0
I0422 07:35:59.940910 134641274300224 pyconfig.py:471] Config param rope_max_timescale: 10000
I0422 07:35:59.940924 134641274300224 pyconfig.py:471] Config param rope_min_timescale: 1
I0422 07:35:59.940940 134641274300224 pyconfig.py:471] Config param rope_theta_for_vit: 10000
I0422 07:35:59.940955 134641274300224 pyconfig.py:471] Config param rope_truncate: True
I0422 07:35:59.940969 134641274300224 pyconfig.py:471] Config param rope_type: RopeType.DEFAULT
I0422 07:35:59.940987 134641274300224 pyconfig.py:471] Config param rope_use_scale: True
I0422 07:35:59.941001 134641274300224 pyconfig.py:471] Config param routed_bias: False
I0422 07:35:59.941017 134641274300224 pyconfig.py:471] Config param routed_bias_update_rate: 0.0
I0422 07:35:59.941032 134641274300224 pyconfig.py:471] Config param routed_scaling_factor: 1.0
I0422 07:35:59.941048 134641274300224 pyconfig.py:471] Config param routed_score_func: 
I0422 07:35:59.941062 134641274300224 pyconfig.py:471] Config param run_name: gpt3-52k_2026-04-22-07-35
I0422 07:35:59.941077 134641274300224 pyconfig.py:471] Config param sa_block_kv: 512
I0422 07:35:59.941093 134641274300224 pyconfig.py:471] Config param sa_block_kv_compute: 512
I0422 07:35:59.941108 134641274300224 pyconfig.py:471] Config param sa_block_kv_dkv: 512
I0422 07:35:59.941123 134641274300224 pyconfig.py:471] Config param sa_block_kv_dkv_compute: 512
I0422 07:35:59.941139 134641274300224 pyconfig.py:471] Config param sa_block_kv_dq: 512
I0422 07:35:59.941154 134641274300224 pyconfig.py:471] Config param sa_block_q: 512
I0422 07:35:59.941170 134641274300224 pyconfig.py:471] Config param sa_block_q_dkv: 512
I0422 07:35:59.941186 134641274300224 pyconfig.py:471] Config param sa_block_q_dq: 512
I0422 07:35:59.941204 134641274300224 pyconfig.py:471] Config param sa_k_layout: HEAD_DIM_MINOR
I0422 07:35:59.941218 134641274300224 pyconfig.py:471] Config param sa_q_layout: HEAD_DIM_MINOR
I0422 07:35:59.941233 134641274300224 pyconfig.py:471] Config param sa_use_fused_bwd_kernel: False
I0422 07:35:59.941249 134641274300224 pyconfig.py:471] Config param sa_v_layout: HEAD_DIM_MINOR
I0422 07:35:59.941264 134641274300224 pyconfig.py:471] Config param sampler_devices_fraction: 0.5
I0422 07:35:59.941279 134641274300224 pyconfig.py:471] Config param save_checkpoint_on_completion: True
I0422 07:35:59.941295 134641274300224 pyconfig.py:471] Config param save_config_to_gcs: False
I0422 07:35:59.941310 134641274300224 pyconfig.py:471] Config param save_quantized_params_path: 
I0422 07:35:59.941325 134641274300224 pyconfig.py:471] Config param scale_embedding_for_audio: True
I0422 07:35:59.941339 134641274300224 pyconfig.py:471] Config param scan_layers: True
I0422 07:35:59.941355 134641274300224 pyconfig.py:471] Config param scan_layers_per_stage: False
I0422 07:35:59.941371 134641274300224 pyconfig.py:471] Config param scan_pipeline_iterations: True
I0422 07:35:59.941385 134641274300224 pyconfig.py:471] Config param scan_pipeline_repeats: False
I0422 07:35:59.941401 134641274300224 pyconfig.py:471] Config param set_remat_policy_on_layers_per_stage: False
I0422 07:35:59.941417 134641274300224 pyconfig.py:471] Config param set_remat_policy_on_pipeline_iterations: True
I0422 07:35:59.941431 134641274300224 pyconfig.py:471] Config param sft_train_on_completion_only: False
I0422 07:35:59.941451 134641274300224 pyconfig.py:471] Config param shard_exp_on_fsdp: False
I0422 07:35:59.941466 134641274300224 pyconfig.py:471] Config param shard_mode: ShardMode.AUTO
I0422 07:35:59.941483 134641274300224 pyconfig.py:471] Config param shard_optimizer_over_data: False
I0422 07:35:59.941499 134641274300224 pyconfig.py:471] Config param sharding_strategy: None
I0422 07:35:59.941514 134641274300224 pyconfig.py:471] Config param sharding_tolerance: 0.02
I0422 07:35:59.941531 134641274300224 pyconfig.py:471] Config param shardy: True
I0422 07:35:59.941545 134641274300224 pyconfig.py:471] Config param share_kv_projections: False
I0422 07:35:59.941560 134641274300224 pyconfig.py:471] Config param shared_experts: 0
I0422 07:35:59.941575 134641274300224 pyconfig.py:471] Config param sinkhorn_iterations: 20
I0422 07:35:59.941591 134641274300224 pyconfig.py:471] Config param skip_first_n_steps_for_profiler: 1
I0422 07:35:59.941607 134641274300224 pyconfig.py:471] Config param skip_jax_distributed_system: False
I0422 07:35:59.941621 134641274300224 pyconfig.py:471] Config param skip_step_interval: 128
I0422 07:35:59.941635 134641274300224 pyconfig.py:471] Config param skip_step_on_spikes: False
I0422 07:35:59.941651 134641274300224 pyconfig.py:471] Config param skip_step_scaling_factor: 6.0
I0422 07:35:59.941666 134641274300224 pyconfig.py:471] Config param sliding_window_size: 0
I0422 07:35:59.941682 134641274300224 pyconfig.py:471] Config param solution_end_token: </answer>
I0422 07:35:59.941697 134641274300224 pyconfig.py:471] Config param solution_start_token: <answer>
I0422 07:35:59.941720 134641274300224 pyconfig.py:471] Config param source_checkpoint_layout: orbax
I0422 07:35:59.941736 134641274300224 pyconfig.py:471] Config param sparse_matmul: True
I0422 07:35:59.941752 134641274300224 pyconfig.py:471] Config param spatial_merge_size_for_vit: 2
I0422 07:35:59.941768 134641274300224 pyconfig.py:471] Config param stack_prefill_result_cache: False
I0422 07:35:59.941782 134641274300224 pyconfig.py:471] Config param stack_trace_interval_seconds: 600
I0422 07:35:59.941798 134641274300224 pyconfig.py:471] Config param stack_trace_to_cloud: False
I0422 07:35:59.941814 134641274300224 pyconfig.py:471] Config param step_deviation_interval_seconds: 30
I0422 07:35:59.941830 134641274300224 pyconfig.py:471] Config param steps: 200000
I0422 07:35:59.941846 134641274300224 pyconfig.py:471] Config param stop_strings: None
I0422 07:35:59.941860 134641274300224 pyconfig.py:471] Config param student_overrides: {'model_name': 'llama3.1-8b'}
I0422 07:35:59.941876 134641274300224 pyconfig.py:471] Config param student_params_to_update: None
I0422 07:35:59.941892 134641274300224 pyconfig.py:471] Config param subslice_shape: 
I0422 07:35:59.941907 134641274300224 pyconfig.py:471] Config param swap_space_vllm_gb: 2
I0422 07:35:59.941923 134641274300224 pyconfig.py:471] Config param system_prompt: 
I0422 07:35:59.941938 134641274300224 pyconfig.py:471] Config param target_eval_loss: 0.0
I0422 07:35:59.941954 134641274300224 pyconfig.py:471] Config param teacher_overrides: {'model_name': 'llama3.1-8b'}
I0422 07:35:59.941969 134641274300224 pyconfig.py:471] Config param temperature_tuning: False
I0422 07:35:59.941985 134641274300224 pyconfig.py:471] Config param temporal_patch_size_for_vit: 2
I0422 07:35:59.942000 134641274300224 pyconfig.py:471] Config param tensorboard_dir: /deps/maxtext_output/gpt3-52k_2026-04-22-07-35/tensorboard/
I0422 07:35:59.942015 134641274300224 pyconfig.py:471] Config param tensors_on_device: None
I0422 07:35:59.942030 134641274300224 pyconfig.py:471] Config param tensors_to_offload: None
I0422 07:35:59.942046 134641274300224 pyconfig.py:471] Config param test_batch_start_index: 0
I0422 07:35:59.942061 134641274300224 pyconfig.py:471] Config param tile_size_for_vit: 336
I0422 07:35:59.942077 134641274300224 pyconfig.py:471] Config param tokenize_eval_data: True
I0422 07:35:59.942092 134641274300224 pyconfig.py:471] Config param tokenize_train_data: True
I0422 07:35:59.942107 134641274300224 pyconfig.py:471] Config param tokenizer_path: meta-llama/Llama-3.1-8B
I0422 07:35:59.942124 134641274300224 pyconfig.py:471] Config param tokenizer_type: TokenizerType.HUGGINGFACE
I0422 07:35:59.942142 134641274300224 pyconfig.py:471] Config param topk_routing_group: -1
I0422 07:35:59.942157 134641274300224 pyconfig.py:471] Config param train_data_columns: ['text']
I0422 07:35:59.942173 134641274300224 pyconfig.py:471] Config param train_fraction: 1.0
I0422 07:35:59.942188 134641274300224 pyconfig.py:471] Config param train_image_column: image
I0422 07:35:59.942203 134641274300224 pyconfig.py:471] Config param train_micro_batch_size: -1
I0422 07:35:59.942219 134641274300224 pyconfig.py:471] Config param train_split: train
I0422 07:35:59.942234 134641274300224 pyconfig.py:471] Config param trainable_parameters_mask: []
I0422 07:35:59.942249 134641274300224 pyconfig.py:471] Config param trainable_position_size: 2048
I0422 07:35:59.942265 134641274300224 pyconfig.py:471] Config param trainer_devices_fraction: 0.5
I0422 07:35:59.942280 134641274300224 pyconfig.py:471] Config param upload_all_profiler_results: False
I0422 07:35:59.942295 134641274300224 pyconfig.py:471] Config param use_2d_fsdp_sharding: False
I0422 07:35:59.942311 134641274300224 pyconfig.py:471] Config param use_agentic_rollout: False
I0422 07:35:59.942327 134641274300224 pyconfig.py:471] Config param use_audio: False
I0422 07:35:59.942341 134641274300224 pyconfig.py:471] Config param use_audio_in_video: False
I0422 07:35:59.942357 134641274300224 pyconfig.py:471] Config param use_batch_split_schedule: False
I0422 07:35:59.942372 134641274300224 pyconfig.py:471] Config param use_chat_template: False
I0422 07:35:59.942388 134641274300224 pyconfig.py:471] Config param use_chunked_prefill: False
I0422 07:35:59.942404 134641274300224 pyconfig.py:471] Config param use_custom_sort_vjp: True
I0422 07:35:59.942420 134641274300224 pyconfig.py:471] Config param use_dpo: False
I0422 07:35:59.942440 134641274300224 pyconfig.py:471] Config param use_gather_mosaic_kernel: False
I0422 07:35:59.942454 134641274300224 pyconfig.py:471] Config param use_grpo: True
I0422 07:35:59.942469 134641274300224 pyconfig.py:471] Config param use_indexer: False
I0422 07:35:59.942485 134641274300224 pyconfig.py:471] Config param use_iota_embed: True
I0422 07:35:59.942501 134641274300224 pyconfig.py:471] Config param use_jax_splash: False
I0422 07:35:59.942517 134641274300224 pyconfig.py:471] Config param use_max_logit_estimate: -1
I0422 07:35:59.942533 134641274300224 pyconfig.py:471] Config param use_mrope: False
I0422 07:35:59.942549 134641274300224 pyconfig.py:471] Config param use_multimodal: False
I0422 07:35:59.942564 134641274300224 pyconfig.py:471] Config param use_pathways: True
I0422 07:35:59.942579 134641274300224 pyconfig.py:471] Config param use_post_attn_norm: False
I0422 07:35:59.942594 134641274300224 pyconfig.py:471] Config param use_post_ffw_norm: False
I0422 07:35:59.942610 134641274300224 pyconfig.py:471] Config param use_qk_clip: False
I0422 07:35:59.942625 134641274300224 pyconfig.py:471] Config param use_qk_norm: False
I0422 07:35:59.942640 134641274300224 pyconfig.py:471] Config param use_qk_norm_in_gdn: True
I0422 07:35:59.942656 134641274300224 pyconfig.py:471] Config param use_qwix_quantization: False
I0422 07:35:59.942671 134641274300224 pyconfig.py:471] Config param use_ragged_attention: False
I0422 07:35:59.942686 134641274300224 pyconfig.py:471] Config param use_random_routing: False
I0422 07:35:59.942701 134641274300224 pyconfig.py:471] Config param use_replicator_service: False
I0422 07:35:59.942727 134641274300224 pyconfig.py:471] Config param use_ring_of_experts: False
I0422 07:35:59.942744 134641274300224 pyconfig.py:471] Config param use_sft: False
I0422 07:35:59.942758 134641274300224 pyconfig.py:471] Config param use_splash_scheduler: False
I0422 07:35:59.942774 134641274300224 pyconfig.py:471] Config param use_tokamax_gmm: False
I0422 07:35:59.942789 134641274300224 pyconfig.py:471] Config param use_tokamax_splash: False
I0422 07:35:59.942804 134641274300224 pyconfig.py:471] Config param use_truncation: True
I0422 07:35:59.942820 134641274300224 pyconfig.py:471] Config param use_tunix_gradient_accumulation: False
I0422 07:35:59.942835 134641274300224 pyconfig.py:471] Config param use_untrainable_positional_embedding: False
I0422 07:35:59.942849 134641274300224 pyconfig.py:471] Config param use_vertex_tensorboard: False
I0422 07:35:59.942864 134641274300224 pyconfig.py:471] Config param using_pipeline_parallelism: False
I0422 07:35:59.942878 134641274300224 pyconfig.py:471] Config param v_head_dim: 128
I0422 07:35:59.942893 134641274300224 pyconfig.py:471] Config param v_norm_with_scale: True
I0422 07:35:59.942907 134641274300224 pyconfig.py:471] Config param value_proj: RematLocation.REMAT
I0422 07:35:59.942923 134641274300224 pyconfig.py:471] Config param vertex_tensorboard_project: 
I0422 07:35:59.942939 134641274300224 pyconfig.py:471] Config param vertex_tensorboard_region: 
I0422 07:35:59.942955 134641274300224 pyconfig.py:471] Config param video_path: 
I0422 07:35:59.942970 134641274300224 pyconfig.py:471] Config param video_placeholder: <|video|>
I0422 07:35:59.942986 134641274300224 pyconfig.py:471] Config param vision_output_dim_for_vit: 4096
I0422 07:35:59.943001 134641274300224 pyconfig.py:471] Config param vision_output_length: -1
I0422 07:35:59.943016 134641274300224 pyconfig.py:471] Config param vllm_additional_config: {}
I0422 07:35:59.943032 134641274300224 pyconfig.py:471] Config param vllm_hf_config_path: 
I0422 07:35:59.943046 134641274300224 pyconfig.py:471] Config param vllm_hf_overrides: {}
I0422 07:35:59.943062 134641274300224 pyconfig.py:471] Config param vocab_size: 32000
I0422 07:35:59.943076 134641274300224 pyconfig.py:471] Config param warmup_steps_fraction: 0.1
I0422 07:35:59.943093 134641274300224 pyconfig.py:471] Config param weight_dtype: float32
I0422 07:35:59.943119 134641274300224 pyconfig.py:471] Config param weight_quantization_calibration_method: absmax
I0422 07:35:59.943135 134641274300224 pyconfig.py:471] Config param wi_tile_dlhs_batch_seq: 512
I0422 07:35:59.943152 134641274300224 pyconfig.py:471] Config param wi_tile_dlhs_embed_dim: 1024
I0422 07:35:59.943167 134641274300224 pyconfig.py:471] Config param wi_tile_dlhs_mlp_dim: 1024
I0422 07:35:59.943183 134641274300224 pyconfig.py:471] Config param wi_tile_drhs_batch_seq: 512
I0422 07:35:59.943198 134641274300224 pyconfig.py:471] Config param wi_tile_drhs_embed_dim: 1024
I0422 07:35:59.943214 134641274300224 pyconfig.py:471] Config param wi_tile_drhs_mlp_dim: 1024
I0422 07:35:59.943228 134641274300224 pyconfig.py:471] Config param wi_tile_fwd_batch_seq: 512
I0422 07:35:59.943244 134641274300224 pyconfig.py:471] Config param wi_tile_fwd_embed_dim: 1024
I0422 07:35:59.943259 134641274300224 pyconfig.py:471] Config param wi_tile_fwd_mlp_dim: 1024
I0422 07:35:59.943282 134641274300224 pyconfig.py:471] Config param wo_tile_dlhs_batch_seq: 512
I0422 07:35:59.943309 134641274300224 pyconfig.py:471] Config param wo_tile_dlhs_embed_dim: 1024
I0422 07:35:59.943337 134641274300224 pyconfig.py:471] Config param wo_tile_dlhs_mlp_dim: 1024
I0422 07:35:59.943363 134641274300224 pyconfig.py:471] Config param wo_tile_drhs_batch_seq: 512
I0422 07:35:59.943390 134641274300224 pyconfig.py:471] Config param wo_tile_drhs_embed_dim: 1024
I0422 07:35:59.943416 134641274300224 pyconfig.py:471] Config param wo_tile_drhs_mlp_dim: 1024
I0422 07:35:59.943448 134641274300224 pyconfig.py:471] Config param wo_tile_fwd_batch_seq: 512
I0422 07:35:59.943474 134641274300224 pyconfig.py:471] Config param wo_tile_fwd_embed_dim: 1024
I0422 07:35:59.943499 134641274300224 pyconfig.py:471] Config param wo_tile_fwd_mlp_dim: 1024
I0422 07:35:59.943522 134641274300224 pyconfig.py:471] Config param wsd_decay_steps_fraction: 0.1
I0422 07:35:59.943549 134641274300224 pyconfig.py:471] Config param wsd_decay_style: WsdDecayStyle.LINEAR
I0422 07:35:59.943579 134641274300224 pyconfig.py:471] Config param xprof_e2e_enable_fw_power_level_event: False
I0422 07:35:59.943604 134641274300224 pyconfig.py:471] Config param xprof_e2e_enable_fw_thermal_event: False
I0422 07:35:59.943630 134641274300224 pyconfig.py:471] Config param xprof_e2e_enable_fw_throttle_event: False
I0422 07:35:59.943656 134641274300224 pyconfig.py:471] Config param xprof_tpu_power_trace_level: 0
I0422 07:35:59.943683 134641274300224 pyconfig.py:471] Config param z_loss_multiplier: 0.0
I0422 07:35:59.944006 134641274300224 tokenizer.py:245] Tokenizer path: meta-llama/Llama-2-7b-chat-hf
I0422 07:35:59.944041 134641274300224 tokenizer.py:224] Loading HF tokenizer: meta-llama/Llama-2-7b-chat-hf
I0422 07:36:03.564426 134641274300224 _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`.
I0422 07:36:03.567562 134641274300224 maxtext_utils.py:1565] Num_devices: 32, shape (1, 4, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1)
I0422 07:36:03.567691 134641274300224 train_distill.py:596] Applying logical axis rules for model initialization and training...
I0422 07:36:03.567777 134641274300224 train_distill.py:600] Loading Student from ...
I0422 07:36:03.567807 134641274300224 train_distill.py:169] --- Student Configuration ---
I0422 07:36:03.567831 134641274300224 train_distill.py:170]   Model Name:      gpt3-52k
I0422 07:36:03.567855 134641274300224 train_distill.py:171]   Dimensions:      1 Layers, 16 Emb Dim, 8 Head Dim
I0422 07:36:03.567875 134641274300224 train_distill.py:174]   Attention Heads: 2 Query, 2 KV
I0422 07:36:03.567894 134641274300224 train_distill.py:175]   Vocab Size:      32000
I0422 07:36:03.567912 134641274300224 train_distill.py:176]   Checkpoint:      
I0422 07:36:03.567931 134641274300224 train_distill.py:465] Initializing model: gpt3-52k...
I0422 07:36:04.842165 134641274300224 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...
I0422 07:36:04.842287 134641274300224 train_distill.py:169] --- Teacher Configuration ---
I0422 07:36:04.842316 134641274300224 train_distill.py:170]   Model Name:      gpt3-52k
I0422 07:36:04.842340 134641274300224 train_distill.py:171]   Dimensions:      1 Layers, 16 Emb Dim, 8 Head Dim
I0422 07:36:04.842362 134641274300224 train_distill.py:174]   Attention Heads: 2 Query, 2 KV
I0422 07:36:04.842382 134641274300224 train_distill.py:175]   Vocab Size:      32000
I0422 07:36:04.842400 134641274300224 train_distill.py:176]   Checkpoint:      gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items
I0422 07:36:04.842420 134641274300224 train_distill.py:465] Initializing model: gpt3-52k...
I0422 07:36:05.979972 134641274300224 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0422 07:36:05.980424 134641274300224 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 0x7a73e5b88c20>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0422 07:36:05.980485 134641274300224 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.28
W0422 07:36:06.517019 134641274300224 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
I0422 07:36:07.053951    2087 google_auth_provider.cc:181] Running on GCE, using service account 562977990677-compute@developer.gserviceaccount.com
I0422 07:36:08.195474 134641274300224 checkpointer.py:304] Restoring checkpoint from gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items.
W0422 07:36:10.743096 134641274300224 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.
I0422 07:36:10.743460 134641274300224 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
I0422 07:36:11.904684 134641274300224 checkpointer.py:318] Finished restoring checkpoint in 4.09 seconds from gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items.
I0422 07:36:12.602456 134641274300224 train_distill.py:640] Initializing Data Iterators via MaxText pipeline...
I0422 07:36:12.666217 134641274300224 config.py:112] TensorFlow version 2.20.0 available.
I0422 07:36:12.666699 134641274300224 config.py:125] JAX version 0.8.3 available.
E0422 07:36:15.062609 134641274300224 packing.py:209] PackAndBatchOperation is deprecated. Please use lazy_dataset.FirstFitPackIterDataset instead.
I0422 07:36:15.062845 134641274300224 data_loader.py:408] Adding CopyNumPyArrayToSharedMemory MapTransform.
I0422 07:36:15.065861 134641274300224 train_distill.py:410] Input Pipeline Checkpointing: DISABLED
I0422 07:36:15.065923 134641274300224 train_distill.py:414] Reason: Iterator 'MultiHostDataLoadIterator' is not recognized as Grain (dataset_type='DatasetType.HF', has_save=False)
I0422 07:36:15.065987 134641274300224 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0422 07:36:15.066073 134641274300224 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 0x7a73e5b88c20>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0422 07:36:15.066114 134641274300224 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0422 07:36:15.066145 134641274300224 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 0x7a73e5b88c20>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0422 07:36:15.066190 134641274300224 checkpoint_manager.py:702] [process=2][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers={'model_params': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a6a0a74aae0>, 'optimizer_state': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a5d942b67b0>, 'custom_metadata': <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc82e9970>}, handler_registry=None
I0422 07:36:15.066386 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "model_params". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a6a0a74aae0>` 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`.
I0422 07:36:15.066427 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "optimizer_state". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a5d942b67b0>` 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`.
I0422 07:36:15.066454 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "custom_metadata". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc82e9970>` 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`.
I0422 07:36:15.066488 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a6a0a8e7e30>` 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`.
I0422 07:36:15.066515 134641274300224 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 0x7a6a0a74aae0>, ('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a6a0a74aae0>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a5d942b67b0>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a5d942b67b0>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc82e9970>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc82e9970>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a6a0a8e7e30>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a6a0a8e7e30>}).
I0422 07:36:15.066959 134641274300224 async_checkpointer.py:177] [process=2][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7a73e97c1620> timeout: 600 secs and primary_host=0 for async checkpoint writes
I0422 07:36:17.547425 134641274300224 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://lance-maxtext/pt_ckpt_xpk_main_20260422_071434/pt_distill_nnx_xpk_main_20260422_071434_07_distill_smoke/checkpoints
I0422 07:36:17.549563 134641274300224 checkpoint_manager.py:921] [process=2][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_20260422_071434/pt_distill_nnx_xpk_main_20260422_071434_07_distill_smoke/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7a5bc82e9af0>
I0422 07:36:17.549678 134641274300224 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0422 07:36:17.549757 134641274300224 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 0x7a73e5b88c20>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0422 07:36:17.549796 134641274300224 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None
I0422 07:36:17.549828 134641274300224 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 0x7a73e5b88c20>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
I0422 07:36:17.549864 134641274300224 checkpoint_manager.py:1983] [process=2][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
I0422 07:36:17.549910 134641274300224 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=134641274300224 count=1 at 0x7a6f77374980>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7a5bc82ea1e0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7a5bc82ea240>, _write_futures=[])
I0422 07:36:17.550281 134641274300224 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=134641274300224 count=1 at 0x7a6f77374980>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7a5bc82ea1e0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7a5bc82ea240>, _write_futures=[])
I0422 07:36:17.550307 134641274300224 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=134641274300224 count=1 at 0x7a6f77374980>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7a5bc82ea1e0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7a5bc82ea240>, _write_futures=[])
I0422 07:36:17.550338 134641274300224 checkpoint_manager.py:702] [process=2][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers={'model_params': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a5bc82e91c0>, 'optimizer_state': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a6a0a8af1d0>, 'custom_metadata': <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc83081d0>, 'iter': <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x7a5bc8308110>}, handler_registry=None
I0422 07:36:17.550437 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "model_params". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a5bc82e91c0>` 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`.
I0422 07:36:17.550470 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "optimizer_state". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a6a0a8af1d0>` 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`.
I0422 07:36:17.550493 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "custom_metadata". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc83081d0>` 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`.
I0422 07:36:17.550519 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "iter". Adding handler `<maxtext.common.checkpointing.GrainCheckpointHandler object at 0x7a5bc8308110>` for item "iter" and save args `<class 'maxtext.common.checkpointing.GrainCheckpointSave'>` and restore args `<class 'maxtext.common.checkpointing.GrainCheckpointRestore'>` to `_handler_registry`.
I0422 07:36:17.550541 134641274300224 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc8308ef0>` 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`.
I0422 07:36:17.550564 134641274300224 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 0x7a5bc82e91c0>, ('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a5bc82e91c0>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a6a0a8af1d0>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x7a6a0a8af1d0>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc83081d0>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc83081d0>, ('iter', <class 'maxtext.common.checkpointing.GrainCheckpointSave'>): <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x7a5bc8308110>, ('iter', <class 'maxtext.common.checkpointing.GrainCheckpointRestore'>): <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x7a5bc8308110>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc8308ef0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7a5bc8308ef0>}).
I0422 07:36:17.550635 134641274300224 async_checkpointer.py:177] [process=2][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7a73e97c18a0> timeout: 600 secs and primary_host=0 for async checkpoint writes
I0422 07:36:17.947366 134641274300224 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://lance-maxtext/pt_ckpt_xpk_main_20260422_071434/pt_distill_nnx_xpk_main_20260422_071434_07_distill_smoke/checkpoints
I0422 07:36:18.374578 134641274300224 checkpoint_manager.py:921] [process=2][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_20260422_071434/pt_distill_nnx_xpk_main_20260422_071434_07_distill_smoke/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7a5d94259430>
I0422 07:36:18.375204 134641274300224 train_distill.py:691] Starting Distillation Training...
I0422 07:36:18.375307 134641274300224 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))
I0422 07:36:18.739296 134641274300224 peft_trainer.py:600] Compiled train_step cache size: 0

Training:   0%|          | 0/5 [00:00<?, ?step/s]I0422 07:36:18.741290 134500583466752 grain_pool.py:367] Grain pool will use 1 processes.
I0422 07:36:18.767706 134500583466752 grain_pool.py:440] Grain pool will start child processes.
I0422 07:36:18.772807 134500583466752 grain_pool.py:448] Grain pool started all child processes.
2026-04-22 07:36:24.809856: 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)
Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'}
`rope_scaling`'s factor field must be a float >= 1, got 40
`rope_scaling`'s beta_fast field must be a float, got 32
`rope_scaling`'s beta_slow field must be a float, got 1
Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'}
Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'}
Unrecognized keys in `rope_scaling` for 'rope_type'='yarn': {'rope_theta'}
I0422 07:36:27.753773 134641274300224 utils.py:86] Train loop finished in: 9.0138 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=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=13, process_index=2, coords=(1,3,0), core_on_chip=0), TpuDevice(id=17, process_index=4, coords=(1,4,0), core_on_chip=0), TpuDevice(id=10, process_index=3, coords=(2,2,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=28, process_index=6, coords=(0,7,0), core_on_chip=0), TpuDevice(id=23, process_index=5, coords=(3,5,0), core_on_chip=0), TpuDevice(id=9, process_index=2, coords=(1,2,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=7, process_index=1, coords=(3,1,0), core_on_chip=0), TpuDevice(id=30, process_index=7, coords=(2,7,0), core_on_chip=0), TpuDevice(id=11, process_index=3, coords=(3,2,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=6, process_index=1, coords=(2,1,0), core_on_chip=0), TpuDevice(id=16, process_index=4, coords=(0,4,0), core_on_chip=0), TpuDevice(id=26, process_index=7, coords=(2,6,0), core_on_chip=0), TpuDevice(id=20, process_index=4, coords=(0,5,0), core_on_chip=0), TpuDevice(id=22, process_index=5, coords=(2,5,0), core_on_chip=0), TpuDevice(id=24, process_index=6, coords=(0,6,0), core_on_chip=0), TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=21, process_index=4, coords=(1,5,0), core_on_chip=0), TpuDevice(id=31, process_index=7, coords=(3,7,0), core_on_chip=0), TpuDevice(id=27, process_index=7, coords=(3,6,0), core_on_chip=0), TpuDevice(id=14, process_index=3, coords=(2,3,0), core_on_chip=0), TpuDevice(id=2, process_index=1, coords=(2,0,0), core_on_chip=0), TpuDevice(id=8, process_index=2, coords=(0,2,0), core_on_chip=0), TpuDevice(id=25, process_index=6, coords=(1,6,0), core_on_chip=0), TpuDevice(id=29, process_index=6, coords=(1,7,0), core_on_chip=0), TpuDevice(id=3, process_index=1, coords=(3,0,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(1,1,0), core_on_chip=0)}
I0422 07:36:28.097763 134500583466752 grain_pool.py:542] Grain pool is exiting.
I0422 07:36:28.097866 134500583466752 grain_pool.py:547] Shutting down multiprocessing system.
I0422 07:36:29.554053 134500583466752 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: Wed Apr 22 07:36:37 UTC 2026
EXIT_CODE=1