checkpoint:
  broadcast_via_filesystem: 'False'
  dcp_allow_mismatched_size: 'False'
  dcp_async_mode_enabled: 'False'
  early_sanity_check: '5'
  jit:
    device: cuda
    dtype: bfloat16
    enabled: 'False'
    input_shape: null
    strict: 'True'
  keys_not_to_resume: []
  load_ema_to_reg: 'False'
  load_path: ''
  load_training_state: 'False'
  only_load_scheduler_state: 'False'
  pretrained: null
  resume: null
  s3_folder: s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/
  save_iter: '1000'
  strict_resume: 'True'
  type:
    _target_: <class 'cosmos_predict2.checkpointer.Checkpointer'>
    callbacks: null
  verbose: 'True'
dataloader_train:
  _target_: <class 'torch.utils.data.dataloader.DataLoader'>
  batch_sampler: null
  batch_size: '4'
  collate_fn: null
  dataset:
    _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
    dataset_dir: custom/config/train/multi/drive_ike3_40h_16g.yaml
    num_frames: '93'
    params: null
    phase: train
    video_size:
    - '704'
    - '1280'
  drop_last: 'True'
  generator: null
  in_order: 'True'
  multiprocessing_context: null
  num_workers: '4'
  persistent_workers: 'False'
  pin_memory: 'True'
  pin_memory_device: ''
  prefetch_factor: null
  sampler:
    _target_: <function get_sampler at 0x7ff46b0e0a40>
    dataset:
      _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
      dataset_dir: custom/config/train/multi/drive_ike3_40h_16g.yaml
      num_frames: '93'
      params: null
      phase: train
      video_size:
      - '704'
      - '1280'
    drop_last: 'True'
    shuffle: 'True'
  shuffle: null
  timeout: '0'
  worker_init_fn: null
