{
 "experiment": {
  "name": "pixeldit_sem1024d24l_pix128d4l_bilineardwgeluproject_unfreeze_3src_const1em5_t128_4stepft",
  "rundir": "",
  "model_type": "pixeldit",
  "resume_from": "",
  "reset_best_valid_loss": true,
  "reset_scheduler_state": true,
  "skip_optimizer_reloading": true
 },
 "dit": {
  "patch_size": 1,
  "patch_size_vert": 1,
  "patch_size_horiz": 4,
  "embed_dim": 1024,
  "n_layers": 24,
  "num_heads": 16,
  "qk_norm": true,
  "qk_norm_elementwise_affine": false,
  "attn_kernel": 9,
  "do_interleaved_dilation": false,
  "na_dilations": 3,
  "na3d_backend": "",
  "gated_attention": true,
  "do_alt_depthwise_attn": true,
  "use_conv_instead_of_na3d": false,
  "conv_kernel_size": 3,
  "do_rope_2d": false,
  "do_rope_2d_stereographic": true,
  "do_rotate_wind": true,
  "do_absolute_pos_embedding": false,
  "index_is_latlon": true,
  "do_arcsinh_input": false,
  "do_te_mlp_last_layer": false,
  "use_cross_layer_concat_skip_connection": false,
  "do_activation_checkpointing": 1.0
 },
 "pixel_dit": {
  "embed_dim": 128,
  "n_layers": 4,
  "num_heads": null,
  "mlp_ratio": 4.0,
  "attn_kernel": 9,
  "gated_attention": false,
  "qk_norm": true,
  "qk_norm_elementwise_affine": false,
  "freeze_semantic": false,
  "freeze_pixel_blocks": false,
  "use_bilinear_dw_gelu_project_adaln": true,
  "do_activation_checkpointing": 0.0
 },
 "pipeline": {
  "target_type": "mixed_state_diff",
  "do_qv_softplus": true,
  "do_precip_relu": true,
  "variables_input_zeroed": []
 },
 "compute": {
  "do_bf16_mixed": true,
  "do_torch_compile": true
 },
 "training": {
  "lr": 1e-05,
  "loss_lr": 0.1,
  "epochs": 50,
  "batch_size": 1,
  "num_workers": 2,
  "num_workers_inference": 2,
  "optimizer": "adamw",
  "beta1": 0.9,
  "beta2": 0.999,
  "weight_decay": 1e-05,
  "seed": 42,
  "scheduler_type": "constant",
  "warmup_steps": 2000,
  "cosine_t_max": 100,
  "cosine_n_cycles": 5,
  "cosine_decay": 0.3162,
  "num_steps": 4,
  "multistep_training_mode": "final_only",
  "do_tf32": true,
  "grad_clip_max_norm": 0.5,
  "backup_every_steps": 100,
  "loss_type": "smooth_l1",
  "resume_from": null,
  "reset_best_valid_loss": false,
  "reset_scheduler_state": false,
  "skip_optimizer_reloading": false,
  "fit_batches_for_norm": 1280,
  "normalization_ref_path_input": "",
  "normalization_ref_path_target": "",
  "do_dual_normalization": false,
  "pooled_loss": null,
  "pooled_loss_channel_weights_csv": null,
  "max_steps_per_epoch": null,
  "validate_steps": 400,
  "validate_min_samples": 1024,
  "validate_total_images": null,
  "validate_num_shards": null,
  "exit_after_first_validation": false,
  "max_validations": 5,
  "no_save": false,
  "save_ckpt_every_epoch": true,
  "use_time_variation_seed": true,
  "do_backtest_inference": false,
  "do_track_gns": false,
  "gns_decay": 0.99,
  "gns_freq": 1
 },
 "data": {
  "grid_type": "cubesphere",
  "cubesphere_ne": 1024,
  "cubesphere_npg": 2,
  "tile_size": 128,
  "nside": 2048,
  "level_start": 8,
  "level_end": 32,
  "plevel": 1,
  "cross_face_tiles": false,
  "use_duo_padding": false,
  "duo_padding_scrip_src_path": null,
  "duo_padding_scrip_tgt_path": null,
  "skip_corner_tiles": false,
  "balance_cross_face": false,
  "variables_prognostic": [
   "PotentialTemperature",
   "U",
   "V",
   "z_mid",
   "omega",
   "qv",
   "T_2m"
  ],
  "variables_prognostic_state": [
   "omega"
  ],
  "variables_forcing": [
   "coszr",
   "phis",
   "sgh30",
   "landfrac"
  ],
  "variables_diagnostic": [
   "precip_ice_surf_mass_flux",
   "precip_liq_surf_mass_flux",
   "ps"
  ],
  "train_main_zarr_paths": [
   "sdecadal.ne1024pg2_ne1024pg2.F20TR-SCREAMv1.c10-sep11.out10min.cubesphere.zarr",
   "sdy2.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr",
   "sdy1.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr"
  ],
  "train_aux_zarr_paths": [
   "sdecadal.ne1024pg2_ne1024pg2.F20TR-SCREAMv1.c10-sep11.out10min.cubesphere.zarr",
   "sdy2.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr",
   "sdy1.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr"
  ],
  "train_zarr_weights": [
   0.5,
   1.0,
   1.0
  ],
  "train_start_indices": [
   144,
   144,
   144
  ],
  "train_end_indices": [
   1435,
   720,
   720
  ],
  "split_by_time": true,
  "train_start_index": 144,
  "train_end_index": 1435,
  "test_start_index": 1872,
  "test_end_index": 2015,
  "test_stride": 36,
  "inference_start_index": 1872,
  "backtest_main_zarr_path": "",
  "backtest_aux_zarr_path": "",
  "backtest_start_index": 864,
  "backtest_end_index": 1007,
  "backtest_stride": 36,
  "latlon_path": "latlon_ne1024pg2.nc",
  "do_fixed_z": false,
  "fixed_z_npy_path": "data/z_coordinate_24levels.npy",
  "phis_npy_path": "data/phis.npy"
 },
 "debug": {
  "print_steps": 5,
  "mock": false,
  "skip_india_ocean_eval": true,
  "use_full_test_set": true,
  "recreate_dataloader_each_validation": true,
  "wandb_disabled": false
 }
}