Upload folder using huggingface_hub

action_decoder/chain/w2a_so101_task1_slow_100ep_lora32_iter2500_75m_2gpu/config.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:38fbe9ece6d1fba6f67ea30d744ba418b7caa40e014e2948360b3c99bda3dd35
-size 77009
+oid sha256:41f56464adc70dae57125145736b494790c89e17ff5171ba498611965579d46d
+size 77007

action_decoder/chain/w2a_so101_task1_slow_100ep_lora32_iter2500_75m_2gpu/config.yaml

@@ -20,7 +20,7 @@ checkpoint:
     callbacks: null
   verbose: 'True'
 data_config:
-  _target_: <function get_data_config at 0x7f3d9e38a9e0>
+  _target_: <function get_data_config at 0x7b8b98c5a950>
   config_name: so101_lerobot
 dataloader_train:
   _target_: <class 'torch.utils.data.dataloader.DataLoader'>
@@ -56,11 +56,11 @@ dataloader_train:
   generator: null
   in_order: 'True'
   multiprocessing_context: null
-  num_workers: '8'
-  persistent_workers: 'True'
+  num_workers: '0'
+  persistent_workers: 'False'
   pin_memory: 'True'
   pin_memory_device: ''
-  prefetch_factor: '8'
+  prefetch_factor: null
   sampler:
     _target_: <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>
     dataset:
@@ -89,11 +89,11 @@ dataloader_train:
       val_ratio: null
       video_backend: pyav
     num_replicas:
-      _target_: <function get_data_parallel_world_size at 0x7f4008f951b0>
+      _target_: <function get_data_parallel_world_size at 0x7b8dff895120>
       partial_data_parallel: 'False'
       with_context_parallel: 'False'
     rank:
-      _target_: <function get_data_parallel_rank at 0x7f4008f952d0>
+      _target_: <function get_data_parallel_rank at 0x7b8dff895240>
       partial_data_parallel: 'False'
       with_context_parallel: 'False'
     seed: '0'
@@ -168,11 +168,11 @@ dataloader_val:
       val_ratio: null
       video_backend: pyav
     num_replicas:
-      _target_: <function get_data_parallel_world_size at 0x7f4008f951b0>
+      _target_: <function get_data_parallel_world_size at 0x7b8dff895120>
       partial_data_parallel: 'False'
       with_context_parallel: 'False'
     rank:
-      _target_: <function get_data_parallel_rank at 0x7f4008f952d0>
+      _target_: <function get_data_parallel_rank at 0x7b8dff895240>
       partial_data_parallel: 'False'
       with_context_parallel: 'False'
     seed: '0'
@@ -208,7 +208,7 @@ model:
     _target_: cosmos_predict2.models.world2action_model.World2ActionModelConfig
     action_dit_path: ''
     data_config:
-      _target_: <function get_data_config at 0x7f3d9e38a9e0>
+      _target_: <function get_data_config at 0x7b8b98c5a950>
       config_name: so101_lerobot
     ema:
       enabled: false
@@ -414,7 +414,7 @@ model_parallel:
   virtual_pipeline_model_parallel_size: null
   wgrad_deferral_limit: 0
 optimizer:
-  _target_: <function get_base_optimizer at 0x7f3d8bda3640>
+  _target_: <function get_base_optimizer at 0x7b8b8276f6d0>
   betas:
   - '0.9'
   - '0.99'

action_decoder/chain/w2a_so101_task1_slow_100ep_lora32_iter2500_75m_2gpu/stdout.log

@@ -95,3 +95,197 @@
 [05-20 02:45:28|INFO|imaginaire/trainer.py:388:validate] Validating at iteration 0...
 [05-20 02:46:08|INFO|imaginaire/trainer.py:193:train] Initial validation done.
 [05-20 02:46:13|WARNING|imaginaire/utils/distributed.py:284:ddp_sync_grad] DDP static_graph=True is incompatible with sync_grad(). Performance will be reduced.