dataloader_val:
  ArgoVQL:
    _target_: <class 'torch.utils.data.dataloader.DataLoader'>
    batch_sampler: null
    batch_size: '1'
    collate_fn: null
    dataset:
      _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
      dataset_dir: custom/config/val/argo_ike3_40h.yaml
      num_frames: '93'
      params: null
      phase: val
      video_size:
      - '704'
      - '1280'
    drop_last: 'False'
    generator: null
    in_order: 'True'
    multiprocessing_context: null
    num_workers: '0'
    persistent_workers: 'False'
    pin_memory: 'False'
    pin_memory_device: ''
    prefetch_factor: null
    sampler:
      _target_: <function get_sampler at 0x7ff46b0e0a40>
      dataset:
        _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
        dataset_dir: custom/config/val/argo_ike3_40h.yaml
        num_frames: '93'
        params: null
        phase: val
        video_size:
        - '704'
        - '1280'
      drop_last: 'False'
      shuffle: 'False'
    shuffle: null
    timeout: '0'
    worker_init_fn: null
  DDADQL:
    _target_: <class 'torch.utils.data.dataloader.DataLoader'>
    batch_sampler: null
    batch_size: '1'
    collate_fn: null
    dataset:
      _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
      dataset_dir: custom/config/val/ddad_ike3_40h.yaml
      num_frames: '93'
      params: null
      phase: val
      video_size:
      - '704'
      - '1280'
    drop_last: 'False'
    generator: null
    in_order: 'True'
    multiprocessing_context: null
    num_workers: '0'
    persistent_workers: 'False'
    pin_memory: 'False'
    pin_memory_device: ''
    prefetch_factor: null
    sampler:
      _target_: <function get_sampler at 0x7ff46b0e0a40>
      dataset:
        _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
        dataset_dir: custom/config/val/ddad_ike3_40h.yaml
        num_frames: '93'
        params: null
        phase: val
        video_size:
        - '704'
        - '1280'
      drop_last: 'False'
      shuffle: 'False'
    shuffle: null
    timeout: '0'
    worker_init_fn: null
  LyftQL:
    _target_: <class 'torch.utils.data.dataloader.DataLoader'>
    batch_sampler: null
    batch_size: '1'
    collate_fn: null
    dataset:
      _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
      dataset_dir: custom/config/val/lyft_ike3_40h.yaml
      num_frames: '93'
      params: null
      phase: val
      video_size:
      - '704'
      - '1280'
    drop_last: 'False'
    generator: null
    in_order: 'True'
    multiprocessing_context: null
    num_workers: '0'
    persistent_workers: 'False'
    pin_memory: 'False'
    pin_memory_device: ''
    prefetch_factor: null
    sampler:
      _target_: <function get_sampler at 0x7ff46b0e0a40>
      dataset:
        _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
        dataset_dir: custom/config/val/lyft_ike3_40h.yaml
        num_frames: '93'
        params: null
        phase: val
        video_size:
        - '704'
        - '1280'
      drop_last: 'False'
      shuffle: 'False'
    shuffle: null
    timeout: '0'
    worker_init_fn: null
  PD4DAllQL:
    _target_: <class 'torch.utils.data.dataloader.DataLoader'>
    batch_sampler: null
    batch_size: '1'
    collate_fn: null
    dataset:
      _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
      dataset_dir: custom/config/val/pd4d_all_ike3_40h.yaml
      num_frames: '93'
      params: null
      phase: val
      video_size:
      - '704'
      - '1280'
    drop_last: 'False'
    generator: null
    in_order: 'True'
    multiprocessing_context: null
    num_workers: '0'
    persistent_workers: 'False'
    pin_memory: 'False'
    pin_memory_device: ''
    prefetch_factor: null
    sampler:
      _target_: <function get_sampler at 0x7ff46b0e0a40>
      dataset:
        _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
        dataset_dir: custom/config/val/pd4d_all_ike3_40h.yaml
        num_frames: '93'
        params: null
        phase: val
        video_size:
        - '704'
        - '1280'
      drop_last: 'False'
      shuffle: 'False'
    shuffle: null
    timeout: '0'
    worker_init_fn: null
  PD4DEgoQL:
    _target_: <class 'torch.utils.data.dataloader.DataLoader'>
    batch_sampler: null
    batch_size: '1'
    collate_fn: null
    dataset:
      _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
      dataset_dir: custom/config/val/pd4d_ego_ike3_40h.yaml
      num_frames: '93'
      params: null
      phase: val
      video_size:
      - '704'
      - '1280'
    drop_last: 'False'
    generator: null
    in_order: 'True'
    multiprocessing_context: null
    num_workers: '0'
    persistent_workers: 'False'
    pin_memory: 'False'
    pin_memory_device: ''
    prefetch_factor: null
    sampler:
      _target_: <function get_sampler at 0x7ff46b0e0a40>
      dataset:
        _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
        dataset_dir: custom/config/val/pd4d_ego_ike3_40h.yaml
        num_frames: '93'
        params: null
        phase: val
        video_size:
        - '704'
        - '1280'
      drop_last: 'False'
      shuffle: 'False'
    shuffle: null
    timeout: '0'
    worker_init_fn: null
  PDv2QL:
    _target_: <class 'torch.utils.data.dataloader.DataLoader'>
    batch_sampler: null
    batch_size: '1'
    collate_fn: null
    dataset:
      _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
      dataset_dir: custom/config/val/pdv2_ike3_40h.yaml
      num_frames: '93'
      params: null
      phase: val
      video_size:
      - '704'
      - '1280'
    drop_last: 'False'
    generator: null
    in_order: 'True'
    multiprocessing_context: null
    num_workers: '0'
    persistent_workers: 'False'
    pin_memory: 'False'
    pin_memory_device: ''
    prefetch_factor: null
    sampler:
      _target_: <function get_sampler at 0x7ff46b0e0a40>
      dataset:
        _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
        dataset_dir: custom/config/val/pdv2_ike3_40h.yaml
        num_frames: '93'
        params: null
        phase: val
        video_size:
        - '704'
        - '1280'
      drop_last: 'False'
      shuffle: 'False'
    shuffle: null
    timeout: '0'
    worker_init_fn: null
  WaymoQL:
    _target_: <class 'torch.utils.data.dataloader.DataLoader'>
    batch_sampler: null
    batch_size: '1'
    collate_fn: null
    dataset:
      _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
      dataset_dir: custom/config/val/waymo_ike3_40h.yaml
      num_frames: '93'
      params: null
      phase: val
      video_size:
      - '704'
      - '1280'
    drop_last: 'False'
    generator: null
    in_order: 'True'
    multiprocessing_context: null
    num_workers: '0'
    persistent_workers: 'False'
    pin_memory: 'False'
    pin_memory_device: ''
    prefetch_factor: null
    sampler:
      _target_: <function get_sampler at 0x7ff46b0e0a40>
      dataset:
        _target_: <class 'cosmos_predict2.data.dataset_video.Dataset'>
        dataset_dir: custom/config/val/waymo_ike3_40h.yaml
        num_frames: '93'
        params: null
        phase: val
        video_size:
        - '704'
        - '1280'
      drop_last: 'False'
      shuffle: 'False'
    shuffle: null
    timeout: '0'
    worker_init_fn: null
