| { |
| "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 |
| } |
| } |