+[05-20 11:48:57|INFO|imaginaire/trainer.py:116:__init__] Config:
+* model: {'config': {'validation_mode': 'loss_only', 'train_architecture': 'base', 'lora_rank': 16, 'lora_alpha': 16, 'lora_target_modules': 'q_proj,k_proj,v_proj,output_proj,mlp.layer1,mlp.layer2', 'init_lora_weights': True, 'precision': 'bfloat16', 'loss_reduce': 'mean', 'loss_scale': 10.0, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, 'action_dit_path': '', 'video_dit_path': '/home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt', 'pipe_config': {'precision': 'bfloat16', 'scheduler': {'alpha': 1.0, 'beta': 1.0, 'num_denoising_steps': 10}, 'net': {'max_horizon': 91, 'in_channels': 6, 'out_channels': 6, 'model_channels': 512, 'num_blocks': 12, 'num_heads': 8, 'mlp_ratio': 4.0, 'atten_backend': 'flash_attn_no_cp', 'crossattn_emb_channels': 2048, 'use_adaln_lora': True, 'adaln_lora_dim': 128, 'pair_timestep_feature_rank': 512, 'sac_config': {'mode': 'none', 'every_n_blocks': 1}, '_target_': <class 'cosmos_predict2.models.world2action_dit.World2ActionDIT'>}, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, '_target_': 'cosmos_predict2.configs.config_world2action.World2ActionPipelineConfig', 'xattn_layer_idx': 20}, 'video_pipe_config': {'adjust_video_noise': True, 'conditioner': {'fps': {'output_key': 'fps', 'dropout_rate': 0.0, 'dtype': None, 'input_key': 'fps', '_target_': <class 'cosmos_predict2.conditioner.ReMapkey'>}, 'padding_mask': {'output_key': 'padding_mask', 'dropout_rate': 0.0, 'dtype': None, 'input_key': 'padding_mask', '_target_': <class 'cosmos_predict2.conditioner.ReMapkey'>}, 'text': {'dropout_rate': 0.0, 'input_key': ['obs/language_embedding'], '_target_': <class 'cosmos_predict2.conditioner.TextAttr'>}, 'use_video_condition': {'output_key': 'use_video_condition', 'dropout_rate': 0.0, 'input_key': 'fps', '_target_': <class 'cosmos_predict2.conditioner.BooleanFlag'>}, '_target_': <class 'cosmos_predict2.conditioner.VideoConditioner'>}, 'conditioning_strategy': 'frame_replace', 'min_num_conditional_frames': 1, 'max_num_conditional_frames': 2, 'sigma_conditional': 0.0001, 'net': {'max_img_h': 240, 'max_img_w': 240, 'max_frames': 128, 'in_channels': 16, 'out_channels': 16, 'patch_spatial': 2, 'patch_temporal': 1, 'concat_padding_mask': True, 'model_channels': 2048, 'num_blocks': 28, 'num_heads': 16, 'atten_backend': 'minimal_a2a', 'pos_emb_cls': 'rope3d', 'pos_emb_learnable': True, 'pos_emb_interpolation': 'crop', 'use_adaln_lora': True, 'adaln_lora_dim': 256, 'rope_h_extrapolation_ratio': 3.0, 'rope_w_extrapolation_ratio': 3.0, 'rope_t_extrapolation_ratio': 1.0, 'extra_per_block_abs_pos_emb': False, 'rope_enable_fps_modulation': False, 'sac_config': {'mode': 'predict2_2b_720', 'every_n_blocks': 1, '_target_': 'cosmos_predict2.models.text2image_dit.SACConfig'}, '_target_': <class 'cosmos_predict2.models.video2world_dit.MinimalV1LVGDiT'>}, 'tokenizer': {'chunk_duration': 81, 'load_mean_std': False, 'temporal_window': 16, 'name': 'tokenizer', 'vae_pth': '/home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth', '_target_': <class 'cosmos_predict2.tokenizers.tokenizer.TokenizerInterface'>}, 'guardrail_config': {'checkpoint_dir': '/home/shadeform/pdt-mimic/mimic-video/model/checkpoints', 'offload_model_to_cpu': True, 'enabled': False}, 'precision': 'bfloat16', 'rectified_flow_t_scaling_factor': 1.0, 'rectified_flow_loss_weight_uniform': True, 'resize_online': False, 'resolution': '480', 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0, '_target_': 'cosmos_predict2.configs.defaults.ema.EMAConfig'}, 'sigma_data': 1.0, 'state_ch': 16, 'state_t': 16, 'text_encoder': {'cls': <TextEncoderClass.T5: 't5'>, 't5': {'ckpt_path': '/home/shadeform/pdt-mimic/mimic-video/model/checkpoints/text_encoder/t5-11b', 'num_tokens': 512, 'embed_dim': 1024}}, 'input_video_key': 'video', 'input_image_key': 'images', 'timestamps': {'nfe': 35, 't_min': 0.002, 't_max': 80.0, 'order': 7.0, 'is_forward': False}}, 'fsdp_shard_size': 0, 'data_config': {'config_name': 'so101_lerobot', '_target_': <function get_data_config at 0x779a85552a70>}, '_target_': 'cosmos_predict2.models.world2action_model.World2ActionModelConfig'}, '_target_': <class 'cosmos_predict2.models.world2action_model.World2ActionModel'>}
+* world2action_pipe: {'precision': 'bfloat16', 'scheduler': {'alpha': 1.0, 'beta': 1.0, 'num_denoising_steps': 10}, 'net': {'max_horizon': 91, 'in_channels': 6, 'out_channels': 6, 'model_channels': 512, 'num_blocks': 12, 'num_heads': 8, 'mlp_ratio': 4.0, 'atten_backend': 'flash_attn_no_cp', 'crossattn_emb_channels': 2048, 'use_adaln_lora': True, 'adaln_lora_dim': 128, 'pair_timestep_feature_rank': 512, 'sac_config': {'mode': 'none', 'every_n_blocks': 1}, '_target_': <class 'cosmos_predict2.models.world2action_dit.World2ActionDIT'>}, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, '_target_': 'cosmos_predict2.configs.config_world2action.World2ActionPipelineConfig'}
+* optimizer: {'optim_type': 'fusedadam', 'model': None, 'lr': 9.999999999999999e-05, 'weight_decay': 0.1, 'betas': [0.9, 0.99], 'eps': 1e-08, 'master_weights': True, 'capturable': True, '_target_': <function get_base_optimizer at 0x779a6ed6b6d0>}
+* scheduler: {'verbosity_interval': 0, 'warm_up_steps': [1000], 'cycle_lengths': [500000], 'f_start': [1e-06], 'f_max': [1], 'f_min': [0.2], '_target_': <class 'cosmos_predict2.functional.lr_scheduler.LambdaLinearScheduler'>}
+* data_config: {'config_name': 'so101_lerobot', '_target_': <function get_data_config at 0x779a85552a70>}
+* video_dataset_train: None
+* video_dataset_val: None
+* dataloader_train: {'batch_size': 2, 'shuffle': None, 'sampler': {'seed': 0, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': None, 'seed': 42, 'train': True, 'cache_dir': None, 'video_backend': 'pyav', 'tolerance_s': 0.04, 'repo_id': 'push-that-thing/task_1', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, 'num_replicas': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_world_size at 0x779cebe89120>}, 'rank': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_rank at 0x779cebe89240>}, 'shuffle': True, '_target_': <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>}, 'batch_sampler': None, 'num_workers': 0, 'collate_fn': None, 'pin_memory': True, 'drop_last': True, 'timeout': 0, 'worker_init_fn': None, 'multiprocessing_context': None, 'generator': None, 'prefetch_factor': None, 'persistent_workers': False, 'pin_memory_device': '', 'in_order': True, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': 0.05, 'seed': 42, 'train': True, 'cache_dir': None, 'video_backend': 'pyav', 'tolerance_s': 0.04, 'repo_id': 'push-that-thing/task_1_slow', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, '_target_': <class 'torch.utils.data.dataloader.DataLoader'>}
+* dataloader_val: {'batch_size': 1, 'shuffle': None, 'sampler': {'seed': 0, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': None, 'seed': 42, 'train': False, 'cache_dir': None, 'video_backend': 'pyav', 'tolerance_s': 0.04, 'repo_id': 'push-that-thing/task_1', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, 'num_replicas': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_world_size at 0x779cebe89120>}, 'rank': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_rank at 0x779cebe89240>}, 'shuffle': False, '_target_': <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>}, 'batch_sampler': None, 'num_workers': 0, 'collate_fn': None, 'pin_memory': False, 'drop_last': False, 'timeout': 0, 'worker_init_fn': None, 'multiprocessing_context': None, 'generator': None, 'prefetch_factor': None, 'persistent_workers': False, 'pin_memory_device': '', 'in_order': True, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': 0.05, 'seed': 42, 'train': False, 'cache_dir': None, 'video_backend': 'pyav', 'tolerance_s': 0.04, 'repo_id': 'push-that-thing/task_1_slow', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, '_target_': <class 'torch.utils.data.dataloader.DataLoader'>}
+* job:
+   * project: vam
+   * group: so101
+   * name: w2a_so101_task1_slow_100ep_lora32_iter2500_75m_2gpu
+* trainer:
+   * type: <class 'imaginaire.trainer.ImaginaireTrainer'>
+   * callbacks: {'ema': {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.EMAModelCallback'>}, 'progress_bar': {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.ProgressBarCallback'>}, 'low_prec': {'config': None, 'trainer': None, 'update_iter': 1, '_target_': <class 'imaginaire.utils.callback.LowPrecisionCallback'>}, 'iter_speed': {'hit_thres': 5, 'every_n': 1000, '_target_': <class 'cosmos_predict2.callbacks.iter_speed.IterSpeed'>}, 'device_monitor': {'every_n': 1000, 'step_size': 1, 'log_memory_detail': True, '_target_': <class 'cosmos_predict2.callbacks.device_monitor.DeviceMonitor'>}, 'manual_gc': {'warm_up': 5, 'every_n': 5, '_target_': <class 'imaginaire.callbacks.manual_gc.ManualGarbageCollection'>}, 'grad_clip': {'clip_norm': 10.0, 'force_finite': True, 'log_wandb': False, '_target_': <class 'cosmos_predict2.callbacks.grad_clip.GradClip'>}, 'wandb': {'mode': None, '_target_': <class 'cosmos_predict2.callbacks.wandb_logger.WandBLogger'>}}
+   * distributed_parallelism: ddp
+   * ddp:
+      * find_unused_parameters: False
+      * static_graph: True
+      * broadcast_buffers: True
+   * cudnn:
+      * deterministic: False
+      * benchmark: True
+   * seed: 0
+   * grad_scaler_args: {'enabled': False}
+   * max_iter: 8000
+   * max_val_iter: 16
+   * logging_iter: 1000
+   * run_validation: True
+   * validation_iter: 100
+   * timeout_period: 999999999
+   * memory_format: torch.preserve_format
+   * grad_accum_iter: 8
+   * profiling:
+      * enable_profiling: False
+      * enable_memory_snapshot: False
+      * profile_freq: 1
+      * first_n_rank: 4
+      * record_shape: True
+      * profile_memory: True
+      * with_stack: True
+      * with_modules: True
+* model_parallel: ModelParallelConfig(tensor_model_parallel_size=1, pipeline_model_parallel_comm_backend=None, pipeline_model_parallel_size=1, virtual_pipeline_model_parallel_size=None, sequence_parallel=False, context_parallel_size=1, hierarchical_context_parallel_sizes=None, expert_model_parallel_size=1, expert_tensor_parallel_size=1, moe_extended_tp=False, perform_initialization=True, use_cpu_initialization=False, fp16=False, bf16=False, params_dtype=torch.float32, timers=None, finalize_model_grads_func=None, grad_scale_func=None, no_sync_func=None, grad_sync_func=None, param_sync_func=None, deterministic_mode=False, enable_autocast=False, autocast_dtype=torch.float32, num_microbatches_with_partial_activation_checkpoints=None, gradient_accumulation_fusion=False, async_tensor_model_parallel_allreduce=False, use_te_rng_tracker=False, tp_comm_overlap=False, tp_comm_bulk_wgrad=True, tp_comm_bulk_dgrad=True, tp_comm_overlap_ag=True, tp_comm_overlap_rs=True, tp_comm_overlap_rs_dgrad=False, tp_comm_split_ag=True, tp_comm_atomic_ag=False, tp_comm_split_rs=True, tp_comm_atomic_rs=False, cross_entropy_loss_fusion=False, cross_entropy_fusion_impl='native', tp_comm_overlap_disable_qkv=False, tp_comm_overlap_disable_fc1=False, tp_comm_bootstrap_backend='nccl', pipeline_dtype=None, variable_seq_lengths=False, overlap_p2p_comm=False, batch_p2p_comm=True, batch_p2p_sync=True, use_ring_exchange_p2p=False, deallocate_pipeline_outputs=False, defer_embedding_wgrad_compute=False, wgrad_deferral_limit=0, pipeline_model_parallel_split_rank=None, overlap_p2p_comm_warmup_flush=False, microbatch_group_size_per_vp_stage=1, cpu_offloading=False, cpu_offloading_num_layers=0, _cpu_offloading_context=None, cpu_offloading_activations=False, cpu_offloading_weights=False, barrier_with_L1_time=True)
+* checkpoint:
+   * type: {'callbacks': None, '_target_': <class 'cosmos_predict2.checkpointer.Checkpointer'>}
+   * dcp_async_mode_enabled: False
+   * save_iter: 100
+   * load_path: 
+   * load_training_state: False
+   * only_load_scheduler_state: False
+   * strict_resume: True
+   * jit:
+      * enabled: False
+      * input_shape: None
+      * device: cuda
+      * dtype: bfloat16
+      * strict: True
+   * verbose: True
+   * keys_not_to_resume: []
+   * broadcast_via_filesystem: False
+   * load_ema_to_reg: False
+   * dcp_allow_mismatched_size: False
+* defaults: ['_self_', {'data_config': None}, {'video_dataset_train': None}, {'video_dataset_val': None}, {'dataloader_train': None}, {'dataloader_val': None}, {'world2action_pipe': None}, {'optimizer': 'fusedadamw'}, {'scheduler': 'constant'}, {'model': None}, {'callbacks': ['basic']}, {'net': None}, {'ema': None}, {'checkpoint': None}, {'ckpt_type': None}, {'experiment': None}]
+[05-20 11:48:57|WARNING|imaginaire/utils/misc.py:127:print_environ_variables] Environment variable TORCH_HOME not set!
+[05-20 11:48:57|WARNING|imaginaire/utils/misc.py:127:print_environ_variables] Environment variable IMAGINAIRE_OUTPUT_ROOT not set!
+[05-20 11:48:57|INFO|imaginaire/utils/misc.py:139:set_random_seed] Using random seed 0.
+[05-20 11:48:57|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback ema: {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.EMAModelCallback'>}
+[05-20 11:48:57|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback progress_bar: {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.ProgressBarCallback'>}
+[05-20 11:48:57|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback low_prec: {'config': None, 'trainer': None, 'update_iter': 1, '_target_': <class 'imaginaire.utils.callback.LowPrecisionCallback'>}
+[05-20 11:48:57|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback iter_speed: {'hit_thres': 5, 'every_n': 1000, '_target_': <class 'cosmos_predict2.callbacks.iter_speed.IterSpeed'>}
+[05-20 11:48:57|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback device_monitor: {'every_n': 1000, 'step_size': 1, 'log_memory_detail': True, '_target_': <class 'cosmos_predict2.callbacks.device_monitor.DeviceMonitor'>}
+[05-20 11:48:57|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback manual_gc: {'warm_up': 5, 'every_n': 5, '_target_': <class 'imaginaire.callbacks.manual_gc.ManualGarbageCollection'>}
+[05-20 11:48:57|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback grad_clip: {'clip_norm': 10.0, 'force_finite': True, 'log_wandb': False, '_target_': <class 'cosmos_predict2.callbacks.grad_clip.GradClip'>}
+[05-20 11:48:57|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback wandb: {'mode': None, '_target_': <class 'cosmos_predict2.callbacks.wandb_logger.WandBLogger'>}
+[05-20 11:48:59|CRITICAL|cosmos_predict2/models/world2action_model.py:106:__init__] Using mean loss reduce with loss scale 10.0
+[05-20 11:48:59|WARNING|cosmos_predict2/pipelines/world2action.py:71:from_config] precision torch.bfloat16
+[05-20 11:48:59|WARNING|cosmos_predict2/pipelines/world2action.py:101:from_config] dit_path not provided, initializing DiT with random weights
+[05-20 11:48:59|WARNING|cosmos_predict2/pipelines/video2world.py:277:from_config] precision torch.bfloat16
+[05-20 11:49:00|INFO|cosmos_predict2/tokenizers/tokenizer.py:687:_video_vae] Loading /home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth
+[05-20 11:49:00|SUCCESS|cosmos_predict2/tokenizers/tokenizer.py:689:_video_vae] Successfully loaded /home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth
+[05-20 11:49:00|INFO|imaginaire/utils/distributed.py:430:sync_model_states] Synchronizing model states from rank 0 to all ranks in process group [0, 1].
+[05-20 11:49:00|INFO|cosmos_predict2/pipelines/video2world.py:335:from_config] Loading DiT from /home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt
+[05-20 11:49:04|SUCCESS|cosmos_predict2/pipelines/video2world.py:354:from_config] Successfully loaded DiT from /home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt
+[05-20 11:49:05|INFO|cosmos_predict2/models/world2action_model.py:148:__init__] Total parameters: 2.04B, Frozen parameters: 1,956,413,440, Trainable parameters: 80,859,000
+[05-20 11:49:05|INFO|cosmos_predict2/models/world2action_model.py:167:__init__] FSDP (Fully Sharded Data Parallel) is disabled.
+[05-20 11:49:05|CRITICAL|cosmos_predict2/utils/optim_instantiate_dtensor.py:49:get_base_optimizer] total num parameters : 80,859,000
+[05-20 11:49:05|WARNING|cosmos_predict2/utils/fused_adam_dtensor.py:103:__init__] FusedAdam master_weights: True capturable: True
+[05-20 11:49:05|INFO|cosmos_predict2/checkpointer.py:288:load] Training from scratch.
+[05-20 11:49:05|INFO|imaginaire/utils/misc.py:266:wrapper] Time spent on checkpoint loading: 0.0002 seconds
+[05-20 11:49:05|CRITICAL|imaginaire/trainer.py:179:train] Distributed parallelism mode: ddp
+[05-20 11:49:06|INFO|imaginaire/trainer.py:187:train] Starting training...
+[05-20 11:49:06|INFO|cosmos_predict2/callbacks/device_monitor.py:92:on_train_start] DeviceMonitor callback: local_dir: checkpoints/vam/so101/w2a_so101_task1_slow_100ep_lora32_iter2500_75m_2gpu/DeviceMonitor
+[05-20 11:49:08|INFO|imaginaire/trainer.py:388:validate] Validating at iteration 0...
+[05-20 11:49:49|INFO|imaginaire/trainer.py:193:train] Initial validation done.
+[05-20 11:49:51|WARNING|imaginaire/utils/distributed.py:284:ddp_sync_grad] DDP static_graph=True is incompatible with sync_grad(). Performance will be reduced.
+[05-20 11:51:40|INFO|imaginaire/trainer.py:116:__init__] Config:
+* model: {'config': {'validation_mode': 'loss_only', 'train_architecture': 'base', 'lora_rank': 16, 'lora_alpha': 16, 'lora_target_modules': 'q_proj,k_proj,v_proj,output_proj,mlp.layer1,mlp.layer2', 'init_lora_weights': True, 'precision': 'bfloat16', 'loss_reduce': 'mean', 'loss_scale': 10.0, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, 'action_dit_path': '', 'video_dit_path': '/home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt', 'pipe_config': {'precision': 'bfloat16', 'scheduler': {'alpha': 1.0, 'beta': 1.0, 'num_denoising_steps': 10}, 'net': {'max_horizon': 91, 'in_channels': 6, 'out_channels': 6, 'model_channels': 512, 'num_blocks': 12, 'num_heads': 8, 'mlp_ratio': 4.0, 'atten_backend': 'flash_attn_no_cp', 'crossattn_emb_channels': 2048, 'use_adaln_lora': True, 'adaln_lora_dim': 128, 'pair_timestep_feature_rank': 512, 'sac_config': {'mode': 'none', 'every_n_blocks': 1}, '_target_': <class 'cosmos_predict2.models.world2action_dit.World2ActionDIT'>}, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, '_target_': 'cosmos_predict2.configs.config_world2action.World2ActionPipelineConfig', 'xattn_layer_idx': 20}, 'video_pipe_config': {'adjust_video_noise': True, 'conditioner': {'fps': {'output_key': 'fps', 'dropout_rate': 0.0, 'dtype': None, 'input_key': 'fps', '_target_': <class 'cosmos_predict2.conditioner.ReMapkey'>}, 'padding_mask': {'output_key': 'padding_mask', 'dropout_rate': 0.0, 'dtype': None, 'input_key': 'padding_mask', '_target_': <class 'cosmos_predict2.conditioner.ReMapkey'>}, 'text': {'dropout_rate': 0.0, 'input_key': ['obs/language_embedding'], '_target_': <class 'cosmos_predict2.conditioner.TextAttr'>}, 'use_video_condition': {'output_key': 'use_video_condition', 'dropout_rate': 0.0, 'input_key': 'fps', '_target_': <class 'cosmos_predict2.conditioner.BooleanFlag'>}, '_target_': <class 'cosmos_predict2.conditioner.VideoConditioner'>}, 'conditioning_strategy': 'frame_replace', 'min_num_conditional_frames': 1, 'max_num_conditional_frames': 2, 'sigma_conditional': 0.0001, 'net': {'max_img_h': 240, 'max_img_w': 240, 'max_frames': 128, 'in_channels': 16, 'out_channels': 16, 'patch_spatial': 2, 'patch_temporal': 1, 'concat_padding_mask': True, 'model_channels': 2048, 'num_blocks': 28, 'num_heads': 16, 'atten_backend': 'minimal_a2a', 'pos_emb_cls': 'rope3d', 'pos_emb_learnable': True, 'pos_emb_interpolation': 'crop', 'use_adaln_lora': True, 'adaln_lora_dim': 256, 'rope_h_extrapolation_ratio': 3.0, 'rope_w_extrapolation_ratio': 3.0, 'rope_t_extrapolation_ratio': 1.0, 'extra_per_block_abs_pos_emb': False, 'rope_enable_fps_modulation': False, 'sac_config': {'mode': 'predict2_2b_720', 'every_n_blocks': 1, '_target_': 'cosmos_predict2.models.text2image_dit.SACConfig'}, '_target_': <class 'cosmos_predict2.models.video2world_dit.MinimalV1LVGDiT'>}, 'tokenizer': {'chunk_duration': 81, 'load_mean_std': False, 'temporal_window': 16, 'name': 'tokenizer', 'vae_pth': '/home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth', '_target_': <class 'cosmos_predict2.tokenizers.tokenizer.TokenizerInterface'>}, 'guardrail_config': {'checkpoint_dir': '/home/shadeform/pdt-mimic/mimic-video/model/checkpoints', 'offload_model_to_cpu': True, 'enabled': False}, 'precision': 'bfloat16', 'rectified_flow_t_scaling_factor': 1.0, 'rectified_flow_loss_weight_uniform': True, 'resize_online': False, 'resolution': '480', 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0, '_target_': 'cosmos_predict2.configs.defaults.ema.EMAConfig'}, 'sigma_data': 1.0, 'state_ch': 16, 'state_t': 16, 'text_encoder': {'cls': <TextEncoderClass.T5: 't5'>, 't5': {'ckpt_path': '/home/shadeform/pdt-mimic/mimic-video/model/checkpoints/text_encoder/t5-11b', 'num_tokens': 512, 'embed_dim': 1024}}, 'input_video_key': 'video', 'input_image_key': 'images', 'timestamps': {'nfe': 35, 't_min': 0.002, 't_max': 80.0, 'order': 7.0, 'is_forward': False}}, 'fsdp_shard_size': 0, 'data_config': {'config_name': 'so101_lerobot', '_target_': <function get_data_config at 0x7b8b98c5a950>}, '_target_': 'cosmos_predict2.models.world2action_model.World2ActionModelConfig'}, '_target_': <class 'cosmos_predict2.models.world2action_model.World2ActionModel'>}
+* world2action_pipe: {'precision': 'bfloat16', 'scheduler': {'alpha': 1.0, 'beta': 1.0, 'num_denoising_steps': 10}, 'net': {'max_horizon': 91, 'in_channels': 6, 'out_channels': 6, 'model_channels': 512, 'num_blocks': 12, 'num_heads': 8, 'mlp_ratio': 4.0, 'atten_backend': 'flash_attn_no_cp', 'crossattn_emb_channels': 2048, 'use_adaln_lora': True, 'adaln_lora_dim': 128, 'pair_timestep_feature_rank': 512, 'sac_config': {'mode': 'none', 'every_n_blocks': 1}, '_target_': <class 'cosmos_predict2.models.world2action_dit.World2ActionDIT'>}, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, '_target_': 'cosmos_predict2.configs.config_world2action.World2ActionPipelineConfig'}
+* optimizer: {'optim_type': 'fusedadam', 'model': None, 'lr': 9.999999999999999e-05, 'weight_decay': 0.1, 'betas': [0.9, 0.99], 'eps': 1e-08, 'master_weights': True, 'capturable': True, '_target_': <function get_base_optimizer at 0x7b8b8276f6d0>}
+* scheduler: {'verbosity_interval': 0, 'warm_up_steps': [1000], 'cycle_lengths': [500000], 'f_start': [1e-06], 'f_max': [1], 'f_min': [0.2], '_target_': <class 'cosmos_predict2.functional.lr_scheduler.LambdaLinearScheduler'>}
+* data_config: {'config_name': 'so101_lerobot', '_target_': <function get_data_config at 0x7b8b98c5a950>}
+* video_dataset_train: None
+* video_dataset_val: None
+* dataloader_train: {'batch_size': 2, 'shuffle': None, 'sampler': {'seed': 0, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': None, 'seed': 42, 'train': True, 'cache_dir': None, 'video_backend': 'pyav', 'tolerance_s': 0.04, 'repo_id': 'push-that-thing/task_1', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, 'num_replicas': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_world_size at 0x7b8dff895120>}, 'rank': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_rank at 0x7b8dff895240>}, 'shuffle': True, '_target_': <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>}, 'batch_sampler': None, 'num_workers': 0, 'collate_fn': None, 'pin_memory': True, 'drop_last': True, 'timeout': 0, 'worker_init_fn': None, 'multiprocessing_context': None, 'generator': None, 'prefetch_factor': None, 'persistent_workers': False, 'pin_memory_device': '', 'in_order': True, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': 0.05, 'seed': 42, 'train': True, 'cache_dir': None, 'video_backend': 'pyav', 'tolerance_s': 0.04, 'repo_id': 'push-that-thing/task_1_slow', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, '_target_': <class 'torch.utils.data.dataloader.DataLoader'>}
+* dataloader_val: {'batch_size': 1, 'shuffle': None, 'sampler': {'seed': 0, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': None, 'seed': 42, 'train': False, 'cache_dir': None, 'video_backend': 'pyav', 'tolerance_s': 0.04, 'repo_id': 'push-that-thing/task_1', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, 'num_replicas': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_world_size at 0x7b8dff895120>}, 'rank': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_rank at 0x7b8dff895240>}, 'shuffle': False, '_target_': <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>}, 'batch_sampler': None, 'num_workers': 0, 'collate_fn': None, 'pin_memory': False, 'drop_last': False, 'timeout': 0, 'worker_init_fn': None, 'multiprocessing_context': None, 'generator': None, 'prefetch_factor': None, 'persistent_workers': False, 'pin_memory_device': '', 'in_order': True, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': 0.05, 'seed': 42, 'train': False, 'cache_dir': None, 'video_backend': 'pyav', 'tolerance_s': 0.04, 'repo_id': 'push-that-thing/task_1_slow', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, '_target_': <class 'torch.utils.data.dataloader.DataLoader'>}
+* job:
+   * project: vam
+   * group: so101
+   * name: w2a_so101_task1_slow_100ep_lora32_iter2500_75m_2gpu
+* trainer:
+   * type: <class 'imaginaire.trainer.ImaginaireTrainer'>
+   * callbacks: {'ema': {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.EMAModelCallback'>}, 'progress_bar': {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.ProgressBarCallback'>}, 'low_prec': {'config': None, 'trainer': None, 'update_iter': 1, '_target_': <class 'imaginaire.utils.callback.LowPrecisionCallback'>}, 'iter_speed': {'hit_thres': 5, 'every_n': 1000, '_target_': <class 'cosmos_predict2.callbacks.iter_speed.IterSpeed'>}, 'device_monitor': {'every_n': 1000, 'step_size': 1, 'log_memory_detail': True, '_target_': <class 'cosmos_predict2.callbacks.device_monitor.DeviceMonitor'>}, 'manual_gc': {'warm_up': 5, 'every_n': 5, '_target_': <class 'imaginaire.callbacks.manual_gc.ManualGarbageCollection'>}, 'grad_clip': {'clip_norm': 10.0, 'force_finite': True, 'log_wandb': False, '_target_': <class 'cosmos_predict2.callbacks.grad_clip.GradClip'>}, 'wandb': {'mode': None, '_target_': <class 'cosmos_predict2.callbacks.wandb_logger.WandBLogger'>}}
+   * distributed_parallelism: ddp
+   * ddp:
+      * find_unused_parameters: False
+      * static_graph: True
+      * broadcast_buffers: True
+   * cudnn:
+      * deterministic: False
+      * benchmark: True
+   * seed: 0
+   * grad_scaler_args: {'enabled': False}
+   * max_iter: 8000
+   * max_val_iter: 16
+   * logging_iter: 1000
+   * run_validation: True
+   * validation_iter: 100
+   * timeout_period: 999999999
+   * memory_format: torch.preserve_format
+   * grad_accum_iter: 8
+   * profiling:
+      * enable_profiling: False
+      * enable_memory_snapshot: False
+      * profile_freq: 1
+      * first_n_rank: 4
+      * record_shape: True
+      * profile_memory: True
+      * with_stack: True
+      * with_modules: True
+* model_parallel: ModelParallelConfig(tensor_model_parallel_size=1, pipeline_model_parallel_comm_backend=None, pipeline_model_parallel_size=1, virtual_pipeline_model_parallel_size=None, sequence_parallel=False, context_parallel_size=1, hierarchical_context_parallel_sizes=None, expert_model_parallel_size=1, expert_tensor_parallel_size=1, moe_extended_tp=False, perform_initialization=True, use_cpu_initialization=False, fp16=False, bf16=False, params_dtype=torch.float32, timers=None, finalize_model_grads_func=None, grad_scale_func=None, no_sync_func=None, grad_sync_func=None, param_sync_func=None, deterministic_mode=False, enable_autocast=False, autocast_dtype=torch.float32, num_microbatches_with_partial_activation_checkpoints=None, gradient_accumulation_fusion=False, async_tensor_model_parallel_allreduce=False, use_te_rng_tracker=False, tp_comm_overlap=False, tp_comm_bulk_wgrad=True, tp_comm_bulk_dgrad=True, tp_comm_overlap_ag=True, tp_comm_overlap_rs=True, tp_comm_overlap_rs_dgrad=False, tp_comm_split_ag=True, tp_comm_atomic_ag=False, tp_comm_split_rs=True, tp_comm_atomic_rs=False, cross_entropy_loss_fusion=False, cross_entropy_fusion_impl='native', tp_comm_overlap_disable_qkv=False, tp_comm_overlap_disable_fc1=False, tp_comm_bootstrap_backend='nccl', pipeline_dtype=None, variable_seq_lengths=False, overlap_p2p_comm=False, batch_p2p_comm=True, batch_p2p_sync=True, use_ring_exchange_p2p=False, deallocate_pipeline_outputs=False, defer_embedding_wgrad_compute=False, wgrad_deferral_limit=0, pipeline_model_parallel_split_rank=None, overlap_p2p_comm_warmup_flush=False, microbatch_group_size_per_vp_stage=1, cpu_offloading=False, cpu_offloading_num_layers=0, _cpu_offloading_context=None, cpu_offloading_activations=False, cpu_offloading_weights=False, barrier_with_L1_time=True)
+* checkpoint:
+   * type: {'callbacks': None, '_target_': <class 'cosmos_predict2.checkpointer.Checkpointer'>}
+   * dcp_async_mode_enabled: False
+   * save_iter: 100
+   * load_path: 
+   * load_training_state: False
+   * only_load_scheduler_state: False
+   * strict_resume: True
+   * jit:
+      * enabled: False
+      * input_shape: None
+      * device: cuda
+      * dtype: bfloat16
+      * strict: True
+   * verbose: True
+   * keys_not_to_resume: []
+   * broadcast_via_filesystem: False
+   * load_ema_to_reg: False
+   * dcp_allow_mismatched_size: False
+* defaults: ['_self_', {'data_config': None}, {'video_dataset_train': None}, {'video_dataset_val': None}, {'dataloader_train': None}, {'dataloader_val': None}, {'world2action_pipe': None}, {'optimizer': 'fusedadamw'}, {'scheduler': 'constant'}, {'model': None}, {'callbacks': ['basic']}, {'net': None}, {'ema': None}, {'checkpoint': None}, {'ckpt_type': None}, {'experiment': None}]
+[05-20 11:51:40|WARNING|imaginaire/utils/misc.py:127:print_environ_variables] Environment variable TORCH_HOME not set!
+[05-20 11:51:40|WARNING|imaginaire/utils/misc.py:127:print_environ_variables] Environment variable IMAGINAIRE_OUTPUT_ROOT not set!
+[05-20 11:51:40|INFO|imaginaire/utils/misc.py:139:set_random_seed] Using random seed 0.
+[05-20 11:51:40|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback ema: {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.EMAModelCallback'>}
+[05-20 11:51:40|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback progress_bar: {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.ProgressBarCallback'>}
+[05-20 11:51:40|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback low_prec: {'config': None, 'trainer': None, 'update_iter': 1, '_target_': <class 'imaginaire.utils.callback.LowPrecisionCallback'>}
+[05-20 11:51:40|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback iter_speed: {'hit_thres': 5, 'every_n': 1000, '_target_': <class 'cosmos_predict2.callbacks.iter_speed.IterSpeed'>}
+[05-20 11:51:40|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback device_monitor: {'every_n': 1000, 'step_size': 1, 'log_memory_detail': True, '_target_': <class 'cosmos_predict2.callbacks.device_monitor.DeviceMonitor'>}
+[05-20 11:51:40|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback manual_gc: {'warm_up': 5, 'every_n': 5, '_target_': <class 'imaginaire.callbacks.manual_gc.ManualGarbageCollection'>}
+[05-20 11:51:40|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback grad_clip: {'clip_norm': 10.0, 'force_finite': True, 'log_wandb': False, '_target_': <class 'cosmos_predict2.callbacks.grad_clip.GradClip'>}
+[05-20 11:51:40|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback wandb: {'mode': None, '_target_': <class 'cosmos_predict2.callbacks.wandb_logger.WandBLogger'>}
+[05-20 11:51:42|CRITICAL|cosmos_predict2/models/world2action_model.py:106:__init__] Using mean loss reduce with loss scale 10.0
+[05-20 11:51:42|WARNING|cosmos_predict2/pipelines/world2action.py:71:from_config] precision torch.bfloat16
+[05-20 11:51:42|WARNING|cosmos_predict2/pipelines/world2action.py:101:from_config] dit_path not provided, initializing DiT with random weights
+[05-20 11:51:42|WARNING|cosmos_predict2/pipelines/video2world.py:277:from_config] precision torch.bfloat16
+[05-20 11:51:43|INFO|cosmos_predict2/tokenizers/tokenizer.py:687:_video_vae] Loading /home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth
+[05-20 11:51:43|SUCCESS|cosmos_predict2/tokenizers/tokenizer.py:689:_video_vae] Successfully loaded /home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth
+[05-20 11:51:43|INFO|imaginaire/utils/distributed.py:430:sync_model_states] Synchronizing model states from rank 0 to all ranks in process group [0, 1].
+[05-20 11:51:43|INFO|cosmos_predict2/pipelines/video2world.py:335:from_config] Loading DiT from /home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt
+[05-20 11:51:48|SUCCESS|cosmos_predict2/pipelines/video2world.py:354:from_config] Successfully loaded DiT from /home/shadeform/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt
+[05-20 11:51:48|INFO|cosmos_predict2/models/world2action_model.py:148:__init__] Total parameters: 2.04B, Frozen parameters: 1,956,413,440, Trainable parameters: 80,859,000
+[05-20 11:51:48|INFO|cosmos_predict2/models/world2action_model.py:167:__init__] FSDP (Fully Sharded Data Parallel) is disabled.
+[05-20 11:51:49|CRITICAL|cosmos_predict2/utils/optim_instantiate_dtensor.py:49:get_base_optimizer] total num parameters : 80,859,000
+[05-20 11:51:49|WARNING|cosmos_predict2/utils/fused_adam_dtensor.py:103:__init__] FusedAdam master_weights: True capturable: True
+[05-20 11:51:49|INFO|cosmos_predict2/checkpointer.py:288:load] Training from scratch.
+[05-20 11:51:49|INFO|imaginaire/utils/misc.py:266:wrapper] Time spent on checkpoint loading: 0.0003 seconds
+[05-20 11:51:49|CRITICAL|imaginaire/trainer.py:179:train] Distributed parallelism mode: ddp
+[05-20 11:51:49|INFO|imaginaire/trainer.py:187:train] Starting training...
+[05-20 11:51:49|INFO|cosmos_predict2/callbacks/device_monitor.py:92:on_train_start] DeviceMonitor callback: local_dir: checkpoints/vam/so101/w2a_so101_task1_slow_100ep_lora32_iter2500_75m_2gpu/DeviceMonitor
+[05-20 11:51:53|INFO|imaginaire/trainer.py:388:validate] Validating at iteration 0...
+[05-20 11:52:36|INFO|imaginaire/trainer.py:193:train] Initial validation done.
+[05-20 11:52:38|WARNING|imaginaire/utils/distributed.py:284:ddp_sync_grad] DDP static_graph=True is incompatible with sync_grad(). Performance will be reduced.