defaults:
- _self_
- data_train: null
- data_val: null
- optimizer: fusedadamw
- scheduler: constant
- model: predict2_video2world_fsdp_2b
- callbacks:
  - basic
- net: null
- ema: null
- checkpoint: null
- ckpt_type: null
- experiment: null
job:
  append_datetime: 'False'
  group: test
  name: wm3
  prepend_datetime: 'True'
  processed: 'False'
  project: a4d2
  undated_name: ''
metrics: null
model:
  _recursive_: 'False'
  _target_: <class 'custom.any4d.a4d_model.Any4DModel'>
  config:
    action_multiplier: 2.0
    adjust_video_noise: true
    all_highdim_entries: []
    all_lowdim_entries: []
    any4d_active: true
    arch: ''
    block_gate_fix: true
    cams_viz_type: comb
    debug_without_randomness: false
    depth_viz_type: viridis
    force_cond_noise_level: false
    fsdp_shard_size: 8
    harmonize_frames: true
    harmonize_streams: true
    has_adaln_stream: false
    has_sattn_stream: false
    has_xattn_stream: false
    high_sigma_ratio: 0.05
    ignore_unused_data: true
    init_lora_weights: true
    input_data_key: video
    input_image_key: images
    job: {}
    load_modals:
    - rgb
    - traj
    lora_alpha: 16
    lora_rank: 16
    lora_target_modules: q_proj,k_proj,v_proj,output_proj,mlp.layer1,mlp.layer2
    loss_reduce: mean
    loss_scale: 10.0
    loss_weights:
      rgb0: 1.0
      rgb1: 1.0
      rgb2: 1.0
    lowdim_adaln_entries:
      traj:
      - '0'
      - '41'
      - '0'
      - '2'
      - zero/load
      traj_input_mask:
      - '0'
      - '41'
      - '2'
      - '3'
      - zero/load
    lowdim_hidden_channels: 64
    lowdim_sattn_entries: {}
    lowdim_sattn_proj_mode: per_token
    lowdim_timestep_mode: per_stream
    lowdim_xattn_entries: {}
    lowdim_xattn_proj_mode: per_token
    model_manager_config:
      _target_: cosmos_predict2.models.video2world_model.Predict2ModelManagerConfig
      dit_ema_path: unused
      dit_path: s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-11-01-56_s27e_rwm3_drive40h/model/iter_000049000_012544000.pt
      text_encoder_path: ''
      tokenizer_chunk_duration: '41'
      vae_path: s3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth
    normalizer_stats: ''
    num_lowdim_adaln_channels: -1
    num_lowdim_adaln_tokens: -1
    num_lowdim_sattn_channels: -1
    num_lowdim_sattn_tokens: -1
    num_lowdim_xattn_channels: -1
    num_lowdim_xattn_tokens: -1
    num_video_in_channels: []
    num_video_out_channels: []
    num_views: 3
    override_temporal_stride_train: '1'
    override_temporal_stride_val: '1'
    override_zero_origin: 'True'
    pipe_config:
      adjust_video_noise: true
      conditioner:
        _target_: <class 'cosmos_predict2.configs.vid2vid.defaults.conditioner.Vid2VidConditioner'>
        fps:
          _target_: <class 'cosmos_predict2.conditioner.ReMapkey'>
          dropout_rate: '0.0'
          dtype: null
          input_key: fps
          output_key: fps
        padding_mask:
          _target_: <class 'cosmos_predict2.conditioner.ReMapkey'>
          dropout_rate: '0.0'
          dtype: null
          input_key: padding_mask
          output_key: padding_mask
        text:
          _target_: <class 'cosmos_predict2.conditioner.TextAttr'>
          dropout_rate: '0.2'
          input_key:
          - t5_text_embeddings
        use_video_condition:
          _target_: <class 'cosmos_predict2.conditioner.BooleanFlag'>
          dropout_rate: '0.0'
          input_key: fps
          output_key: use_video_condition
      conditioning_strategy: frame_replace
      ema:
        _target_: cosmos_predict2.configs.base.defaults.ema.EMAConfig
        enabled: 'False'
        iteration_shift: '0'
        rate: '0.1'
      extra_nets: {}
      guardrail_config:
        checkpoint_dir: /basile/ws/repos/cpred2-basile/logs_a4d/checkpoints/
        enabled: false
        offload_model_to_cpu: true
      input_data_key: video
      input_image_key: images
      max_num_conditional_frames: 2
      min_num_conditional_frames: 2
      net:
        _target_: <class 'custom.any4d.a4d_network.Any4DDiT'>
        action_dim: '-1'
        adaln_lora_dim: '256'
        any4d_config: null
        atten_backend: minimal_a2a
        concat_padding_mask: 'True'
        extra_per_block_abs_pos_emb: 'False'
        in_channels: '16'
        max_frames: '128'
        max_img_h: '240'
        max_img_w: '240'
        model_channels: '2048'
        num_blocks: '28'
        num_heads: '16'
        out_channels: '16'
        patch_spatial: '2'
        patch_temporal: '1'
        pos_emb_cls: rope3d
        pos_emb_interpolation: crop
        pos_emb_learnable: 'True'
        rope_enable_fps_modulation: 'False'
        rope_h_extrapolation_ratio: '3.0'
        rope_t_extrapolation_ratio: '1.0'
        rope_w_extrapolation_ratio: '3.0'
        sac_config:
          _target_: cosmos_predict2.models.text2image_dit.SACConfig
          every_n_blocks: '1'
          mode: predict2_2b_720
        use_adaln_lora: 'True'
      precision: bfloat16
      prompt_refiner_config:
        checkpoint_dir: /basile/ws/repos/cpred2-basile/logs_a4d/checkpoints/nvidia/Cosmos-Reason1-7B
        enabled: false
        offload_model_to_cpu: true
      rectified_flow_t_scaling_factor: 1.0
      remove_dit: false
      resize_online: false
      resolution: '720'
      sigma_conditional: 0.0001
      sigma_data: 1.0
      state_ch: 16
      state_t: 24
      text_encoder_class: T5
      timestamps:
        is_forward: false
        nfe: 35
        order: 7.0
        t_max: 80.0
        t_min: 0.002
      tokenizer:
        _target_: <class 'cosmos_predict2.tokenizers.tokenizer.TokenizerInterface'>
        chunk_duration: '81'
        load_mean_std: 'False'
        name: tokenizer
        temporal_window: '4'
        vae_pth: ''
    precision: bfloat16
    relative_actions: false
    shuffle_cams2views: false
    strict_populate_data: false
    task_probs:
      cross_modal: 0.0
      dyn_view_synth: 0.0
      forecast: 1.0
      inverse_dynamics: 0.0
      policy: 0.0
      pose_estimation: 0.0
      world_model: 0.9
    track_metrics:
      rgb0:
      - psnr
      - ssim
      rgb1:
      - psnr
      - ssim
    train_architecture: base
    train_directives: {}
    train_visuals_detail: 1
    train_visuals_interval: 99
    transforms: default
    use_views: all
    vae: a4d_vae
    val_cfg_scale: 0.0
    val_cfg_zero_entries: []
    val_directives:
      perturb_traj: basile3
      remove_cond: 'True'
      tasks: forecast,world_model
    val_num_steps: 35
    val_sigma_max: 80.0
    val_sigma_min: 0.002
    val_visuals_detail: 1
    validated: false
    vidar2a4d: basile
    vidar_active: true
    video_concat_mode: view
    video_entries:
    - rgb0:
      - '0'
      - '16'
      - load
      - load
      rgb0_input_mask:
      - '16'
      - '17'
      - load
      - null
      rgb0_output_mask:
      - '17'
      - '18'
      - load
      - null
    - rgb1:
      - '0'
      - '16'
      - copy:rgb0/load
      - copy:rgb0/load
      rgb1_input_mask:
      - '16'
      - '17'
      - copy:rgb0_input_mask/load
      - null
      rgb1_output_mask:
      - '17'
      - '18'
      - copy:rgb0_output_mask/load
      - null
    - rgb2:
      - '0'
      - '16'
      - copy:rgb0/load
      - copy:rgb0/load
      rgb2_input_mask:
      - '16'
      - '17'
      - copy:rgb0_input_mask/load
      - null
      rgb2_output_mask:
      - '17'
      - '18'
      - copy:rgb0_output_mask/load
      - null
    video_proj_mode: per_view
    view_emb_std: 0.06
    view_timestep_mode: per_view
    visuals_quality: 8
    viz_extra_modes: []
    viz_input_blacklist: []
    viz_mask_border_width: 2
    wandb:
      dir: ''
      enabled: 'True'
      entity: tri
      name: wm3
      num_validation_logs: '5'
      project: a4d2
      tags:
      - a4d2
      - debug
      - wm3
