extended_config.yaml

data_workers: 10
name: Walrus_ft_CNS3D_128_Rand_realglobalnorm_rusty-CNS3D-delta-Isotr[Space-Adapt-Space-FullA-Axial]-AdamW-0.0001
finetune: true
automatic_setup: true
trainer:
  _target_: walrus.trainer.Trainer
  max_epoch: 50
  val_frequency: 5
  rollout_val_frequency: 5
  short_validation_length: 20
  max_rollout_steps: 200
  num_time_intervals: 5
  enable_amp: false
  loss_fn:
    _target_: the_well.benchmark.metrics.MAE
  formatter:
    _target_: hydra.utils.get_class
    path: walrus.data.well_to_multi_transformer.ChannelsFirstWithTimeFormatter
  revin:
    _target_: walrus.trainer.normalization_strat.GlobalRevNormalization
    _partial_: true
  prediction_type: delta
  grad_acc_steps: 1
  image_validation: true
  video_validation: true
  gradient_log_level: 0
  clip_gradient: 10
  log_interval: 200
  loss_multiplier: 100.0
  lr_scheduler_per_step: false
  skip_spectral_metrics: true
optimizer:
  _target_: torch.optim.AdamW
  lr: 0.0001
  weight_decay: 0.0001
  eps: 1.0e-10
lr_scheduler:
  _target_: walrus.optim.schedulers.InverseSqrtLinearWarmupSqrtCooldown
  warmup_epochs: 10
  cooldown_epochs: 10
  warmup_lr_factor: 0.1
  cooldown_lr_factor: 0.001
model:
  encoder:
    _partial_: true
    _target_: walrus.models.encoders.vstride_encoder.SpaceBagAdaptiveDVstrideEncoder
    learned_pad: true
    base_kernel_size1d:
    - - 4
      - 4
    base_kernel_size2d:
    - - 8
      - 4
    - - 8
      - 4
    base_kernel_size3d:
    - - 8
      - 4
    - - 8
      - 4
    - - 8
      - 4
    groups: 12
    kernel_scales_seq:
    - - 2
      - 2
    - - 4
      - 2
    - - 4
      - 4
    - - 8
      - 4
    variable_downsample: true
    variable_deterministic_ds: true
    activation:
      _partial_: true
      _target_: torch.nn.SiLU
  decoder:
    _partial_: true
    _target_: walrus.models.decoders.vstride_decoder.AdaptiveDVstrideDecoder
    learned_pad: true
    base_kernel_size1d:
    - - 4
      - 4
    base_kernel_size2d:
    - - 8
      - 4
    - - 8
      - 4
    base_kernel_size3d:
    - - 8
      - 4
    - - 8
      - 4
    - - 8
      - 4
    groups: 12
    activation:
      _partial_: true
      _target_: torch.nn.SiLU
  processor:
    space_mixing:
      _partial_: true
      _target_: walrus.models.spatial_blocks.full_attention.FullAttention
      num_heads: 16
      mlp_dim: null
    time_mixing:
      _partial_: true
      _target_: walrus.models.temporal_blocks.axial_time_attention.AxialTimeAttention
      num_heads: 16
      bias_type: rel
    channel_mixing:
      _partial_: true
      _target_: torch.nn.Identity
    _partial_: true
    _target_: walrus.models.spatiotemporal_blocks.space_time_split.SpaceTimeSplitBlock
  norm_layer:
    _partial_: true
    _target_: walrus.models.shared_utils.normalization.RMSGroupNorm
  _target_: walrus.models.IsotropicModel
  hidden_dim: 1408
  projection_dim: 48
  intermediate_dim: 352
  processor_blocks: 40
  drop_path: 0.0
  groups: 16
  max_d: 3
  static_axes: true
  weight_tied_axes: false
  causal_in_time: true
  include_d:
  - 2
  - 3
  override_dimensionality: 0
  jitter_patches: true
  gradient_checkpointing_freq: 2
  use_periodic_fixed_jitter: true
  input_field_drop: 0
data:
  field_index_map_override:
    closed_boundary: 0
    open_boundary: 1
    bias_correction: 2
    pressure: 3
    velocity_x: 4
    velocity_y: 5
    velocity_z: 6
    zeros_like_density: 7
    speed_of_sound: 8
    concentration: 9
    D_xx: 10
    D_xy: 11
    D_xz: 12
    D_yx: 13
    D_yy: 14
    D_yz: 15
    D_zx: 16
    D_zy: 17
    D_zz: 18
    E_xx: 19
    E_xy: 20
    E_xz: 21
    E_yx: 22
    E_yy: 23
    E_yz: 24
    E_zx: 25
    E_zy: 26
    E_zz: 27
    density: 28
    energy: 29
    velocity_r: 30
    velocity_theta: 31
    velocity_phi: 32
    momentum_x: 33
    momentum_y: 34
    momentum_z: 35
    pressure_re: 36
    pressure_im: 37
    mask: 38
    magnetic_field_x: 39
    magnetic_field_y: 40
    magnetic_field_z: 41
    A: 42
    B: 43
    height: 44
    internal_energy: 45
    temperature: 46
    electron_fraction: 47
    entropy: 48
    magnetic_field_log_r: 49
    magnetic_field_theta: 50
    magnetic_field_phi: 51
    velocity_log_r: 52
    buoyancy: 53
    tracer: 54
    log10_density: 55
    log10_temperature: 56
    c_zz: 57
    C_xx: 58
    C_xy: 59
    C_xz: 60
    C_yx: 61
    C_yy: 62
    C_yz: 63
    C_zx: 64
    C_zy: 65
    C_zz: 66
  well_base_path: /mnt/home/polymathic/ceph/the_well/datasets/
  wandb_data_name: CNS3D_128_Rand
  module_parameters:
    _target_: walrus.data.MixedWellDataModule
    batch_size: 1
    n_steps_input: 3
    n_steps_output: 1
    min_dt_stride: 1
    max_dt_stride: 1
    max_samples: 2000
    well_dataset_info:
      CNS3D_128_Rand:
        include_filters: []
        exclude_filters: []
        path: /mnt/home/polymathic/ceph/WellFormattedExternalData/PDEBench/CNS3D_128_Rand/
        normalization_path: logged_stats.yaml
auto_resume: true
folder_override: ''
checkpoint_override: ''
config_override: /mnt/home/polymathic/ceph/walrus_logging/platinum_checkpoints/extended_config.yaml
validation_mode: false
frozen_components:
- model
distribution:
  distribution_type: fsdp
  local_size: null
logger:
  wandb: true
  wandb_project_name: walrus_Finetuning_Runs
checkpoint:
  _target_: walrus.trainer.checkpoints.CheckPointer
  save_dir: /mnt/home/polymathic/ceph/walrus_logging/runs/Walrus_ft_CNS3D_128_Rand_realglobalnorm_rusty-CNS3D-delta-Isotr[Space-Adapt-Space-FullA-Axial]-AdamW-0.0001/finetune/0/checkpoints
  load_checkpoint_path: null
  coalesced_checkpoint_path: /mnt/home/polymathic/ceph/walrus_logging/platinum_checkpoints/final_base_model/walrus.pt
  save_best: true
  checkpoint_frequency: 20
finetuning_mods:
  learnable_rope: true
  rope_per_axis: true
  ape_shape:
  - 16
  - 16
  - 16
experiment_dir: /mnt/home/polymathic/ceph/walrus_logging/runs