initial commit post neutron star

coalesced.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:118dc6735c440435c7765096c46c117cb28b18cdb90e8dbb3457bc63b406f91a
+size 5179778007

extended_config.yaml

@@ -0,0 +1,269 @@
+data_workers: 10
+name: WalrusFT_pns_1step_realglobalnorm-postn-delta-Isotr[Space-Adapt-]-AdamW-8e-05
+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: 30
+  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
+  validation_suite:
+  - _target_: the_well.benchmark.metrics.NRMSE
+  - _target_: the_well.benchmark.metrics.VRMSE
+  - _target_: the_well.benchmark.metrics.PearsonR
+  batch_aggregation_fns:
+  - torch.mean
+  - torch.median
+  - torch.std
+  skip_spectral_metrics: true
+optimizer:
+  _target_: torch.optim.AdamW
+  weight_decay: 0.0001
+  eps: 1.0e-10
+  lr: 8.0e-05
+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
+    log10_internal_energy: 67
+    log10_pressure: 68
+    log10_entropy: 69
+  well_base_path: /mnt/gpuxl/polymathic/the_well/datasets/
+  wandb_data_name: post_neutron_star_merger
+  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:
+      post_neutron_star_merger:
+        include_filters: []
+        exclude_filters: []
+        normalization_path: logged_stats.yaml
+        field_transforms:
+          density: torch.log10
+          temperature: torch.log10
+          pressure: torch.log10
+          entropy: torch.log10
+          internal_energy: torch.log10
+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/WalrusFT_pns_1step_realglobalnorm-postn-delta-Isotr[Space-Adapt-]-AdamW-8e-05/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
+  align_fields: true
+  load_chkpt_after_finetuning_expansion: false
+finetuning_mods:
+  learnable_rope: true
+  rope_per_axis: true
+  ape_shape:
+  - 17
+  - 17
+  - 16
+experiment_dir: /mnt/home/polymathic/ceph/walrus_logging/runs