model_parallel:
  _cpu_offloading_context: null
  async_tensor_model_parallel_allreduce: false
  autocast_dtype: torch.float32
  barrier_with_L1_time: true
  batch_p2p_comm: true
  batch_p2p_sync: true
  bf16: false
  context_parallel_size: 1
  cpu_offloading: false
  cpu_offloading_activations: true
  cpu_offloading_num_layers: 0
  cpu_offloading_weights: true
  cross_entropy_fusion_impl: native
  cross_entropy_loss_fusion: false
  deallocate_pipeline_outputs: false
  defer_embedding_wgrad_compute: false
  deterministic_mode: false
  enable_autocast: false
  expert_model_parallel_size: 1
  expert_tensor_parallel_size: 1
  finalize_model_grads_func: null
  fp16: false
  grad_scale_func: null
  grad_sync_func: null
  gradient_accumulation_fusion: false
  hierarchical_context_parallel_sizes: null
  microbatch_group_size_per_vp_stage: 1
  moe_extended_tp: false
  no_sync_func: null
  num_microbatches_with_partial_activation_checkpoints: null
  overlap_p2p_comm: false
  overlap_p2p_comm_warmup_flush: false
  param_sync_func: null
  params_dtype: torch.float32
  perform_initialization: true
  pipeline_dtype: null
  pipeline_model_parallel_comm_backend: null
  pipeline_model_parallel_size: 1
  pipeline_model_parallel_split_rank: null
  sequence_parallel: false
  tensor_model_parallel_size: 1
  timers: null
  tp_comm_atomic_ag: false
  tp_comm_atomic_rs: false
  tp_comm_bootstrap_backend: nccl
  tp_comm_bulk_dgrad: true
  tp_comm_bulk_wgrad: true
  tp_comm_overlap: false
  tp_comm_overlap_ag: true
  tp_comm_overlap_disable_fc1: false
  tp_comm_overlap_disable_qkv: false
  tp_comm_overlap_rs: true
  tp_comm_overlap_rs_dgrad: false
  tp_comm_split_ag: true
  tp_comm_split_rs: true
  use_cpu_initialization: false
  use_ring_exchange_p2p: false
  use_te_rng_tracker: false
  variable_seq_lengths: false
  virtual_pipeline_model_parallel_size: null
  wgrad_deferral_limit: 0
