XPK Start: Tue Apr 21 09:00:16 UTC 2026 2026-04-21 09:00:33.706925: E external/local_xla/xla/stream_executor/cuda/cuda_platform.cc:51] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303) I0421 09:00:37.266352 133802761832256 max_utils.py:273] Attempting to initialize the jax distributed system... INFO:2026-04-21 09:00:46,305:jax._src.distributed:149: Starting JAX distributed service on [::]:8482 I0421 09:00:46.305887 133802761832256 distributed.py:149] Starting JAX distributed service on [::]:8482 INFO:2026-04-21 09:00:46,308:jax._src.distributed:166: Connecting to JAX distributed service on mt-07-distill-smoke-fpmrv-slice-job-0-0.mt-07-distill-smoke-fpmrv:8482 I0421 09:00:46.308248 133802761832256 distributed.py:166] Connecting to JAX distributed service on mt-07-distill-smoke-fpmrv-slice-job-0-0.mt-07-distill-smoke-fpmrv:8482 I0421 09:00:47.817356 133802761832256 max_utils.py:284] Jax distributed system initialized! I0421 09:00:54.077345 133802761832256 max_utils.py:244] Jax distributed system is already initialized. I0421 09:00:54.541831 133802761832256 max_utils.py:244] Jax distributed system is already initialized. I0421 09:00:54.543003 133802761832256 pyconfig.py:432] Config param abort_on_inf_loss: True I0421 09:00:54.543053 133802761832256 pyconfig.py:432] Config param abort_on_nan_loss: True I0421 09:00:54.543080 133802761832256 pyconfig.py:432] Config param act_quantization_calibration_method: absmax I0421 09:00:54.543101 133802761832256 pyconfig.py:432] Config param activation_dropout_for_audio: 0.0 I0421 09:00:54.543120 133802761832256 pyconfig.py:432] Config param activation_function_for_audio: gelu I0421 09:00:54.543138 133802761832256 pyconfig.py:432] Config param activations_in_float32: False I0421 09:00:54.543156 133802761832256 pyconfig.py:432] Config param adam_b1: 0.9 I0421 09:00:54.543175 133802761832256 pyconfig.py:432] Config param adam_b2: 0.95 I0421 09:00:54.543191 133802761832256 pyconfig.py:432] Config param adam_eps: 1e-08 I0421 09:00:54.543214 133802761832256 pyconfig.py:432] Config param adam_eps_root: 0.0 I0421 09:00:54.543229 133802761832256 pyconfig.py:432] Config param adam_weight_decay: 0.1 I0421 09:00:54.543246 133802761832256 pyconfig.py:432] Config param adamw_mask: [] I0421 09:00:54.543262 133802761832256 pyconfig.py:432] Config param add_bos: True I0421 09:00:54.543277 133802761832256 pyconfig.py:432] Config param add_eos: True I0421 09:00:54.543293 133802761832256 pyconfig.py:432] Config param allow_split_physical_axes: False I0421 09:00:54.543308 133802761832256 pyconfig.py:432] Config param ar_cache_axis_order: 1,2,0,3 I0421 09:00:54.543324 133802761832256 pyconfig.py:432] Config param async_checkpointing: True I0421 09:00:54.543339 133802761832256 pyconfig.py:432] Config param async_scheduling: False I0421 09:00:54.543355 133802761832256 pyconfig.py:432] Config param attention: dot_product I0421 09:00:54.543370 133802761832256 pyconfig.py:432] Config param attention_bias: False I0421 09:00:54.543386 133802761832256 pyconfig.py:432] Config param attention_dropout_for_audio: 0.0 I0421 09:00:54.543402 133802761832256 pyconfig.py:432] Config param attention_out: RematLocation.REMAT I0421 09:00:54.543422 133802761832256 pyconfig.py:432] Config param attention_output_dim: -1 I0421 09:00:54.543437 133802761832256 pyconfig.py:432] Config param attention_sink: False I0421 09:00:54.543452 133802761832256 pyconfig.py:432] Config param attention_type: global I0421 09:00:54.543468 133802761832256 pyconfig.py:432] Config param attn_logits_soft_cap: None I0421 09:00:54.543483 133802761832256 pyconfig.py:432] Config param audio_path: I0421 09:00:54.543499 133802761832256 pyconfig.py:432] Config param audio_placeholder: <|audio|> I0421 09:00:54.543515 133802761832256 pyconfig.py:432] Config param autoregressive_decode_assert: I0421 09:00:54.543531 133802761832256 pyconfig.py:432] Config param base_config: base.yml I0421 09:00:54.543545 133802761832256 pyconfig.py:432] Config param base_emb_dim: 16 I0421 09:00:54.543565 133802761832256 pyconfig.py:432] Config param base_mlp_dim: 64 I0421 09:00:54.543580 133802761832256 pyconfig.py:432] Config param base_moe_mlp_dim: -1 I0421 09:00:54.543603 133802761832256 pyconfig.py:432] Config param base_num_decoder_layers: 1 I0421 09:00:54.543625 133802761832256 pyconfig.py:432] Config param base_num_kv_heads: 2 I0421 09:00:54.543664 133802761832256 pyconfig.py:432] Config param base_num_query_heads: 2 I0421 09:00:54.543687 133802761832256 pyconfig.py:432] Config param base_output_directory: I0421 09:00:54.543709 133802761832256 pyconfig.py:432] Config param batch_size: 1 I0421 09:00:54.543730 133802761832256 pyconfig.py:432] Config param batch_split_factor: 1 I0421 09:00:54.543751 133802761832256 pyconfig.py:432] Config param beta_fast: 32 I0421 09:00:54.543773 133802761832256 pyconfig.py:432] Config param beta_slow: 1 I0421 09:00:54.543795 133802761832256 pyconfig.py:432] Config param bwd_quantization_calibration_method: absmax I0421 09:00:54.543817 133802761832256 pyconfig.py:432] Config param capacity_factor: -1.0 I0421 09:00:54.543841 133802761832256 pyconfig.py:432] Config param cast_logits_to_fp32: True I0421 09:00:54.543862 133802761832256 pyconfig.py:432] Config param chat_template: I0421 09:00:54.543884 133802761832256 pyconfig.py:432] Config param chat_template_path: I0421 09:00:54.543908 133802761832256 pyconfig.py:432] Config param checkpoint_conversion_fn: None I0421 09:00:54.543930 133802761832256 pyconfig.py:432] Config param checkpoint_dir: None I0421 09:00:54.543952 133802761832256 pyconfig.py:432] Config param checkpoint_is_quantized: False I0421 09:00:54.543973 133802761832256 pyconfig.py:432] Config param checkpoint_period: 2000 I0421 09:00:54.543993 133802761832256 pyconfig.py:432] Config param checkpoint_storage_concurrent_gb: 96 I0421 09:00:54.544013 133802761832256 pyconfig.py:432] Config param checkpoint_storage_target_data_file_size_bytes: 2147483648 I0421 09:00:54.544036 133802761832256 pyconfig.py:432] Config param checkpoint_storage_use_ocdbt: True I0421 09:00:54.544056 133802761832256 pyconfig.py:432] Config param checkpoint_storage_use_zarr3: True I0421 09:00:54.544077 133802761832256 pyconfig.py:432] Config param checkpoint_todelete_full_path: None I0421 09:00:54.544096 133802761832256 pyconfig.py:432] Config param checkpoint_todelete_subdir: None I0421 09:00:54.544116 133802761832256 pyconfig.py:432] Config param chips_per_vm: 4 I0421 09:00:54.544136 133802761832256 pyconfig.py:432] Config param chunk_attn_window_size: 0 I0421 09:00:54.544157 133802761832256 pyconfig.py:432] Config param collect_stack_trace: False I0421 09:00:54.544176 133802761832256 pyconfig.py:432] Config param colocated_python_checkpointing: False I0421 09:00:54.544197 133802761832256 pyconfig.py:432] Config param colocated_python_data_input: False I0421 09:00:54.544216 133802761832256 pyconfig.py:432] Config param compile_topology: I0421 09:00:54.544235 133802761832256 pyconfig.py:432] Config param compile_topology_num_slices: -1 I0421 09:00:54.544254 133802761832256 pyconfig.py:432] Config param compile_xla_flags: I0421 09:00:54.544273 133802761832256 pyconfig.py:432] Config param compiled_trainstep_file: I0421 09:00:54.544291 133802761832256 pyconfig.py:432] Config param compute_axis_order: 0,1,2,3 I0421 09:00:54.544311 133802761832256 pyconfig.py:432] Config param constant_bound_config: [] I0421 09:00:54.544331 133802761832256 pyconfig.py:432] Config param context: RematLocation.REMAT I0421 09:00:54.544352 133802761832256 pyconfig.py:432] Config param context_parallel_load_balance: True I0421 09:00:54.544369 133802761832256 pyconfig.py:432] Config param context_parallel_size: 1 I0421 09:00:54.544389 133802761832256 pyconfig.py:432] Config param context_parallel_strategy: all_gather I0421 09:00:54.544406 133802761832256 pyconfig.py:432] Config param context_sharding: context I0421 09:00:54.544426 133802761832256 pyconfig.py:432] Config param conv_chunksize_for_audio: 500 I0421 09:00:54.544449 133802761832256 pyconfig.py:432] Config param conv_stride_for_vit: 14 I0421 09:00:54.544466 133802761832256 pyconfig.py:432] Config param cost_estimate_flops_bwd: -1 I0421 09:00:54.544486 133802761832256 pyconfig.py:432] Config param cost_estimate_flops_fwd: -1 I0421 09:00:54.544506 133802761832256 pyconfig.py:432] Config param custom_mesh: I0421 09:00:54.544526 133802761832256 pyconfig.py:432] Config param custom_mesh_and_rule: I0421 09:00:54.544546 133802761832256 pyconfig.py:432] Config param d_model_for_audio: 256 I0421 09:00:54.544568 133802761832256 pyconfig.py:432] Config param data_sharding: (('data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive'),) I0421 09:00:54.544592 133802761832256 pyconfig.py:432] Config param data_shuffle_seed: 0 I0421 09:00:54.544612 133802761832256 pyconfig.py:432] Config param dataset_name: c4/en:3.0.1 I0421 09:00:54.544631 133802761832256 pyconfig.py:432] Config param dataset_path: I0421 09:00:54.544665 133802761832256 pyconfig.py:432] Config param dataset_type: DatasetType.HF I0421 09:00:54.544685 133802761832256 pyconfig.py:432] Config param dcn_autoregressive_parallelism: 1 I0421 09:00:54.544703 133802761832256 pyconfig.py:432] Config param dcn_context_autoregressive_parallelism: 1 I0421 09:00:54.544723 133802761832256 pyconfig.py:432] Config param dcn_context_parallelism: 1 I0421 09:00:54.544741 133802761832256 pyconfig.py:432] Config param dcn_data_parallelism: -1 I0421 09:00:54.544764 133802761832256 pyconfig.py:432] Config param dcn_diloco_parallelism: 1 I0421 09:00:54.544784 133802761832256 pyconfig.py:432] Config param dcn_expert_parallelism: 1 I0421 09:00:54.544829 133802761832256 pyconfig.py:432] Config param dcn_fsdp_parallelism: 1 I0421 09:00:54.544850 133802761832256 pyconfig.py:432] Config param dcn_fsdp_transpose_parallelism: 1 I0421 09:00:54.544867 133802761832256 pyconfig.py:432] Config param dcn_parallelism: [1, -1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] I0421 09:00:54.544883 133802761832256 pyconfig.py:432] Config param dcn_pipeline_parallelism: 1 I0421 09:00:54.544898 133802761832256 pyconfig.py:432] Config param dcn_sequence_parallelism: 1 I0421 09:00:54.544920 133802761832256 pyconfig.py:432] Config param dcn_tensor_parallelism: 1 I0421 09:00:54.544944 133802761832256 pyconfig.py:432] Config param dcn_tensor_sequence_parallelism: 1 I0421 09:00:54.544967 133802761832256 pyconfig.py:432] Config param dcn_tensor_transpose_parallelism: 1 I0421 09:00:54.544990 133802761832256 pyconfig.py:432] Config param debug: {'rl': False} I0421 09:00:54.545016 133802761832256 pyconfig.py:432] Config param debug_sharding: False I0421 09:00:54.545039 133802761832256 pyconfig.py:432] Config param decode_sampling_nucleus_p: -1 I0421 09:00:54.545057 133802761832256 pyconfig.py:432] Config param decode_sampling_strategy: SamplingStrategy.GREEDY I0421 09:00:54.545075 133802761832256 pyconfig.py:432] Config param decode_sampling_temperature: 1.0 I0421 09:00:54.545092 133802761832256 pyconfig.py:432] Config param decode_sampling_top_k: 0 I0421 09:00:54.545116 133802761832256 pyconfig.py:432] Config param decoder_block: DecoderBlockType.GPT3 I0421 09:00:54.545140 133802761832256 pyconfig.py:432] Config param decoder_layer_input: RematLocation.DEVICE I0421 09:00:54.545163 133802761832256 pyconfig.py:432] Config param deepstack_visual_indexes_for_vit: [] I0421 09:00:54.545185 133802761832256 pyconfig.py:432] Config param degenerate_group_masking: True I0421 09:00:54.545206 133802761832256 pyconfig.py:432] Config param dense_init_scale: 1.0 I0421 09:00:54.545230 133802761832256 pyconfig.py:432] Config param diloco_outer_lr: 0.3 I0421 09:00:54.545253 133802761832256 pyconfig.py:432] Config param diloco_outer_momentum: 0.9 I0421 09:00:54.545276 133802761832256 pyconfig.py:432] Config param diloco_sync_period: 36 I0421 09:00:54.545296 133802761832256 pyconfig.py:432] Config param distill_alpha: 0.5 I0421 09:00:54.545320 133802761832256 pyconfig.py:432] Config param distill_alpha_end: None I0421 09:00:54.545347 133802761832256 pyconfig.py:432] Config param distill_alpha_schedule: constant I0421 09:00:54.545371 133802761832256 pyconfig.py:432] Config param distill_beta: 0.0 I0421 09:00:54.545393 133802761832256 pyconfig.py:432] Config param distill_beta_end: None I0421 09:00:54.545412 133802761832256 pyconfig.py:432] Config param distill_beta_schedule: constant I0421 09:00:54.545428 133802761832256 pyconfig.py:432] Config param distill_feature_loss_type: cosine I0421 09:00:54.545442 133802761832256 pyconfig.py:432] Config param distill_layer_indices: None I0421 09:00:54.545457 133802761832256 pyconfig.py:432] Config param distill_temperature: 1.0 I0421 09:00:54.545473 133802761832256 pyconfig.py:432] Config param distill_temperature_end: None I0421 09:00:54.545487 133802761832256 pyconfig.py:432] Config param distill_temperature_schedule: constant I0421 09:00:54.545503 133802761832256 pyconfig.py:432] Config param downsample_hidden_size_for_audio: 256 I0421 09:00:54.545519 133802761832256 pyconfig.py:432] Config param dpo_beta: 0.1 I0421 09:00:54.545535 133802761832256 pyconfig.py:432] Config param dpo_label_smoothing: 0.0 I0421 09:00:54.545550 133802761832256 pyconfig.py:432] Config param dq_reduction_steps: 0 I0421 09:00:54.545569 133802761832256 pyconfig.py:432] Config param dropout_rate: 0.0 I0421 09:00:54.545583 133802761832256 pyconfig.py:432] Config param dtype: bfloat16 I0421 09:00:54.545613 133802761832256 pyconfig.py:432] Config param dtype_mm: float32 I0421 09:00:54.545629 133802761832256 pyconfig.py:432] Config param dump_hlo: False I0421 09:00:54.545653 133802761832256 pyconfig.py:432] Config param dump_hlo_delete_local_after: True I0421 09:00:54.545669 133802761832256 pyconfig.py:432] Config param dump_hlo_gcs_dir: gpt3-52k_2026-04-21-09-00/xla_dump I0421 09:00:54.545685 133802761832256 pyconfig.py:432] Config param dump_hlo_local_dir: /tmp/xla_dump/ I0421 09:00:54.545698 133802761832256 pyconfig.py:432] Config param dump_hlo_local_module_name: jit_train_step I0421 09:00:54.545713 133802761832256 pyconfig.py:432] Config param dump_hlo_module_name: jit_train_step I0421 09:00:54.545729 133802761832256 pyconfig.py:432] Config param dump_hlo_upload_all: False I0421 09:00:54.545742 133802761832256 pyconfig.py:432] Config param dump_hlo_xla_flags: I0421 09:00:54.545758 133802761832256 pyconfig.py:432] Config param dump_jaxpr: False I0421 09:00:54.545772 133802761832256 pyconfig.py:432] Config param dump_jaxpr_delete_local_after: True I0421 09:00:54.545787 133802761832256 pyconfig.py:432] Config param dump_jaxpr_gcs_dir: gpt3-52k_2026-04-21-09-00/jaxpr_dump I0421 09:00:54.545801 133802761832256 pyconfig.py:432] Config param dump_jaxpr_local_dir: /tmp/jaxpr_dump/ I0421 09:00:54.545815 133802761832256 pyconfig.py:432] Config param dump_step: -1 I0421 09:00:54.545831 133802761832256 pyconfig.py:432] Config param elastic_enabled: False I0421 09:00:54.545844 133802761832256 pyconfig.py:432] Config param elastic_max_retries: 10 I0421 09:00:54.545860 133802761832256 pyconfig.py:432] Config param elastic_timeout_seconds: 300 I0421 09:00:54.545876 133802761832256 pyconfig.py:432] Config param emb_dim: 16 I0421 09:00:54.545892 133802761832256 pyconfig.py:432] Config param enable_autocheckpoint: False I0421 09:00:54.545905 133802761832256 pyconfig.py:432] Config param enable_checkpoint_cloud_logger: False I0421 09:00:54.545919 133802761832256 pyconfig.py:432] Config param enable_checkpointing: True I0421 09:00:54.545935 133802761832256 pyconfig.py:432] Config param enable_continuous_checkpointing: False I0421 09:00:54.545949 133802761832256 pyconfig.py:432] Config param enable_data_shuffling: True I0421 09:00:54.545964 133802761832256 pyconfig.py:432] Config param enable_diloco: False I0421 09:00:54.545978 133802761832256 pyconfig.py:432] Config param enable_dp_attention: False I0421 09:00:54.545994 133802761832256 pyconfig.py:432] Config param enable_dropout: False I0421 09:00:54.546008 133802761832256 pyconfig.py:432] Config param enable_emergency_checkpoint: False I0421 09:00:54.546023 133802761832256 pyconfig.py:432] Config param enable_expert_parallel: False I0421 09:00:54.546039 133802761832256 pyconfig.py:432] Config param enable_gcp_goodput_metrics: True I0421 09:00:54.546052 133802761832256 pyconfig.py:432] Config param enable_gcp_step_deviation_metrics: True I0421 09:00:54.546068 133802761832256 pyconfig.py:432] Config param enable_goodput_recording: False I0421 09:00:54.546082 133802761832256 pyconfig.py:432] Config param enable_jax_profiler: False I0421 09:00:54.546096 133802761832256 pyconfig.py:432] Config param enable_llm_inference_pool: False I0421 09:00:54.546112 133802761832256 pyconfig.py:432] Config param enable_model_warmup: False I0421 09:00:54.546125 133802761832256 pyconfig.py:432] Config param enable_multi_tier_checkpointing: False I0421 09:00:54.546140 133802761832256 pyconfig.py:432] Config param enable_nnx: False I0421 09:00:54.546156 133802761832256 pyconfig.py:432] Config param enable_orbax_v1: False I0421 09:00:54.546169 133802761832256 pyconfig.py:432] Config param enable_padding_causal_mask: True I0421 09:00:54.546185 133802761832256 pyconfig.py:432] Config param enable_pathways_goodput: False I0421 09:00:54.546199 133802761832256 pyconfig.py:432] Config param enable_prefix_caching: False I0421 09:00:54.546214 133802761832256 pyconfig.py:432] Config param enable_rampup_batch_size: False I0421 09:00:54.546239 133802761832256 pyconfig.py:432] Config param enable_single_controller: False I0421 09:00:54.546255 133802761832256 pyconfig.py:432] Config param enable_single_replica_ckpt_restoring: False I0421 09:00:54.546269 133802761832256 pyconfig.py:432] Config param enable_tensorboard: True I0421 09:00:54.546284 133802761832256 pyconfig.py:432] Config param enable_tunix_perf_metrics: False I0421 09:00:54.546300 133802761832256 pyconfig.py:432] Config param encoder_attention_heads_for_audio: 4 I0421 09:00:54.546318 133802761832256 pyconfig.py:432] Config param encoder_ffn_dim_for_audio: 512 I0421 09:00:54.546332 133802761832256 pyconfig.py:432] Config param encoder_layers_for_audio: 2 I0421 09:00:54.546347 133802761832256 pyconfig.py:432] Config param engram: RematLocation.REMAT I0421 09:00:54.546362 133802761832256 pyconfig.py:432] Config param engram_head_dim: 1280 I0421 09:00:54.546377 133802761832256 pyconfig.py:432] Config param engram_kernel_size: 4 I0421 09:00:54.546393 133802761832256 pyconfig.py:432] Config param engram_layers: [] I0421 09:00:54.546416 133802761832256 pyconfig.py:432] Config param engram_max_ngram_size: 3 I0421 09:00:54.546432 133802761832256 pyconfig.py:432] Config param engram_num_heads: 8 I0421 09:00:54.546446 133802761832256 pyconfig.py:432] Config param engram_seed: 0 I0421 09:00:54.546462 133802761832256 pyconfig.py:432] Config param engram_vocab_bases: [] I0421 09:00:54.546477 133802761832256 pyconfig.py:432] Config param epsilon_high: None I0421 09:00:54.546491 133802761832256 pyconfig.py:432] Config param eval_corr_lst: False I0421 09:00:54.546507 133802761832256 pyconfig.py:432] Config param eval_data_columns: ['text'] I0421 09:00:54.546521 133802761832256 pyconfig.py:432] Config param eval_dataset_name: c4/en:3.0.1 I0421 09:00:54.546537 133802761832256 pyconfig.py:432] Config param eval_image_column: image I0421 09:00:54.546551 133802761832256 pyconfig.py:432] Config param eval_interval: -1 I0421 09:00:54.546570 133802761832256 pyconfig.py:432] Config param eval_make_lst: False I0421 09:00:54.546585 133802761832256 pyconfig.py:432] Config param eval_per_device_batch_size: 2 I0421 09:00:54.546599 133802761832256 pyconfig.py:432] Config param eval_sampling_strategy: greedy I0421 09:00:54.546614 133802761832256 pyconfig.py:432] Config param eval_split: validation I0421 09:00:54.546629 133802761832256 pyconfig.py:432] Config param eval_steps: -1 I0421 09:00:54.546654 133802761832256 pyconfig.py:432] Config param expansion_factor_real_data: -1.0 I0421 09:00:54.546670 133802761832256 pyconfig.py:432] Config param final_logits_soft_cap: None I0421 09:00:54.546685 133802761832256 pyconfig.py:432] Config param first_num_dense_layers: 0 I0421 09:00:54.546699 133802761832256 pyconfig.py:432] Config param float32_gate_logits: False I0421 09:00:54.546714 133802761832256 pyconfig.py:432] Config param float32_logits: False I0421 09:00:54.546728 133802761832256 pyconfig.py:432] Config param float32_qk_product: False I0421 09:00:54.546744 133802761832256 pyconfig.py:432] Config param float32_weight_sum: True I0421 09:00:54.546758 133802761832256 pyconfig.py:432] Config param force_q_layout: False I0421 09:00:54.546773 133802761832256 pyconfig.py:432] Config param force_unroll: False I0421 09:00:54.546787 133802761832256 pyconfig.py:432] Config param freeze_audio_encoder_params: True I0421 09:00:54.546803 133802761832256 pyconfig.py:432] Config param freeze_vision_encoder_params: True I0421 09:00:54.546818 133802761832256 pyconfig.py:432] Config param fused_mlp: False I0421 09:00:54.546833 133802761832256 pyconfig.py:432] Config param fused_qkv: True I0421 09:00:54.546848 133802761832256 pyconfig.py:432] Config param gcs_metrics: False I0421 09:00:54.546864 133802761832256 pyconfig.py:432] Config param gdn_chunk_size: 64 I0421 09:00:54.546880 133802761832256 pyconfig.py:432] Config param gdn_conv_kernel_dim: 4 I0421 09:00:54.546894 133802761832256 pyconfig.py:432] Config param gdn_key_head_dim: 128 I0421 09:00:54.546909 133802761832256 pyconfig.py:432] Config param gdn_num_key_heads: 16 I0421 09:00:54.546925 133802761832256 pyconfig.py:432] Config param gdn_num_value_heads: 32 I0421 09:00:54.546942 133802761832256 pyconfig.py:432] Config param gdn_value_head_dim: 128 I0421 09:00:54.546957 133802761832256 pyconfig.py:432] Config param generate_padding_batch_eval: False I0421 09:00:54.546972 133802761832256 pyconfig.py:432] Config param generate_padding_batch_train: False I0421 09:00:54.546986 133802761832256 pyconfig.py:432] Config param generate_slice: v5e-16 I0421 09:00:54.547001 133802761832256 pyconfig.py:432] Config param generation_configs: {} I0421 09:00:54.547015 133802761832256 pyconfig.py:432] Config param global_batch_size_to_eval_on: 64 I0421 09:00:54.547030 133802761832256 pyconfig.py:432] Config param global_batch_size_to_load: 512 I0421 09:00:54.547044 133802761832256 pyconfig.py:432] Config param global_batch_size_to_load_eval: 64 I0421 09:00:54.547059 133802761832256 pyconfig.py:432] Config param global_batch_size_to_load_increment: None I0421 09:00:54.547074 133802761832256 pyconfig.py:432] Config param global_batch_size_to_load_start: None I0421 09:00:54.547088 133802761832256 pyconfig.py:432] Config param global_batch_size_to_train_on: 512 I0421 09:00:54.547104 133802761832256 pyconfig.py:432] Config param global_head_dim: 0 I0421 09:00:54.547118 133802761832256 pyconfig.py:432] Config param global_num_kv_heads: 0 I0421 09:00:54.547133 133802761832256 pyconfig.py:432] Config param global_parameter_scale: 1 I0421 09:00:54.547148 133802761832256 pyconfig.py:432] Config param global_rampup_samples: 500 I0421 09:00:54.547162 133802761832256 pyconfig.py:432] Config param global_rope_max_timescale: -1 I0421 09:00:54.547177 133802761832256 pyconfig.py:432] Config param global_rope_proportion: 0.25 I0421 09:00:54.547192 133802761832256 pyconfig.py:432] Config param goodput_upload_interval_seconds: 30 I0421 09:00:54.547206 133802761832256 pyconfig.py:432] Config param grad_dtype: float32 I0421 09:00:54.547241 133802761832256 pyconfig.py:432] Config param gradient_accumulation_steps: 8 I0421 09:00:54.547256 133802761832256 pyconfig.py:432] Config param gradient_clipping_threshold: 1.0 I0421 09:00:54.547272 133802761832256 pyconfig.py:432] Config param grain_data_source_max_workers: 16 I0421 09:00:54.547286 133802761832256 pyconfig.py:432] Config param grain_eval_files: I0421 09:00:54.547301 133802761832256 pyconfig.py:432] Config param grain_file_type: arrayrecord I0421 09:00:54.547316 133802761832256 pyconfig.py:432] Config param grain_num_threads: 16 I0421 09:00:54.547331 133802761832256 pyconfig.py:432] Config param grain_num_threads_eval: 16 I0421 09:00:54.547346 133802761832256 pyconfig.py:432] Config param grain_packing_type: first_fit I0421 09:00:54.547362 133802761832256 pyconfig.py:432] Config param grain_per_worker_buffer_size: 1 I0421 09:00:54.547376 133802761832256 pyconfig.py:432] Config param grain_per_worker_buffer_size_eval: 1 I0421 09:00:54.547392 133802761832256 pyconfig.py:432] Config param grain_prefetch_buffer_size: 500 I0421 09:00:54.547406 133802761832256 pyconfig.py:432] Config param grain_prefetch_buffer_size_eval: 500 I0421 09:00:54.547421 133802761832256 pyconfig.py:432] Config param grain_ram_budget_mb: 1024 I0421 09:00:54.547437 133802761832256 pyconfig.py:432] Config param grain_shuffle_buffer_size: 100 I0421 09:00:54.547451 133802761832256 pyconfig.py:432] Config param grain_train_files: I0421 09:00:54.547466 133802761832256 pyconfig.py:432] Config param grain_train_mixture_config_path: I0421 09:00:54.547479 133802761832256 pyconfig.py:432] Config param grain_worker_count: 1 I0421 09:00:54.547494 133802761832256 pyconfig.py:432] Config param grain_worker_count_eval: 1 I0421 09:00:54.547509 133802761832256 pyconfig.py:432] Config param grpo_beta: 0.08 I0421 09:00:54.547525 133802761832256 pyconfig.py:432] Config param grpo_epsilon: 0.2 I0421 09:00:54.547540 133802761832256 pyconfig.py:432] Config param hardware: tpu I0421 09:00:54.547555 133802761832256 pyconfig.py:432] Config param hbm_utilization_vllm: 0.72 I0421 09:00:54.547573 133802761832256 pyconfig.py:432] Config param head_dim: 8 I0421 09:00:54.547588 133802761832256 pyconfig.py:432] Config param heartbeat_reporting_interval_in_seconds: 5 I0421 09:00:54.547603 133802761832256 pyconfig.py:432] Config param hf_data_dir: None I0421 09:00:54.547617 133802761832256 pyconfig.py:432] Config param hf_eval_files: None I0421 09:00:54.547632 133802761832256 pyconfig.py:432] Config param hf_eval_split: None I0421 09:00:54.547660 133802761832256 pyconfig.py:432] Config param hf_name: None I0421 09:00:54.547674 133802761832256 pyconfig.py:432] Config param hf_path: OptimalScale/ClimbMix I0421 09:00:54.547689 133802761832256 pyconfig.py:432] Config param hf_train_files: None I0421 09:00:54.547705 133802761832256 pyconfig.py:432] Config param hidden_size_for_vit: 1408 I0421 09:00:54.547720 133802761832256 pyconfig.py:432] Config param hide_profiler_step_metric: False I0421 09:00:54.547735 133802761832256 pyconfig.py:432] Config param ici_autoregressive_parallelism: 1 I0421 09:00:54.547751 133802761832256 pyconfig.py:432] Config param ici_context_autoregressive_parallelism: 1 I0421 09:00:54.547765 133802761832256 pyconfig.py:432] Config param ici_context_parallelism: 1 I0421 09:00:54.547780 133802761832256 pyconfig.py:432] Config param ici_data_parallelism: 1 I0421 09:00:54.547795 133802761832256 pyconfig.py:432] Config param ici_diloco_parallelism: 1 I0421 09:00:54.547809 133802761832256 pyconfig.py:432] Config param ici_expert_parallelism: 1 I0421 09:00:54.547824 133802761832256 pyconfig.py:432] Config param ici_fsdp_parallelism: -1 I0421 09:00:54.547838 133802761832256 pyconfig.py:432] Config param ici_fsdp_transpose_parallelism: 1 I0421 09:00:54.547854 133802761832256 pyconfig.py:432] Config param ici_parallelism: [1, 1, 1, -1, 1, 1, 1, 1, 1, 1, 1, 1, 1] I0421 09:00:54.547869 133802761832256 pyconfig.py:432] Config param ici_pipeline_parallelism: 1 I0421 09:00:54.547884 133802761832256 pyconfig.py:432] Config param ici_sequence_parallelism: 1 I0421 09:00:54.547898 133802761832256 pyconfig.py:432] Config param ici_tensor_parallelism: 1 I0421 09:00:54.547913 133802761832256 pyconfig.py:432] Config param ici_tensor_sequence_parallelism: 1 I0421 09:00:54.547927 133802761832256 pyconfig.py:432] Config param ici_tensor_transpose_parallelism: 1 I0421 09:00:54.547943 133802761832256 pyconfig.py:432] Config param image_path: I0421 09:00:54.547958 133802761832256 pyconfig.py:432] Config param image_placeholder: <|image|> I0421 09:00:54.547972 133802761832256 pyconfig.py:432] Config param image_size_for_vit: 896 I0421 09:00:54.547987 133802761832256 pyconfig.py:432] Config param indexer_head_dim: 128 I0421 09:00:54.548002 133802761832256 pyconfig.py:432] Config param indexer_loss_scaling_factor: 0.0 I0421 09:00:54.548018 133802761832256 pyconfig.py:432] Config param indexer_n_heads: 64 I0421 09:00:54.548032 133802761832256 pyconfig.py:432] Config param indexer_sparse_training: False I0421 09:00:54.548047 133802761832256 pyconfig.py:432] Config param indexer_topk: 2048 I0421 09:00:54.548061 133802761832256 pyconfig.py:432] Config param inference_benchmark_test: False I0421 09:00:54.548076 133802761832256 pyconfig.py:432] Config param inference_metadata_file: I0421 09:00:54.548091 133802761832256 pyconfig.py:432] Config param inference_microbenchmark_log_file_path: I0421 09:00:54.548106 133802761832256 pyconfig.py:432] Config param inference_microbenchmark_loop_iters: 10 I0421 09:00:54.548120 133802761832256 pyconfig.py:432] Config param inference_microbenchmark_num_samples: [1, 2, 3, 4, 5] I0421 09:00:54.548136 133802761832256 pyconfig.py:432] Config param inference_microbenchmark_prefill_lengths: 64,128,256,512,1024 I0421 09:00:54.548151 133802761832256 pyconfig.py:432] Config param inference_microbenchmark_stages: prefill,generate I0421 09:00:54.548165 133802761832256 pyconfig.py:432] Config param inference_server: MaxtextInterleavedServer I0421 09:00:54.548181 133802761832256 pyconfig.py:432] Config param inhomogeneous_layer_cycle_interval: 1 I0421 09:00:54.548195 133802761832256 pyconfig.py:432] Config param init_weights_seed: 0 I0421 09:00:54.548211 133802761832256 pyconfig.py:432] Config param input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activation_norm_length'] I0421 09:00:54.548226 133802761832256 pyconfig.py:432] Config param interleave_moe_layer_step: 1 I0421 09:00:54.548241 133802761832256 pyconfig.py:432] Config param intermediate_size_for_vit: 5632 I0421 09:00:54.548255 133802761832256 pyconfig.py:432] Config param internal_compile: False I0421 09:00:54.548269 133802761832256 pyconfig.py:432] Config param internal_compile_num_devices: -1 I0421 09:00:54.548284 133802761832256 pyconfig.py:432] Config param jax_cache_dir: ~/jax_cache I0421 09:00:54.548298 133802761832256 pyconfig.py:432] Config param jax_debug_log_modules: I0421 09:00:54.548313 133802761832256 pyconfig.py:432] Config param jax_distributed_initialization_timeout: 300 I0421 09:00:54.548328 133802761832256 pyconfig.py:432] Config param jax_profiler_port: 9999 I0421 09:00:54.548343 133802761832256 pyconfig.py:432] Config param key_proj: RematLocation.REMAT I0421 09:00:54.548358 133802761832256 pyconfig.py:432] Config param kv_cache_buffer: 256 I0421 09:00:54.548373 133802761832256 pyconfig.py:432] Config param kv_lora_rank: 512 I0421 09:00:54.548389 133802761832256 pyconfig.py:432] Config param kv_quant_axis: KvQuantAxis.HEADS_AND_DKV I0421 09:00:54.548406 133802761832256 pyconfig.py:432] Config param kv_quant_dtype: int8 I0421 09:00:54.548421 133802761832256 pyconfig.py:432] Config param kv_wa_proj: RematLocation.REMAT I0421 09:00:54.548436 133802761832256 pyconfig.py:432] Config param learning_rate: 0.0002 I0421 09:00:54.548451 133802761832256 pyconfig.py:432] Config param learning_rate_final_fraction: 0.1 I0421 09:00:54.548466 133802761832256 pyconfig.py:432] Config param learning_rate_schedule_steps: 200000 I0421 09:00:54.548480 133802761832256 pyconfig.py:432] Config param load_balance_loss_weight: 0.0 I0421 09:00:54.548496 133802761832256 pyconfig.py:432] Config param load_checkpoint_only_once: False I0421 09:00:54.548510 133802761832256 pyconfig.py:432] Config param load_from_prefill_dir: False I0421 09:00:54.548526 133802761832256 pyconfig.py:432] Config param load_full_state_path: I0421 09:00:54.548541 133802761832256 pyconfig.py:432] Config param load_parameters_path: gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items I0421 09:00:54.548557 133802761832256 pyconfig.py:432] Config param local_checkpoint_directory: I0421 09:00:54.548574 133802761832256 pyconfig.py:432] Config param local_checkpoint_period: 0 I0421 09:00:54.548589 133802761832256 pyconfig.py:432] Config param local_rope_max_timescale: -1 I0421 09:00:54.548604 133802761832256 pyconfig.py:432] Config param local_rope_proportion: 1.0 I0421 09:00:54.548619 133802761832256 pyconfig.py:432] Config param log_config: True I0421 09:00:54.548635 133802761832256 pyconfig.py:432] Config param log_period: 10 I0421 09:00:54.548664 133802761832256 pyconfig.py:432] Config param logical_axis_rules: (('activation_embed_and_logits_batch', ('data', 'stage', 'fsdp', 'fsdp_transpose', 'expert')), ('activation_embed_and_logits_batch_sequence', ('data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('activation_vocab', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('activation_vocab', ('tensor', 'tensor_transpose')), ('activation_vocab', 'tensor_sequence'), ('activation_vocab', ('sequence', 'context')), ('vocab', ('tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive')), ('embed_vocab', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('activation_heads', ('tensor', 'tensor_transpose', 'sequence', 'tensor_sequence', 'autoregressive')), ('activation_kv_heads', ('tensor', 'tensor_transpose', 'sequence', 'tensor_sequence')), ('activation_attn_length', ('sequence', 'context')), ('activation_attn_length', ('context',)), ('activation_q_length', ('context',)), ('activation_kv_length', ()), ('activation_attn_embed', ('tensor', 'tensor_transpose')), ('activation_kv', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('activation_kv_batch', ('data', 'fsdp', 'fsdp_transpose', 'expert')), ('activation_kv_head_dim', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('heads', ('tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive')), ('q_heads', ('tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive')), ('kv_heads', ('tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive')), ('qkv', ()), ('kv', ()), ('kv_head_dim', ()), ('q_lora', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'tensor_transpose', 'expert')), ('q_lora', ('fsdp', 'sequence', 'context', 'tensor_transpose', 'expert')), ('q_lora', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('q_lora', ('fsdp', 'sequence', 'context', 'expert')), ('q_lora_up_proj', ()), ('kv_lora', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'tensor_transpose', 'expert')), ('kv_lora', ('fsdp', 'sequence', 'context', 'tensor_transpose', 'expert')), ('kv_lora', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('kv_lora', ('fsdp', 'sequence', 'context', 'expert')), ('kv_lora_up_proj', ()), ('activation_batch_moe', ('data', 'fsdp', 'fsdp_transpose')), ('activation_length_moe', ('sequence', 'context')), ('activation_length_moe', ('context',)), ('activation_norm_length_moe', ('tensor_sequence', 'context', 'sequence')), ('activation_embed_moe', ('tensor', 'tensor_transpose')), ('activation_mlp_moe', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('activation_exp', ('expert',)), ('exp', 'expert'), ('mlp_moe', ('fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive')), ('embed_moe', ('fsdp', 'fsdp_transpose', 'sequence', 'tensor_transpose', 'context')), ('embed_moe', ('fsdp', 'sequence', 'tensor_transpose', 'context')), ('embed_moe', ('fsdp', 'fsdp_transpose', 'sequence', 'context')), ('embed_moe', ('fsdp', 'sequence', 'context')), ('activation_mlp', ('tensor', 'tensor_transpose', 'tensor_sequence')), ('activation_batch', ('data', 'fsdp', 'fsdp_transpose', 'expert')), ('activation_length', ('sequence', 'context')), ('activation_length', ('context',)), ('activation_norm_length', ('tensor_sequence', 'context', 'sequence')), ('activation_embed', ('tensor', 'tensor_transpose')), ('activation_stage', 'stage'), ('mlp', ('fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive')), ('embed', ('fsdp', 'fsdp_transpose', 'sequence', 'tensor_transpose', 'context', 'expert')), ('embed', ('fsdp', 'sequence', 'tensor_transpose', 'context', 'expert')), ('embed', ('fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert')), ('embed', ('fsdp', 'sequence', 'context', 'expert')), ('norm', ('tensor', 'tensor_transpose')), ('layers', 'stage'), ('diloco', 'diloco'), ('engram_dim', ('tensor',)), ('dense_layers', ()), ('moe_layers', ()), ('mhc', ()), ('prefill_activation_length', ('sequence', 'context')), ('prefill_activation_norm_length', ('tensor_sequence', 'context', 'sequence')), ('activation_prefill_kv_batch', ('data', 'fsdp', 'fsdp_transpose', 'expert')), ('decode_batch', ('data', 'fsdp', 'fsdp_transpose', 'expert')), ('decode_length', ('sequence',)), ('cache_heads', ('autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence')), ('cache_heads', ('autoregressive', 'tensor', 'tensor_sequence')), ('paged_kv_heads', ('tensor',)), ('cache_batch_prefill', ()), ('cache_batch', ()), ('cache_heads_none', ()), ('cache_kv', ()), ('cache_sequence', ()), ('num_pages', ()), ('tokens_per_page', ()), ('paged_kv_head_dim_size', ()), ('mlp_no_fsdp', ('tensor', 'tensor_sequence', 'autoregressive')), ('embed_tensor_transpose', ('tensor_transpose',)), ('exp_with_fsdp', 'fsdp')) I0421 09:00:54.548740 133802761832256 pyconfig.py:432] Config param logits_dot_in_fp32: False I0421 09:00:54.548755 133802761832256 pyconfig.py:432] Config param logits_via_embedding: True I0421 09:00:54.548770 133802761832256 pyconfig.py:432] Config param lora_input_adapters_path: I0421 09:00:54.548785 133802761832256 pyconfig.py:432] Config param loss_algo: grpo I0421 09:00:54.548800 133802761832256 pyconfig.py:432] Config param lr_schedule_type: LearningRateScheduleType.COSINE I0421 09:00:54.548817 133802761832256 pyconfig.py:432] Config param managed_mldiagnostics: False I0421 09:00:54.548832 133802761832256 pyconfig.py:432] Config param managed_mldiagnostics_dir: None I0421 09:00:54.548846 133802761832256 pyconfig.py:432] Config param managed_mldiagnostics_run_group: I0421 09:00:54.548862 133802761832256 pyconfig.py:432] Config param matmul_precision: MatmulPrecision.DEFAULT I0421 09:00:54.548880 133802761832256 pyconfig.py:432] Config param max_checkify: False I0421 09:00:54.548894 133802761832256 pyconfig.py:432] Config param max_concurrency: 256 I0421 09:00:54.548910 133802761832256 pyconfig.py:432] Config param max_corpus_chars: 10000000 I0421 09:00:54.548925 133802761832256 pyconfig.py:432] Config param max_num_batched_tokens: None I0421 09:00:54.548940 133802761832256 pyconfig.py:432] Config param max_num_checkpoints_to_keep: None I0421 09:00:54.548955 133802761832256 pyconfig.py:432] Config param max_num_images_per_example: -1 I0421 09:00:54.548970 133802761832256 pyconfig.py:432] Config param max_num_seqs: None I0421 09:00:54.548985 133802761832256 pyconfig.py:432] Config param max_position_embeddings: 163840 I0421 09:00:54.549000 133802761832256 pyconfig.py:432] Config param max_prefill_predict_length: 64 I0421 09:00:54.549014 133802761832256 pyconfig.py:432] Config param max_sample_len_for_audio: 10000 I0421 09:00:54.549029 133802761832256 pyconfig.py:432] Config param max_segments_per_seq: -1 I0421 09:00:54.549045 133802761832256 pyconfig.py:432] Config param max_source_positions_for_audio: 1500 I0421 09:00:54.549060 133802761832256 pyconfig.py:432] Config param max_target_length: 2048 I0421 09:00:54.549076 133802761832256 pyconfig.py:432] Config param max_timescale_for_audio: 10000.0 I0421 09:00:54.549091 133802761832256 pyconfig.py:432] Config param megablox: True I0421 09:00:54.549105 133802761832256 pyconfig.py:432] Config param merge_gating_gmm: False I0421 09:00:54.549120 133802761832256 pyconfig.py:432] Config param mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive'] I0421 09:00:54.549138 133802761832256 pyconfig.py:432] Config param metrics_dir: None I0421 09:00:54.549152 133802761832256 pyconfig.py:432] Config param metrics_file: I0421 09:00:54.549166 133802761832256 pyconfig.py:432] Config param mhc_expansion_rate: 1 I0421 09:00:54.549181 133802761832256 pyconfig.py:432] Config param micro_batch_size_to_eval_on: 64 I0421 09:00:54.549197 133802761832256 pyconfig.py:432] Config param micro_batch_size_to_train_on: 64 I0421 09:00:54.549212 133802761832256 pyconfig.py:432] Config param mla_kv: RematLocation.REMAT I0421 09:00:54.549228 133802761832256 pyconfig.py:432] Config param mla_naive_kvcache: True I0421 09:00:54.549246 133802761832256 pyconfig.py:432] Config param mla_q: RematLocation.REMAT I0421 09:00:54.549264 133802761832256 pyconfig.py:432] Config param mlp_activations: ['gelu'] I0421 09:00:54.549282 133802761832256 pyconfig.py:432] Config param mlp_activations_limit: -1.0 I0421 09:00:54.549296 133802761832256 pyconfig.py:432] Config param mlp_bias: False I0421 09:00:54.549312 133802761832256 pyconfig.py:432] Config param mlp_dim: 64 I0421 09:00:54.549325 133802761832256 pyconfig.py:432] Config param mlpwi: RematLocation.REMAT I0421 09:00:54.549344 133802761832256 pyconfig.py:432] Config param mlpwi_0: RematLocation.REMAT I0421 09:00:54.549359 133802761832256 pyconfig.py:432] Config param mlpwi_1: RematLocation.REMAT I0421 09:00:54.549373 133802761832256 pyconfig.py:432] Config param mlpwo: RematLocation.REMAT I0421 09:00:54.549389 133802761832256 pyconfig.py:432] Config param moba: False I0421 09:00:54.549403 133802761832256 pyconfig.py:432] Config param moba_chunk_size: 1024 I0421 09:00:54.549417 133802761832256 pyconfig.py:432] Config param moba_topk: 8 I0421 09:00:54.549433 133802761832256 pyconfig.py:432] Config param model_call_mode: I0421 09:00:54.549447 133802761832256 pyconfig.py:432] Config param model_name: gpt3-52k I0421 09:00:54.549462 133802761832256 pyconfig.py:432] Config param moe_expert_input_dim: -1 I0421 09:00:54.549477 133802761832256 pyconfig.py:432] Config param moe_fsdp_use_two_stage_all_gather: False I0421 09:00:54.549491 133802761832256 pyconfig.py:432] Config param moe_mlp_dim: -1 I0421 09:00:54.549505 133802761832256 pyconfig.py:432] Config param moe_mlpwi_0: RematLocation.REMAT I0421 09:00:54.549521 133802761832256 pyconfig.py:432] Config param moe_mlpwi_1: RematLocation.REMAT I0421 09:00:54.549535 133802761832256 pyconfig.py:432] Config param moe_mlpwo: RematLocation.REMAT I0421 09:00:54.549550 133802761832256 pyconfig.py:432] Config param monitor_goodput: False I0421 09:00:54.549568 133802761832256 pyconfig.py:432] Config param monitor_step_time_deviation: True I0421 09:00:54.549583 133802761832256 pyconfig.py:432] Config param mrope_section: [24, 20, 20] I0421 09:00:54.549597 133802761832256 pyconfig.py:432] Config param mscale: 1.0 I0421 09:00:54.549613 133802761832256 pyconfig.py:432] Config param mtc_data_parallelism: 0 I0421 09:00:54.549627 133802761832256 pyconfig.py:432] Config param mtp_eval_target_module: 0 I0421 09:00:54.549652 133802761832256 pyconfig.py:432] Config param mtp_loss_scaling_factor: 0.1 I0421 09:00:54.549668 133802761832256 pyconfig.py:432] Config param mtp_num_layers: 0 I0421 09:00:54.549682 133802761832256 pyconfig.py:432] Config param mu_dtype: float32 I0421 09:00:54.549705 133802761832256 pyconfig.py:432] Config param multi_sampling: False I0421 09:00:54.549719 133802761832256 pyconfig.py:432] Config param multi_tier_checkpointing_backup_interval_minutes: 0 I0421 09:00:54.549734 133802761832256 pyconfig.py:432] Config param muon_beta: 0.95 I0421 09:00:54.549749 133802761832256 pyconfig.py:432] Config param muon_consistent_rms: None I0421 09:00:54.549765 133802761832256 pyconfig.py:432] Config param muon_weight_decay: 0.0 I0421 09:00:54.549780 133802761832256 pyconfig.py:432] Config param n_routing_groups: -1 I0421 09:00:54.549795 133802761832256 pyconfig.py:432] Config param n_window_for_audio: 50 I0421 09:00:54.549809 133802761832256 pyconfig.py:432] Config param n_window_infer_for_audio: 800 I0421 09:00:54.549823 133802761832256 pyconfig.py:432] Config param nope_layer_interval: -1 I0421 09:00:54.549838 133802761832256 pyconfig.py:432] Config param norm_topk_prob: False I0421 09:00:54.549853 133802761832256 pyconfig.py:432] Config param normalization_layer_epsilon: 1e-05 I0421 09:00:54.549871 133802761832256 pyconfig.py:432] Config param normalize_embedding_logits: False I0421 09:00:54.549885 133802761832256 pyconfig.py:432] Config param num_attention_heads_for_vit: 16 I0421 09:00:54.549901 133802761832256 pyconfig.py:432] Config param num_batches: 4 I0421 09:00:54.549916 133802761832256 pyconfig.py:432] Config param num_channels_for_vit: 3 I0421 09:00:54.549930 133802761832256 pyconfig.py:432] Config param num_conv_layers_for_audio: 3 I0421 09:00:54.549945 133802761832256 pyconfig.py:432] Config param num_decoder_layers: 1 I0421 09:00:54.549959 133802761832256 pyconfig.py:432] Config param num_diloco_replicas: 1 I0421 09:00:54.549974 133802761832256 pyconfig.py:432] Config param num_epoch: 1 I0421 09:00:54.549988 133802761832256 pyconfig.py:432] Config param num_eval_passes: 1 I0421 09:00:54.550004 133802761832256 pyconfig.py:432] Config param num_experts: 1 I0421 09:00:54.550018 133802761832256 pyconfig.py:432] Config param num_experts_per_tok: 1 I0421 09:00:54.550033 133802761832256 pyconfig.py:432] Config param num_generations: 2 I0421 09:00:54.550047 133802761832256 pyconfig.py:432] Config param num_hidden_layers_for_vit: 34 I0421 09:00:54.550063 133802761832256 pyconfig.py:432] Config param num_iterations: 1 I0421 09:00:54.550077 133802761832256 pyconfig.py:432] Config param num_kv_heads: 2 I0421 09:00:54.550093 133802761832256 pyconfig.py:432] Config param num_layers_per_pipeline_stage: 1 I0421 09:00:54.550108 133802761832256 pyconfig.py:432] Config param num_mel_bins_for_audio: 128 I0421 09:00:54.550122 133802761832256 pyconfig.py:432] Config param num_pipeline_microbatches: -1 I0421 09:00:54.550137 133802761832256 pyconfig.py:432] Config param num_pipeline_repeats: -1 I0421 09:00:54.550151 133802761832256 pyconfig.py:432] Config param num_position_embeddings_for_vit: 1024 I0421 09:00:54.550167 133802761832256 pyconfig.py:432] Config param num_query_heads: 2 I0421 09:00:54.550181 133802761832256 pyconfig.py:432] Config param num_samplers_slices: -1 I0421 09:00:54.550196 133802761832256 pyconfig.py:432] Config param num_slices: 1 I0421 09:00:54.550211 133802761832256 pyconfig.py:432] Config param num_target_devices: 32 I0421 09:00:54.550226 133802761832256 pyconfig.py:432] Config param num_test_batches: 5 I0421 09:00:54.550240 133802761832256 pyconfig.py:432] Config param num_trainer_slices: -1 I0421 09:00:54.550255 133802761832256 pyconfig.py:432] Config param num_vocab_tiling: 1 I0421 09:00:54.550269 133802761832256 pyconfig.py:432] Config param off_policy_steps: 0 I0421 09:00:54.550285 133802761832256 pyconfig.py:432] Config param offline_data_dir: None I0421 09:00:54.550299 133802761832256 pyconfig.py:432] Config param opt_type: OptimizerType.ADAM_PAX I0421 09:00:54.550320 133802761832256 pyconfig.py:432] Config param optimize_mesh_for_tpu_v6e: False I0421 09:00:54.550334 133802761832256 pyconfig.py:432] Config param optimizer_memory_host_offload: False I0421 09:00:54.550348 133802761832256 pyconfig.py:432] Config param original_max_position_embeddings: 4096 I0421 09:00:54.550362 133802761832256 pyconfig.py:432] Config param out_hidden_size_for_vit: 512 I0421 09:00:54.550375 133802761832256 pyconfig.py:432] Config param out_proj: RematLocation.REMAT I0421 09:00:54.550391 133802761832256 pyconfig.py:432] Config param output_dim_for_audio: 512 I0421 09:00:54.550405 133802761832256 pyconfig.py:432] Config param override_logical_axis_rules: False I0421 09:00:54.550420 133802761832256 pyconfig.py:432] Config param override_model_config: True I0421 09:00:54.550435 133802761832256 pyconfig.py:432] Config param packing: True I0421 09:00:54.550451 133802761832256 pyconfig.py:432] Config param pagedattn_head_dim_alignment: 128 I0421 09:00:54.550464 133802761832256 pyconfig.py:432] Config param pagedattn_max_pages_per_group: -1 I0421 09:00:54.550478 133802761832256 pyconfig.py:432] Config param pagedattn_num_pages: 64 I0421 09:00:54.550494 133802761832256 pyconfig.py:432] Config param pagedattn_pages_per_compute_block: 4 I0421 09:00:54.550508 133802761832256 pyconfig.py:432] Config param pagedattn_tokens_per_page: 32 I0421 09:00:54.550531 133802761832256 pyconfig.py:432] Config param param_scan_axis: 1 I0421 09:00:54.550544 133802761832256 pyconfig.py:432] Config param parameter_memory_host_offload: False I0421 09:00:54.550563 133802761832256 pyconfig.py:432] Config param partial_rotary_factor: 1.0 I0421 09:00:54.550579 133802761832256 pyconfig.py:432] Config param patch_size_for_vit: 14 I0421 09:00:54.550594 133802761832256 pyconfig.py:432] Config param penalty_incorrect_answer: -1.0 I0421 09:00:54.550609 133802761832256 pyconfig.py:432] Config param penalty_incorrect_format: -0.5 I0421 09:00:54.550624 133802761832256 pyconfig.py:432] Config param per_device_batch_size: 2 I0421 09:00:54.550649 133802761832256 pyconfig.py:432] Config param per_device_batch_size_increment: 2.0 I0421 09:00:54.550665 133802761832256 pyconfig.py:432] Config param per_device_batch_size_start: 4.0 I0421 09:00:54.550680 133802761832256 pyconfig.py:432] Config param pipeline_delay_activation_forwarding: False I0421 09:00:54.550695 133802761832256 pyconfig.py:432] Config param pipeline_fsdp_ag_once: False I0421 09:00:54.550709 133802761832256 pyconfig.py:432] Config param pipeline_fsdp_ag_per_repeat: False I0421 09:00:54.550725 133802761832256 pyconfig.py:432] Config param pipeline_parallel_layers: 1 I0421 09:00:54.550739 133802761832256 pyconfig.py:432] Config param pixel_shuffle_ratio_for_vit: 0.5 I0421 09:00:54.550754 133802761832256 pyconfig.py:432] Config param posemb_type_for_vit: learn I0421 09:00:54.550768 133802761832256 pyconfig.py:432] Config param position_id_per_seconds: 25 I0421 09:00:54.550782 133802761832256 pyconfig.py:432] Config param prefill_cache_axis_order: 1,2,0,3 I0421 09:00:54.550798 133802761832256 pyconfig.py:432] Config param prefill_cache_dir: I0421 09:00:54.550814 133802761832256 pyconfig.py:432] Config param prefill_chunk_size: 256 I0421 09:00:54.550828 133802761832256 pyconfig.py:432] Config param prefill_slice: v5e-16 I0421 09:00:54.550843 133802761832256 pyconfig.py:432] Config param prefix_caching_dram_byte: 100000000000 I0421 09:00:54.550857 133802761832256 pyconfig.py:432] Config param prefix_caching_hbm_byte: 10000000000 I0421 09:00:54.550872 133802761832256 pyconfig.py:432] Config param profile_cleanly: True I0421 09:00:54.550886 133802761832256 pyconfig.py:432] Config param profile_periodically_period: -1 I0421 09:00:54.550901 133802761832256 pyconfig.py:432] Config param profile_power_events: False I0421 09:00:54.550915 133802761832256 pyconfig.py:432] Config param profiler: ProfilerType.NONE I0421 09:00:54.550932 133802761832256 pyconfig.py:432] Config param profiler_steps: 5 I0421 09:00:54.550946 133802761832256 pyconfig.py:432] Config param projector_dropout_for_vit: 0.0 I0421 09:00:54.550961 133802761832256 pyconfig.py:432] Config param projector_input_dim_for_vit: 4096 I0421 09:00:54.550975 133802761832256 pyconfig.py:432] Config param projector_output_dim_for_vit: 4096 I0421 09:00:54.550990 133802761832256 pyconfig.py:432] Config param prometheus_port: 0 I0421 09:00:54.551004 133802761832256 pyconfig.py:432] Config param prompt: I love to I0421 09:00:54.551019 133802761832256 pyconfig.py:432] Config param pure_nnx: False I0421 09:00:54.551033 133802761832256 pyconfig.py:432] Config param pure_nnx_decoder: False I0421 09:00:54.551048 133802761832256 pyconfig.py:432] Config param q_lora_rank: 0 I0421 09:00:54.551063 133802761832256 pyconfig.py:432] Config param qk_clip_threshold: 100.0 I0421 09:00:54.551077 133802761832256 pyconfig.py:432] Config param qk_nope_head_dim: 128 I0421 09:00:54.551092 133802761832256 pyconfig.py:432] Config param qk_norm_with_scale: True I0421 09:00:54.551107 133802761832256 pyconfig.py:432] Config param qk_rope_head_dim: 64 I0421 09:00:54.551122 133802761832256 pyconfig.py:432] Config param qkv_proj: RematLocation.REMAT I0421 09:00:54.551136 133802761832256 pyconfig.py:432] Config param quant_cfg_path: I0421 09:00:54.551152 133802761832256 pyconfig.py:432] Config param quantization: QuantizationType.NONE I0421 09:00:54.551170 133802761832256 pyconfig.py:432] Config param quantization_local_shard_count: 4 I0421 09:00:54.551184 133802761832256 pyconfig.py:432] Config param quantize_kvcache: False I0421 09:00:54.551198 133802761832256 pyconfig.py:432] Config param query_proj: RematLocation.REMAT I0421 09:00:54.551214 133802761832256 pyconfig.py:432] Config param query_wa_proj: RematLocation.REMAT I0421 09:00:54.551230 133802761832256 pyconfig.py:432] Config param ragged_block_size: 256 I0421 09:00:54.551244 133802761832256 pyconfig.py:432] Config param ragged_buffer_factor: -1.0 I0421 09:00:54.551259 133802761832256 pyconfig.py:432] Config param rampup_end_step: 0 I0421 09:00:54.551274 133802761832256 pyconfig.py:432] Config param rampup_samples_per_increment_to_load: None I0421 09:00:54.551290 133802761832256 pyconfig.py:432] Config param reasoning_end_token: </reasoning> I0421 09:00:54.551305 133802761832256 pyconfig.py:432] Config param reasoning_start_token: <reasoning> I0421 09:00:54.551320 133802761832256 pyconfig.py:432] Config param record_internal_nn_metrics: 0 I0421 09:00:54.551335 133802761832256 pyconfig.py:432] Config param remat_policy: full I0421 09:00:54.551350 133802761832256 pyconfig.py:432] Config param remat_policy_for_vit: minimal I0421 09:00:54.551365 133802761832256 pyconfig.py:432] Config param remove_size_one_mesh_axis_from_type: True I0421 09:00:54.551379 133802761832256 pyconfig.py:432] Config param replicate_quant_scale: False I0421 09:00:54.551394 133802761832256 pyconfig.py:432] Config param replicator_backup_interval_minutes: 0 I0421 09:00:54.551409 133802761832256 pyconfig.py:432] Config param report_heartbeat_metric_for_gcp_monitoring: False I0421 09:00:54.551424 133802761832256 pyconfig.py:432] Config param report_performance_metric_for_gcp_monitoring: False I0421 09:00:54.551439 133802761832256 pyconfig.py:432] Config param reshape_q: False I0421 09:00:54.551454 133802761832256 pyconfig.py:432] Config param return_log_prob: False I0421 09:00:54.551468 133802761832256 pyconfig.py:432] Config param reuse_example_batch: 0 I0421 09:00:54.551482 133802761832256 pyconfig.py:432] Config param reward_exact_answer: 5.0 I0421 09:00:54.551497 133802761832256 pyconfig.py:432] Config param reward_exact_format_match: 3.0 I0421 09:00:54.551512 133802761832256 pyconfig.py:432] Config param reward_partial_format_match: 0.5 I0421 09:00:54.551527 133802761832256 pyconfig.py:432] Config param reward_ratio_guess_to_answer_high: 0.5 I0421 09:00:54.551543 133802761832256 pyconfig.py:432] Config param reward_ratio_guess_to_answer_low: 0.25 I0421 09:00:54.551560 133802761832256 pyconfig.py:432] Config param reward_white_space_format_match: 1.5 I0421 09:00:54.551577 133802761832256 pyconfig.py:432] Config param rl: {'num_generations': 2, 'num_iterations': 1, 'grpo_beta': 0.08, 'grpo_epsilon': 0.2, 'loss_algo': 'grpo', 'use_agentic_rollout': False, 'max_concurrency': 256, 'off_policy_steps': 0, 'system_prompt': '', 'degenerate_group_masking': True, 'epsilon_high': None} I0421 09:00:54.551596 133802761832256 pyconfig.py:432] Config param rollout_data_parallelism: -1 I0421 09:00:54.551611 133802761832256 pyconfig.py:432] Config param rollout_expert_parallelism: 1 I0421 09:00:54.551626 133802761832256 pyconfig.py:432] Config param rollout_micro_batch_size: -1 I0421 09:00:54.551652 133802761832256 pyconfig.py:432] Config param rollout_tensor_parallelism: -1 I0421 09:00:54.551667 133802761832256 pyconfig.py:432] Config param rope_attention_scaling: False I0421 09:00:54.551683 133802761832256 pyconfig.py:432] Config param rope_factor: 40 I0421 09:00:54.551697 133802761832256 pyconfig.py:432] Config param rope_interleave: True I0421 09:00:54.551712 133802761832256 pyconfig.py:432] Config param rope_linear_scaling_factor: 1.0 I0421 09:00:54.551726 133802761832256 pyconfig.py:432] Config param rope_max_timescale: 10000 I0421 09:00:54.551741 133802761832256 pyconfig.py:432] Config param rope_min_timescale: 1 I0421 09:00:54.551755 133802761832256 pyconfig.py:432] Config param rope_theta_for_vit: 10000 I0421 09:00:54.551770 133802761832256 pyconfig.py:432] Config param rope_truncate: True I0421 09:00:54.551785 133802761832256 pyconfig.py:432] Config param rope_type: RopeType.DEFAULT I0421 09:00:54.551801 133802761832256 pyconfig.py:432] Config param rope_use_scale: True I0421 09:00:54.551817 133802761832256 pyconfig.py:432] Config param routed_bias: False I0421 09:00:54.551831 133802761832256 pyconfig.py:432] Config param routed_bias_update_rate: 0.0 I0421 09:00:54.551846 133802761832256 pyconfig.py:432] Config param routed_scaling_factor: 1.0 I0421 09:00:54.551861 133802761832256 pyconfig.py:432] Config param routed_score_func: I0421 09:00:54.551876 133802761832256 pyconfig.py:432] Config param run_name: gpt3-52k_2026-04-21-09-00 I0421 09:00:54.551890 133802761832256 pyconfig.py:432] Config param sa_block_kv: 512 I0421 09:00:54.551904 133802761832256 pyconfig.py:432] Config param sa_block_kv_compute: 512 I0421 09:00:54.551920 133802761832256 pyconfig.py:432] Config param sa_block_kv_dkv: 512 I0421 09:00:54.551933 133802761832256 pyconfig.py:432] Config param sa_block_kv_dkv_compute: 512 I0421 09:00:54.551948 133802761832256 pyconfig.py:432] Config param sa_block_kv_dq: 512 I0421 09:00:54.551962 133802761832256 pyconfig.py:432] Config param sa_block_q: 512 I0421 09:00:54.551977 133802761832256 pyconfig.py:432] Config param sa_block_q_dkv: 512 I0421 09:00:54.551991 133802761832256 pyconfig.py:432] Config param sa_block_q_dq: 512 I0421 09:00:54.552006 133802761832256 pyconfig.py:432] Config param sa_k_layout: HEAD_DIM_MINOR I0421 09:00:54.552021 133802761832256 pyconfig.py:432] Config param sa_q_layout: HEAD_DIM_MINOR I0421 09:00:54.552036 133802761832256 pyconfig.py:432] Config param sa_use_fused_bwd_kernel: False I0421 09:00:54.552050 133802761832256 pyconfig.py:432] Config param sa_v_layout: HEAD_DIM_MINOR I0421 09:00:54.552066 133802761832256 pyconfig.py:432] Config param sampler_devices_fraction: 0.5 I0421 09:00:54.552080 133802761832256 pyconfig.py:432] Config param save_checkpoint_on_completion: True I0421 09:00:54.552095 133802761832256 pyconfig.py:432] Config param save_config_to_gcs: False I0421 09:00:54.552109 133802761832256 pyconfig.py:432] Config param save_quantized_params_path: I0421 09:00:54.552125 133802761832256 pyconfig.py:432] Config param scale_embedding_for_audio: True I0421 09:00:54.552139 133802761832256 pyconfig.py:432] Config param scan_layers: True I0421 09:00:54.552154 133802761832256 pyconfig.py:432] Config param scan_layers_per_stage: False I0421 09:00:54.552168 133802761832256 pyconfig.py:432] Config param scan_pipeline_iterations: True I0421 09:00:54.552183 133802761832256 pyconfig.py:432] Config param scan_pipeline_repeats: False I0421 09:00:54.552197 133802761832256 pyconfig.py:432] Config param set_remat_policy_on_layers_per_stage: False I0421 09:00:54.552212 133802761832256 pyconfig.py:432] Config param set_remat_policy_on_pipeline_iterations: True I0421 09:00:54.552226 133802761832256 pyconfig.py:432] Config param sft_train_on_completion_only: False I0421 09:00:54.552241 133802761832256 pyconfig.py:432] Config param shard_exp_on_fsdp: False I0421 09:00:54.552254 133802761832256 pyconfig.py:432] Config param shard_mode: ShardMode.AUTO I0421 09:00:54.552271 133802761832256 pyconfig.py:432] Config param shard_optimizer_over_data: False I0421 09:00:54.552285 133802761832256 pyconfig.py:432] Config param sharding_strategy: None I0421 09:00:54.552301 133802761832256 pyconfig.py:432] Config param sharding_tolerance: 0.02 I0421 09:00:54.552314 133802761832256 pyconfig.py:432] Config param shardy: True I0421 09:00:54.552330 133802761832256 pyconfig.py:432] Config param share_kv_projections: False I0421 09:00:54.552344 133802761832256 pyconfig.py:432] Config param shared_experts: 0 I0421 09:00:54.552359 133802761832256 pyconfig.py:432] Config param sinkhorn_iterations: 20 I0421 09:00:54.552373 133802761832256 pyconfig.py:432] Config param skip_first_n_steps_for_profiler: 1 I0421 09:00:54.552388 133802761832256 pyconfig.py:432] Config param skip_jax_distributed_system: False I0421 09:00:54.552402 133802761832256 pyconfig.py:432] Config param skip_step_interval: 128 I0421 09:00:54.552417 133802761832256 pyconfig.py:432] Config param skip_step_on_spikes: False I0421 09:00:54.552431 133802761832256 pyconfig.py:432] Config param skip_step_scaling_factor: 6.0 I0421 09:00:54.552446 133802761832256 pyconfig.py:432] Config param sliding_window_size: 0 I0421 09:00:54.552460 133802761832256 pyconfig.py:432] Config param solution_end_token: </answer> I0421 09:00:54.552475 133802761832256 pyconfig.py:432] Config param solution_start_token: <answer> I0421 09:00:54.552489 133802761832256 pyconfig.py:432] Config param source_checkpoint_layout: orbax I0421 09:00:54.552505 133802761832256 pyconfig.py:432] Config param sparse_matmul: True I0421 09:00:54.552518 133802761832256 pyconfig.py:432] Config param spatial_merge_size_for_vit: 2 I0421 09:00:54.552534 133802761832256 pyconfig.py:432] Config param stack_prefill_result_cache: False I0421 09:00:54.552548 133802761832256 pyconfig.py:432] Config param stack_trace_interval_seconds: 600 I0421 09:00:54.552566 133802761832256 pyconfig.py:432] Config param stack_trace_to_cloud: False I0421 09:00:54.552580 133802761832256 pyconfig.py:432] Config param step_deviation_interval_seconds: 30 I0421 09:00:54.552595 133802761832256 pyconfig.py:432] Config param steps: 200000 I0421 09:00:54.552610 133802761832256 pyconfig.py:432] Config param stop_strings: None I0421 09:00:54.552626 133802761832256 pyconfig.py:432] Config param student_overrides: {'model_name': 'llama3.1-8b'} I0421 09:00:54.552652 133802761832256 pyconfig.py:432] Config param student_params_to_update: None I0421 09:00:54.552669 133802761832256 pyconfig.py:432] Config param subslice_shape: I0421 09:00:54.552683 133802761832256 pyconfig.py:432] Config param swap_space_vllm_gb: 2 I0421 09:00:54.552699 133802761832256 pyconfig.py:432] Config param system_prompt: I0421 09:00:54.552715 133802761832256 pyconfig.py:432] Config param target_eval_loss: 0.0 I0421 09:00:54.552731 133802761832256 pyconfig.py:432] Config param teacher_overrides: {'model_name': 'llama3.1-8b'} I0421 09:00:54.552748 133802761832256 pyconfig.py:432] Config param temperature_tuning: False I0421 09:00:54.552763 133802761832256 pyconfig.py:432] Config param temporal_patch_size_for_vit: 2 I0421 09:00:54.552778 133802761832256 pyconfig.py:432] Config param tensorboard_dir: None I0421 09:00:54.552794 133802761832256 pyconfig.py:432] Config param tensors_on_device: None I0421 09:00:54.552809 133802761832256 pyconfig.py:432] Config param tensors_to_offload: None I0421 09:00:54.552823 133802761832256 pyconfig.py:432] Config param test_batch_start_index: 0 I0421 09:00:54.552837 133802761832256 pyconfig.py:432] Config param tile_size_for_vit: 336 I0421 09:00:54.552853 133802761832256 pyconfig.py:432] Config param tokenize_eval_data: True I0421 09:00:54.552867 133802761832256 pyconfig.py:432] Config param tokenize_train_data: True I0421 09:00:54.552883 133802761832256 pyconfig.py:432] Config param tokenizer_path: meta-llama/Llama-3.1-8B I0421 09:00:54.552898 133802761832256 pyconfig.py:432] Config param tokenizer_type: TokenizerType.HUGGINGFACE I0421 09:00:54.552916 133802761832256 pyconfig.py:432] Config param topk_routing_group: -1 I0421 09:00:54.552930 133802761832256 pyconfig.py:432] Config param train_data_columns: ['text'] I0421 09:00:54.552945 133802761832256 pyconfig.py:432] Config param train_fraction: 1.0 I0421 09:00:54.552961 133802761832256 pyconfig.py:432] Config param train_image_column: image I0421 09:00:54.552975 133802761832256 pyconfig.py:432] Config param train_micro_batch_size: -1 I0421 09:00:54.552989 133802761832256 pyconfig.py:432] Config param train_split: train I0421 09:00:54.553004 133802761832256 pyconfig.py:432] Config param trainable_parameters_mask: [] I0421 09:00:54.553018 133802761832256 pyconfig.py:432] Config param trainable_position_size: 2048 I0421 09:00:54.553033 133802761832256 pyconfig.py:432] Config param trainer_devices_fraction: 0.5 I0421 09:00:54.553048 133802761832256 pyconfig.py:432] Config param upload_all_profiler_results: False I0421 09:00:54.553063 133802761832256 pyconfig.py:432] Config param use_2d_fsdp_sharding: False I0421 09:00:54.553079 133802761832256 pyconfig.py:432] Config param use_agentic_rollout: False I0421 09:00:54.553093 133802761832256 pyconfig.py:432] Config param use_audio: False I0421 09:00:54.553108 133802761832256 pyconfig.py:432] Config param use_audio_in_video: False I0421 09:00:54.553122 133802761832256 pyconfig.py:432] Config param use_batch_split_schedule: False I0421 09:00:54.553137 133802761832256 pyconfig.py:432] Config param use_chat_template: False I0421 09:00:54.553152 133802761832256 pyconfig.py:432] Config param use_chunked_prefill: False I0421 09:00:54.553167 133802761832256 pyconfig.py:432] Config param use_custom_sort_vjp: True I0421 09:00:54.553180 133802761832256 pyconfig.py:432] Config param use_dpo: False I0421 09:00:54.553196 133802761832256 pyconfig.py:432] Config param use_gather_mosaic_kernel: False I0421 09:00:54.553210 133802761832256 pyconfig.py:432] Config param use_grpo: True I0421 09:00:54.553224 133802761832256 pyconfig.py:432] Config param use_indexer: False I0421 09:00:54.553240 133802761832256 pyconfig.py:432] Config param use_iota_embed: True I0421 09:00:54.553254 133802761832256 pyconfig.py:432] Config param use_jax_splash: False I0421 09:00:54.553268 133802761832256 pyconfig.py:432] Config param use_max_logit_estimate: -1 I0421 09:00:54.553282 133802761832256 pyconfig.py:432] Config param use_mrope: False I0421 09:00:54.553297 133802761832256 pyconfig.py:432] Config param use_multimodal: False I0421 09:00:54.553312 133802761832256 pyconfig.py:432] Config param use_pathways: True I0421 09:00:54.553326 133802761832256 pyconfig.py:432] Config param use_post_attn_norm: False I0421 09:00:54.553341 133802761832256 pyconfig.py:432] Config param use_post_ffw_norm: False I0421 09:00:54.553356 133802761832256 pyconfig.py:432] Config param use_qk_clip: False I0421 09:00:54.553370 133802761832256 pyconfig.py:432] Config param use_qk_norm: False I0421 09:00:54.553385 133802761832256 pyconfig.py:432] Config param use_qk_norm_in_gdn: True I0421 09:00:54.553400 133802761832256 pyconfig.py:432] Config param use_qwix_quantization: False I0421 09:00:54.553416 133802761832256 pyconfig.py:432] Config param use_ragged_attention: False I0421 09:00:54.553431 133802761832256 pyconfig.py:432] Config param use_random_routing: False I0421 09:00:54.553445 133802761832256 pyconfig.py:432] Config param use_replicator_service: False I0421 09:00:54.553460 133802761832256 pyconfig.py:432] Config param use_ring_of_experts: False I0421 09:00:54.553475 133802761832256 pyconfig.py:432] Config param use_sft: False I0421 09:00:54.553489 133802761832256 pyconfig.py:432] Config param use_splash_scheduler: False I0421 09:00:54.553504 133802761832256 pyconfig.py:432] Config param use_tokamax_gmm: False I0421 09:00:54.553518 133802761832256 pyconfig.py:432] Config param use_tokamax_splash: False I0421 09:00:54.553531 133802761832256 pyconfig.py:432] Config param use_truncation: True I0421 09:00:54.553546 133802761832256 pyconfig.py:432] Config param use_tunix_gradient_accumulation: False I0421 09:00:54.553564 133802761832256 pyconfig.py:432] Config param use_untrainable_positional_embedding: False I0421 09:00:54.553579 133802761832256 pyconfig.py:432] Config param use_vertex_tensorboard: False I0421 09:00:54.553593 133802761832256 pyconfig.py:432] Config param using_pipeline_parallelism: False I0421 09:00:54.553609 133802761832256 pyconfig.py:432] Config param v_head_dim: 128 I0421 09:00:54.553622 133802761832256 pyconfig.py:432] Config param v_norm_with_scale: True I0421 09:00:54.553645 133802761832256 pyconfig.py:432] Config param value_proj: RematLocation.REMAT I0421 09:00:54.553662 133802761832256 pyconfig.py:432] Config param vertex_tensorboard_project: I0421 09:00:54.553676 133802761832256 pyconfig.py:432] Config param vertex_tensorboard_region: I0421 09:00:54.553690 133802761832256 pyconfig.py:432] Config param video_path: I0421 09:00:54.553705 133802761832256 pyconfig.py:432] Config param video_placeholder: <|video|> I0421 09:00:54.553719 133802761832256 pyconfig.py:432] Config param vision_output_dim_for_vit: 4096 I0421 09:00:54.553733 133802761832256 pyconfig.py:432] Config param vision_output_length: -1 I0421 09:00:54.553749 133802761832256 pyconfig.py:432] Config param vllm_additional_config: {} I0421 09:00:54.553763 133802761832256 pyconfig.py:432] Config param vllm_hf_config_path: I0421 09:00:54.553778 133802761832256 pyconfig.py:432] Config param vllm_hf_overrides: {} I0421 09:00:54.553792 133802761832256 pyconfig.py:432] Config param vocab_size: 32000 I0421 09:00:54.553807 133802761832256 pyconfig.py:432] Config param warmup_steps_fraction: 0.1 I0421 09:00:54.553821 133802761832256 pyconfig.py:432] Config param weight_dtype: float32 I0421 09:00:54.553844 133802761832256 pyconfig.py:432] Config param weight_quantization_calibration_method: absmax I0421 09:00:54.553858 133802761832256 pyconfig.py:432] Config param wi_tile_dlhs_batch_seq: 512 I0421 09:00:54.553874 133802761832256 pyconfig.py:432] Config param wi_tile_dlhs_embed_dim: 1024 I0421 09:00:54.553888 133802761832256 pyconfig.py:432] Config param wi_tile_dlhs_mlp_dim: 1024 I0421 09:00:54.553903 133802761832256 pyconfig.py:432] Config param wi_tile_drhs_batch_seq: 512 I0421 09:00:54.553916 133802761832256 pyconfig.py:432] Config param wi_tile_drhs_embed_dim: 1024 I0421 09:00:54.553932 133802761832256 pyconfig.py:432] Config param wi_tile_drhs_mlp_dim: 1024 I0421 09:00:54.553945 133802761832256 pyconfig.py:432] Config param wi_tile_fwd_batch_seq: 512 I0421 09:00:54.553961 133802761832256 pyconfig.py:432] Config param wi_tile_fwd_embed_dim: 1024 I0421 09:00:54.553975 133802761832256 pyconfig.py:432] Config param wi_tile_fwd_mlp_dim: 1024 I0421 09:00:54.553990 133802761832256 pyconfig.py:432] Config param wo_tile_dlhs_batch_seq: 512 I0421 09:00:54.554004 133802761832256 pyconfig.py:432] Config param wo_tile_dlhs_embed_dim: 1024 I0421 09:00:54.554020 133802761832256 pyconfig.py:432] Config param wo_tile_dlhs_mlp_dim: 1024 I0421 09:00:54.554034 133802761832256 pyconfig.py:432] Config param wo_tile_drhs_batch_seq: 512 I0421 09:00:54.554049 133802761832256 pyconfig.py:432] Config param wo_tile_drhs_embed_dim: 1024 I0421 09:00:54.554063 133802761832256 pyconfig.py:432] Config param wo_tile_drhs_mlp_dim: 1024 I0421 09:00:54.554078 133802761832256 pyconfig.py:432] Config param wo_tile_fwd_batch_seq: 512 I0421 09:00:54.554092 133802761832256 pyconfig.py:432] Config param wo_tile_fwd_embed_dim: 1024 I0421 09:00:54.554108 133802761832256 pyconfig.py:432] Config param wo_tile_fwd_mlp_dim: 1024 I0421 09:00:54.554124 133802761832256 pyconfig.py:432] Config param wsd_decay_steps_fraction: 0.1 I0421 09:00:54.554142 133802761832256 pyconfig.py:432] Config param wsd_decay_style: WsdDecayStyle.LINEAR I0421 09:00:54.554160 133802761832256 pyconfig.py:432] Config param xprof_e2e_enable_fw_power_level_event: False I0421 09:00:54.554174 133802761832256 pyconfig.py:432] Config param xprof_e2e_enable_fw_thermal_event: False I0421 09:00:54.554190 133802761832256 pyconfig.py:432] Config param xprof_e2e_enable_fw_throttle_event: False I0421 09:00:54.554204 133802761832256 pyconfig.py:432] Config param xprof_tpu_power_trace_level: 0 I0421 09:00:54.554223 133802761832256 pyconfig.py:432] Config param z_loss_multiplier: 0.0 I0421 09:00:54.554541 133802761832256 tokenizer.py:245] Tokenizer path: meta-llama/Llama-2-7b-chat-hf I0421 09:00:54.554581 133802761832256 tokenizer.py:224] Loading HF tokenizer: meta-llama/Llama-2-7b-chat-hf I0421 09:00:58.677709 133802761832256 _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`. I0421 09:00:58.680812 133802761832256 maxtext_utils.py:1718] Num_devices: 32, shape (1, 4, 1, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1) I0421 09:00:58.680935 133802761832256 train_distill.py:596] Applying logical axis rules for model initialization and training... I0421 09:00:58.681010 133802761832256 train_distill.py:600] Loading Student from ... I0421 09:00:58.681038 133802761832256 train_distill.py:169] --- Student Configuration --- I0421 09:00:58.681061 133802761832256 train_distill.py:170] Model Name: gpt3-52k I0421 09:00:58.681083 133802761832256 train_distill.py:171] Dimensions: 1 Layers, 16 Emb Dim, 8 Head Dim I0421 09:00:58.681101 133802761832256 train_distill.py:174] Attention Heads: 2 Query, 2 KV I0421 09:00:58.681119 133802761832256 train_distill.py:175] Vocab Size: 32000 I0421 09:00:58.681136 133802761832256 train_distill.py:176] Checkpoint: I0421 09:00:58.681153 133802761832256 train_distill.py:465] Initializing model: gpt3-52k... I0421 09:00:59.947936 133802761832256 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... I0421 09:00:59.948042 133802761832256 train_distill.py:169] --- Teacher Configuration --- I0421 09:00:59.948070 133802761832256 train_distill.py:170] Model Name: gpt3-52k I0421 09:00:59.948092 133802761832256 train_distill.py:171] Dimensions: 1 Layers, 16 Emb Dim, 8 Head Dim I0421 09:00:59.948117 133802761832256 train_distill.py:174] Attention Heads: 2 Query, 2 KV I0421 09:00:59.948133 133802761832256 train_distill.py:175] Vocab Size: 32000 I0421 09:00:59.948150 133802761832256 train_distill.py:176] Checkpoint: gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items I0421 09:00:59.948168 133802761832256 train_distill.py:465] Initializing model: gpt3-52k... I0421 09:01:01.095861 133802761832256 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0421 09:01:01.096292 133802761832256 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 0x79b0aa928140>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0421 09:01:01.096360 133802761832256 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.28 W0421 09:01:02.056916 133802761832256 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 I0421 09:01:02.584193 2146 google_auth_provider.cc:181] Running on GCE, using service account 562977990677-compute@developer.gserviceaccount.com I0421 09:01:03.735723 133802761832256 checkpointer.py:304] Restoring checkpoint from gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items. W0421 09:01:05.785310 133802761832256 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. I0421 09:01:05.785710 133802761832256 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 I0421 09:01:06.789961 133802761832256 checkpointer.py:318] Finished restoring checkpoint in 3.44 seconds from gs://lance-maxtext/pt_seed_ckpts/pt_seed_ckpts/pt_seed_ckpt_gpt352k_v32k_linen/checkpoints/4/items. I0421 09:01:07.479608 133802761832256 train_distill.py:640] Initializing Data Iterators via MaxText pipeline... I0421 09:01:07.542737 133802761832256 config.py:112] TensorFlow version 2.20.0 available. I0421 09:01:07.543216 133802761832256 config.py:125] JAX version 0.8.3 available. E0421 09:01:09.845675 133802761832256 packing.py:209] PackAndBatchOperation is deprecated. Please use lazy_dataset.FirstFitPackIterDataset instead. I0421 09:01:09.845891 133802761832256 data_loader.py:408] Adding CopyNumPyArrayToSharedMemory MapTransform. I0421 09:01:09.848937 133802761832256 train_distill.py:410] Input Pipeline Checkpointing: DISABLED I0421 09:01:09.848997 133802761832256 train_distill.py:414] Reason: Iterator 'MultiHostDataLoadIterator' is not recognized as Grain (dataset_type='DatasetType.HF', has_save=False) I0421 09:01:09.849060 133802761832256 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0421 09:01:09.849135 133802761832256 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 0x79b0aa928140>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0421 09:01:09.849177 133802761832256 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0421 09:01:09.849209 133802761832256 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 0x79b0aa928140>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0421 09:01:09.849251 133802761832256 checkpoint_manager.py:702] [process=4][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers={'model_params': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f17230>, 'optimizer_state': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x79ac3c38f950>, 'custom_metadata': <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x79a6b0677bc0>}, handler_registry=None I0421 09:01:09.849440 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "model_params". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f17230>` 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`. I0421 09:01:09.849480 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "optimizer_state". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x79ac3c38f950>` 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`. I0421 09:01:09.849505 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "custom_metadata". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x79a6b0677bc0>` 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`. I0421 09:01:09.849528 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x79a6b0e05b80>` for item "metrics" and save args `<class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>` to `_handler_registry`. I0421 09:01:09.849553 133802761832256 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 0x799938f17230>, ('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f17230>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x79ac3c38f950>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x79ac3c38f950>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x79a6b0677bc0>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x79a6b0677bc0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x79a6b0e05b80>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x79a6b0e05b80>}). I0421 09:01:09.849993 133802761832256 async_checkpointer.py:177] [process=4][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7999387dc9a0> timeout: 600 secs and primary_host=0 for async checkpoint writes I0421 09:01:12.322866 133802761832256 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://lance-maxtext/pt_ckpt_xpk_feat_nnx_trainstate_and_training_loop_20260421_083247/pt_distill_nnx_xpk_feat_nnx_trainstate_and_training_loop_20260421_083247_07_distill_smoke/checkpoints I0421 09:01:12.767292 133802761832256 checkpoint_manager.py:921] [process=4][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_feat_nnx_trainstate_and_training_loop_20260421_083247/pt_distill_nnx_xpk_feat_nnx_trainstate_and_training_loop_20260421_083247_07_distill_smoke/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x799938f171a0> I0421 09:01:12.767463 133802761832256 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0421 09:01:12.767531 133802761832256 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 0x79b0aa928140>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0421 09:01:12.767576 133802761832256 pytree_checkpoint_handler.py:577] save_device_host_concurrent_bytes=None I0421 09:01:12.767610 133802761832256 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 0x79b0aa928140>, enable_pinned_host_transfer=False, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB) I0421 09:01:12.767662 133802761832256 checkpoint_manager.py:1983] [process=4][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning. I0421 09:01:12.767719 133802761832256 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=133802761832256 count=1 at 0x79993a899880>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x799938f16fc0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x79a6b0687440>, _write_futures=[]) I0421 09:01:12.768080 133802761832256 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=133802761832256 count=1 at 0x79993a899880>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x799938f16fc0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x79a6b0687440>, _write_futures=[]) I0421 09:01:12.768106 133802761832256 checkpoint.py:459] Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=133802761832256 count=1 at 0x79993a899880>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x799938f16fc0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x79a6b0687440>, _write_futures=[]) I0421 09:01:12.768139 133802761832256 checkpoint_manager.py:702] [process=4][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers={'model_params': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f17170>, 'optimizer_state': <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f16450>, 'custom_metadata': <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x799938f16720>, 'iter': <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x799938f16ae0>}, handler_registry=None I0421 09:01:12.768249 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "model_params". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f17170>` 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`. I0421 09:01:12.768283 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "optimizer_state". Adding handler `<orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f16450>` 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`. I0421 09:01:12.768310 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "custom_metadata". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x799938f16720>` 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`. I0421 09:01:12.768339 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "iter". Adding handler `<maxtext.common.checkpointing.GrainCheckpointHandler object at 0x799938f16ae0>` for item "iter" and save args `<class 'maxtext.common.checkpointing.GrainCheckpointSave'>` and restore args `<class 'maxtext.common.checkpointing.GrainCheckpointRestore'>` to `_handler_registry`. I0421 09:01:12.768363 133802761832256 composite_checkpoint_handler.py:237] Deferred registration for item: "metrics". Adding handler `<orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x799938766270>` for item "metrics" and save args `<class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>` to `_handler_registry`. I0421 09:01:12.768389 133802761832256 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 0x799938f17170>, ('model_params', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f17170>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeSaveArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f16450>, ('optimizer_state', <class 'orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeRestoreArgs'>): <orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler object at 0x799938f16450>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x799938f16720>, ('custom_metadata', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x799938f16720>, ('iter', <class 'maxtext.common.checkpointing.GrainCheckpointSave'>): <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x799938f16ae0>, ('iter', <class 'maxtext.common.checkpointing.GrainCheckpointRestore'>): <maxtext.common.checkpointing.GrainCheckpointHandler object at 0x799938f16ae0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x799938766270>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x799938766270>}). I0421 09:01:12.768461 133802761832256 async_checkpointer.py:177] [process=4][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>._fn at 0x7999387dcae0> timeout: 600 secs and primary_host=0 for async checkpoint writes I0421 09:01:13.569192 133802761832256 checkpoint_manager.py:1788] Found 0 checkpoint steps in gs://lance-maxtext/pt_ckpt_xpk_feat_nnx_trainstate_and_training_loop_20260421_083247/pt_distill_nnx_xpk_feat_nnx_trainstate_and_training_loop_20260421_083247_07_distill_smoke/checkpoints I0421 09:01:13.998022 133802761832256 checkpoint_manager.py:921] [process=4][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_feat_nnx_trainstate_and_training_loop_20260421_083247/pt_distill_nnx_xpk_feat_nnx_trainstate_and_training_loop_20260421_083247_07_distill_smoke/checkpoints: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x79ac3c38ec90> I0421 09:01:13.998604 133802761832256 train_distill.py:691] Starting Distillation Training... I0421 09:01:13.998726 133802761832256 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)) I0421 09:01:14.121504 133802761832256 peft_trainer.py:600] Compiled train_step cache size: 0 Training: 0%| | 0/5 [00:00<?, ?step/s]I0421 09:01:14.123289 133657964570368 grain_pool.py:367] Grain pool will use 1 processes. I0421 09:01:14.149531 133657964570368 grain_pool.py:440] Grain pool will start child processes. I0421 09:01:14.154778 133657964570368 grain_pool.py:448] Grain pool started all child processes. 2026-04-21 09:01:20.177039: E external/local_xla/xla/stream_executor/cuda/cuda_platform.cc:51] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303) I0421 09:01:23.428210 133802761832256 utils.py:86] Train loop finished in: 9.3061 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=9, process_index=2, coords=(1,2,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=24, process_index=6, coords=(0,6,0), core_on_chip=0), TpuDevice(id=18, process_index=5, coords=(2,4,0), core_on_chip=0), TpuDevice(id=3, process_index=1, coords=(3,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=10, process_index=3, coords=(2,2,0), core_on_chip=0), TpuDevice(id=17, process_index=4, coords=(1,4,0), core_on_chip=0), TpuDevice(id=25, process_index=6, coords=(1,6,0), core_on_chip=0), TpuDevice(id=11, process_index=3, coords=(3,2,0), core_on_chip=0), TpuDevice(id=19, process_index=5, coords=(3,4,0), core_on_chip=0), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=26, process_index=7, coords=(2,6,0), core_on_chip=0), TpuDevice(id=31, process_index=7, coords=(3,7,0), core_on_chip=0), TpuDevice(id=7, process_index=1, coords=(3,1,0), core_on_chip=0), TpuDevice(id=12, process_index=2, coords=(0,3,0), core_on_chip=0), TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=20, process_index=4, coords=(0,5,0), core_on_chip=0), TpuDevice(id=28, process_index=6, coords=(0,7,0), core_on_chip=0), TpuDevice(id=8, process_index=2, coords=(0,2,0), core_on_chip=0), TpuDevice(id=2, process_index=1, coords=(2,0,0), core_on_chip=0), TpuDevice(id=14, process_index=3, coords=(2,3,0), core_on_chip=0), TpuDevice(id=22, process_index=5, coords=(2,5,0), core_on_chip=0), TpuDevice(id=27, process_index=7, coords=(3,6,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=30, process_index=7, coords=(2,7,0), core_on_chip=0), TpuDevice(id=21, process_index=4, coords=(1,5,0), core_on_chip=0), TpuDevice(id=15, process_index=3, coords=(3,3,0), core_on_chip=0), TpuDevice(id=23, process_index=5, coords=(3,5,0), core_on_chip=0), TpuDevice(id=29, process_index=6, coords=(1,7,0), core_on_chip=0), TpuDevice(id=13, process_index=2, coords=(1,3,0), core_on_chip=0)} I0421 09:01:23.779594 133657964570368 grain_pool.py:542] Grain pool is exiting. I0421 09:01:23.779724 133657964570368 grain_pool.py:547] Shutting down multiprocessing system. I0421 09:01:25.225557 133657964570368 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: Tue Apr 21 09:01:36 UTC 2026 EXIT_CODE=1