optimizer:
  _target_: <function get_base_optimizer at 0x7ff46b0e09a0>
  betas:
  - '0.9'
  - '0.99'
  capturable: 'True'
  eps: 1e-08
  lr: 3e-05
  master_weights: 'True'
  model: null
  optim_type: fusedadam
  weight_decay: '0.1'
params_for_wandb:
  any4d:
    action_multiplier: '2.0'
    block_gate_fix: 'True'
    harmonize_frames: 'True'
    harmonize_streams: 'True'
    load_modals:
    - rgb
    - traj
    loss_weights:
      rgb0: '1.0'
      rgb1: '1.0'
      rgb2: '1.0'
    lowdim_adaln_entries:
      traj:
      - '0'
      - '41'
      - '0'
      - '2'
      - zero/load
      traj_input_mask:
      - '0'
      - '41'
      - '2'
      - '3'
      - zero/load
    normalizer_stats: ''
    num_views: '3'
    override_temporal_stride_train: '1'
    override_temporal_stride_val: '1'
    override_zero_origin: 'True'
    relative_actions: 'False'
    shuffle_cams2views: 'False'
    task_probs:
      cross_modal: '0.0'
      dyn_view_synth: '0.0'
      forecast: '1.0'
      inverse_dynamics: '0.0'
      policy: '0.0'
      pose_estimation: '0.0'
      world_model: '0.9'
    track_metrics:
      rgb0:
      - psnr
      - ssim
      rgb1:
      - psnr
      - ssim
    train_directives: {}
    train_visuals_detail: '1'
    train_visuals_interval: '99'
    transforms: default
    use_views: all
    vae: a4d_vae
    val_directives:
      perturb_traj: basile3
      remove_cond: 'True'
      tasks: forecast,world_model
    val_num_steps: '35'
    val_visuals_detail: '1'
    vidar2a4d: basile
    video_concat_mode: view
    video_entries:
    - rgb0:
      - '0'
      - '16'
      - load
      - load
      rgb0_input_mask:
      - '16'
      - '17'
      - load
      - null
      rgb0_output_mask:
      - '17'
      - '18'
      - load
      - null
    - rgb1:
      - '0'
      - '16'
      - copy:rgb0/load
      - copy:rgb0/load
      rgb1_input_mask:
      - '16'
      - '17'
      - copy:rgb0_input_mask/load
      - null
      rgb1_output_mask:
      - '17'
      - '18'
      - copy:rgb0_output_mask/load
      - null
    - rgb2:
      - '0'
      - '16'
      - copy:rgb0/load
      - copy:rgb0/load
      rgb2_input_mask:
      - '16'
      - '17'
      - copy:rgb0_input_mask/load
      - null
      rgb2_output_mask:
      - '17'
      - '18'
      - copy:rgb0_output_mask/load
      - null
    video_proj_mode: per_view
    view_timestep_mode: per_view
    visuals_quality: '8'
    viz_input_blacklist: []
    viz_mask_border_width: '2'
  checkpoint:
    early_sanity_check: '5'
    s3_folder: s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/
    save_iter: '1000'
  dataset_train:
    batch_size: '4'
    config: custom/config/train/multi/drive_ike3_40h_16g.yaml
    num_workers: '4'
  dataset_val:
    batch_size: '1'
    config:
      ArgoVQL: custom/config/val/argo_ike3_40h.yaml
      DDADQL: custom/config/val/ddad_ike3_40h.yaml
      LyftQL: custom/config/val/lyft_ike3_40h.yaml
      PD4DAllQL: custom/config/val/pd4d_all_ike3_40h.yaml
      PD4DEgoQL: custom/config/val/pd4d_ego_ike3_40h.yaml
      PDv2QL: custom/config/val/pdv2_ike3_40h.yaml
      WaymoQL: custom/config/val/waymo_ike3_40h.yaml
    num_workers: '0'
  job:
    group: debug
    name: wm3
    prepend_datetime: 'True'
    project: a4d2
  model:
    context_parallel_size: '1'
    device_monitor: '0'
    dit_path: s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-11-01-56_s27e_rwm3_drive40h/model/iter_000049000_012544000.pt
    ema_enabled: 'False'
    extra_nets: {}
    fsdp_shard_size: '8'
    grad_accum_iter: '1'
    manual_gc: '500'
    manual_gc_iter: '288'
    manual_gc_warm_up: '-1'
    max_iter: '80000'
    model: predict2_video2world_fsdp_2b
    net: {}
    remove_dit: 'False'
    run_validation: 'True'
    skip_first_validation: delay
    state_t: '11'
    text_encoder_path: ''
    tokenizer_chunk_duration: '41'
    vae_path: s3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth
    validation_iter: '500'
  optimizer:
    lr: 3e-05
  scheduler:
    cycle_lengths:
    - '80000'
    f_max:
    - '1.0'
    f_min:
    - '0.05'
    f_start:
    - '0.01'
    warm_up_steps:
    - '500'
  wandb:
    dir: ''
    enabled: 'True'
    entity: tri
    name: wm3
    num_validation_logs: '5'
    project: a4d2
    tags:
    - a4d2
    - debug
    - wm3
scheduler:
  _target_: <class 'cosmos_predict2.functional.lr_scheduler.LambdaLinearScheduler'>
  cycle_lengths:
  - '80000'
  f_max:
  - '1.0'
  f_min:
  - '0.05'
  f_start:
  - '0.01'
  verbosity_interval: '0'
  warm_up_steps:
  - '500'
trainer:
  callbacks:
    device_monitor:
      _target_: <class 'cosmos_predict2.callbacks.device_monitor.DeviceMonitor'>
      every_n: '0'
      log_memory_detail: 'True'
      step_size: '1'
    iter_speed:
      _target_: <class 'cosmos_predict2.callbacks.iter_speed.IterSpeed'>
      every_n: '10'
      hit_thres: '0'
    loss_log:
      _target_: <class 'cosmos_predict2.callbacks.loss_log.LossLog'>
      logging_iter_multipler: '1'
    low_prec:
      _target_: <class 'imaginaire.utils.callback.LowPrecisionCallback'>
      config: null
      trainer: null
      update_iter: '1'
    manual_gc:
      _target_: <class 'imaginaire.callbacks.manual_gc.ManualGarbageCollection'>
      every_n: '288'
      warm_up: '-1'
  cudnn:
    benchmark: 'True'
    deterministic: 'False'
  ddp:
    broadcast_buffers: 'True'
    find_unused_parameters: 'False'
    static_graph: 'True'
  distributed_parallelism: fsdp
  grad_accum_iter: '1'
  grad_scaler_args:
    enabled: 'False'
  logging_iter: '10'
  max_iter: '80000'
  max_val_iter: null
  memory_format: torch.preserve_format
  profiling:
    enable_memory_snapshot: 'False'
    enable_profiling: 'False'
    first_n_rank: '8'
    profile_freq: '1'
    profile_memory: 'True'
    record_shape: 'True'
    with_modules: 'True'
    with_stack: 'True'
  run_validation: 'True'
  seed: '0'
  skip_first_validation: delay
  timeout_period: '999999999'
  type: <class 'imaginaire.trainer.ImaginaireTrainer'>
  validation_iter: '500'
