Upload folder using huggingface_hub

config.json

@@ -0,0 +1,70 @@
+{
+  "action_horizon": 50,
+  "add_pos_embed": true,
+  "apply_sincos_state_encoding": true,
+  "architectures": [
+    "Gr00tN1d6"
+  ],
+  "attn_dropout": 0.2,
+  "attn_implementation": null,
+  "backbone_embedding_dim": 2048,
+  "backbone_model_type": "eagle",
+  "backbone_trainable_params_fp32": true,
+  "collator_overwrite_image_inputs": false,
+  "color_jitter_params": {
+    "brightness": 0.1,
+    "contrast": 0.1,
+    "hue": 0.1,
+    "saturation": 0.1
+  },
+  "crop_fraction": 0.95,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "eagle_collator": true,
+  "formalize_language": true,
+  "gemma_collator": false,
+  "hidden_size": 1024,
+  "image_crop_size": null,
+  "image_target_size": null,
+  "input_embedding_dim": 1536,
+  "load_bf16": true,
+  "max_action_dim": 128,
+  "max_num_embodiments": 32,
+  "max_seq_len": 1024,
+  "max_state_dim": 128,
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "model_type": "Gr00tN1d6",
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_inference_timesteps": 4,
+  "num_timestep_buckets": 1000,
+  "random_rotation_angle": null,
+  "reproject_vision": false,
+  "select_layer": 16,
+  "shortest_image_edge": 256,
+  "state_dropout_prob": 0.0,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "tune_diffusion_model": true,
+  "tune_llm": false,
+  "tune_projector": true,
+  "tune_top_llm_layers": 4,
+  "tune_visual": false,
+  "tune_vlln": true,
+  "use_albumentations_transforms": true,
+  "use_alternate_vl_dit": true,
+  "use_flash_attention": true,
+  "use_relative_action": true,
+  "use_vlln": true
+}

embodiment_id.json

@@ -0,0 +1,11 @@
+{
+  "robocasa_panda_omron": 13,
+  "gr1": 20,
+  "behavior_r1_pro": 24,
+  "unitree_g1": 8,
+  "oxe_google": 0,
+  "oxe_widowx": 1,
+  "libero_panda": 2,
+  "oxe_droid": 16,
+  "new_embodiment": 10
+}

experiment_cfg/conf.yaml

@@ -0,0 +1,210 @@
+load_config_path: null
+model:
+  model_type: Gr00tN1d6
+  model_dtype: bfloat16
+  model_name: nvidia/Eagle-Block2A-2B-v2
+  backbone_model_type: eagle
+  model_revision: null
+  tune_top_llm_layers: 4
+  backbone_embedding_dim: 2048
+  tune_llm: false
+  tune_visual: false
+  select_layer: 16
+  reproject_vision: false
+  use_flash_attention: true
+  load_bf16: false
+  collator_overwrite_image_inputs: false
+  eagle_collator: true
+  backbone_trainable_params_fp32: true
+  image_crop_size: null
+  image_target_size: null
+  shortest_image_edge: 256
+  crop_fraction: 0.95
+  random_rotation_angle: null
+  color_jitter_params:
+    brightness: 0.3
+    contrast: 0.4
+    saturation: 0.5
+    hue: 0.08
+  use_albumentations_transforms: true
+  formalize_language: true
+  apply_sincos_state_encoding: false
+  use_relative_action: true
+  max_state_dim: 29
+  max_action_dim: 29
+  action_horizon: 16
+  hidden_size: 1024
+  input_embedding_dim: 1536
+  add_pos_embed: true
+  attn_dropout: 0.2
+  use_vlln: true
+  max_seq_len: 1024
+  use_alternate_vl_dit: true
+  attend_text_every_n_blocks: 2
+  diffusion_model_cfg:
+    positional_embeddings: null
+    num_layers: 32
+    num_attention_heads: 32
+    attention_head_dim: 48
+    norm_type: ada_norm
+    dropout: 0.2
+    final_dropout: true
+    output_dim: 1024
+    interleave_self_attention: true
+  num_inference_timesteps: 4
+  noise_beta_alpha: 1.5
+  noise_beta_beta: 1.0
+  noise_s: 0.999
+  num_timestep_buckets: 1000
+  tune_projector: true
+  tune_diffusion_model: true
+  tune_vlln: true
+  state_dropout_prob: 0.0
+  state_additive_noise_scale: 0.0
+  max_num_embodiments: 32
+data:
+  datasets:
+  - dataset_paths:
+    - /workspaces/isaaclab_arena/train_data
+    embodiment_tag: new_embodiment
+    mix_ratio: 1.0
+    dataset_type: physical_embodiment
+    val_dataset_path: null
+  modality_configs:
+    new_embodiment:
+      video:
+        delta_indices:
+        - 0
+        modality_keys:
+        - ego_view
+        - left_view
+        - right_view
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - single_arm
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      action:
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        modality_keys:
+        - single_arm
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs:
+        - rep: RELATIVE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: single_arm
+        - rep: ABSOLUTE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: null
+      language:
+        delta_indices:
+        - 0
+        modality_keys:
+        - annotation.human.task_description
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+  download_cache: false
+  shard_size: 1024
+  episode_sampling_rate: 0.1
+  num_shards_per_epoch: 100000
+  override_pretraining_statistics: false
+  mode: single_turn
+  random_chop: 0.0
+  mock_dataset_mode: false
+  shuffle: true
+  seed: 42
+  multiprocessing_context: fork
+  allow_padding: false
+  subsample_ratio: 1.0
+  image_crop_size:
+  - 244
+  - 244
+  image_target_size:
+  - 224
+  - 224
+  video_backend: torchcodec
+training:
+  output_dir: /workspaces/isaaclab_arena/models
+  experiment_name: null
+  max_steps: 10000
+  global_batch_size: 32
+  batch_size: null
+  gradient_accumulation_steps: 1
+  learning_rate: 0.0001
+  lr_scheduler_type: cosine
+  weight_decay: 1.0e-05
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  max_grad_norm: 1.0
+  optim: adamw_torch
+  start_from_checkpoint: nvidia/GR00T-N1.6-3B
+  tf32: true
+  fp16: false
+  bf16: true
+  eval_bf16: true
+  logging_steps: 10
+  save_steps: 200
+  save_total_limit: 5
+  save_vl_model: false
+  upload_checkpoints: false
+  upload_every: 1000
+  upload_last_n_checkpoints: 5
+  max_concurrent_uploads: 2
+  eval_strategy: 'no'
+  eval_steps: 500
+  eval_set_split_ratio: 0.1
+  eval_batch_size: 2
+  save_best_eval_metric_name: ''
+  save_best_eval_metric_greater_is_better: true
+  deepspeed_stage: 2
+  gradient_checkpointing: false
+  transformers_trust_remote_code: true
+  transformers_local_files_only: false
+  transformers_cache_dir: null
+  transformers_access_token: null
+  use_ddp: false
+  ddp_bucket_cap_mb: 100
+  num_gpus: 4
+  dataloader_num_workers: 32
+  remove_unused_columns: false
+  use_wandb: true
+  wandb_project: finetune-gr00t-n1d6
+  enable_profiling: false
+  max_retries: 3
+  assert_loss_less_than: null
+  add_rl_callback: false
+  enable_open_loop_eval: false
+  open_loop_eval_traj_ids:
+  - 0
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_plot_indices: null
+max_steps: 10000
+save_steps: 200

experiment_cfg/config.yaml

@@ -0,0 +1,244 @@
+!!python/object:gr00t.configs.base_config.Config
+data: !!python/object:gr00t.configs.data.data_config.DataConfig
+  allow_padding: false
+  datasets:
+  - !!python/object:gr00t.configs.data.data_config.SingleDatasetConfig
+    dataset_paths:
+    - /workspaces/isaaclab_arena/train_data
+    dataset_type: physical_embodiment
+    embodiment_tag: new_embodiment
+    mix_ratio: 1.0
+    val_dataset_path: null
+  download_cache: false
+  episode_sampling_rate: 0.1
+  image_crop_size:
+  - 244
+  - 244
+  image_target_size:
+  - 224
+  - 224
+  mock_dataset_mode: false
+  modality_configs:
+    new_embodiment:
+      action: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs:
+        - !!python/object:gr00t.data.types.ActionConfig
+          format: &id001 !!python/object/apply:gr00t.data.types.ActionFormat
+          - default
+          rep: !!python/object/apply:gr00t.data.types.ActionRepresentation
+          - relative
+          state_key: single_arm
+          type: &id002 !!python/object/apply:gr00t.data.types.ActionType
+          - non_eef
+        - !!python/object:gr00t.data.types.ActionConfig
+          format: *id001
+          rep: !!python/object/apply:gr00t.data.types.ActionRepresentation
+          - absolute
+          state_key: null
+          type: *id002
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        mean_std_embedding_keys: null
+        modality_keys:
+        - single_arm
+        - gripper
+        sin_cos_embedding_keys: null
+      language: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - annotation.human.task_description
+        sin_cos_embedding_keys: null
+      state: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - single_arm
+        - gripper
+        sin_cos_embedding_keys: null
+      video: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - ego_view
+        - left_view
+        - right_view
+        sin_cos_embedding_keys: null
+  mode: single_turn
+  multiprocessing_context: fork
+  num_shards_per_epoch: 100000
+  override_pretraining_statistics: false
+  random_chop: 0.0
+  seed: 42
+  shard_size: 1024
+  shuffle: true
+  subsample_ratio: 1.0
+  video_backend: torchcodec
+load_config_path: null
+model: !!python/object:gr00t.configs.model.gr00t_n1d6.Gr00tN1d6Config
+  _attn_implementation_autoset: false
+  _attn_implementation_internal: null
+  _commit_hash: null
+  _name_or_path: ''
+  add_cross_attention: false
+  architectures: null
+  backbone_model_type: eagle
+  backbone_trainable_params_fp32: true
+  bad_words_ids: null
+  begin_suppress_tokens: null
+  bos_token_id: null
+  chunk_size_feed_forward: 0
+  color_jitter_params:
+    brightness: 0.3
+    contrast: 0.4
+    hue: 0.08
+    saturation: 0.5
+  cross_attention_hidden_size: null
+  decoder_start_token_id: null
+  diffusion_model_cfg:
+    attention_head_dim: 48
+    dropout: 0.2
+    final_dropout: true
+    interleave_self_attention: true
+    norm_type: ada_norm
+    num_attention_heads: 32
+    num_layers: 32
+    output_dim: 1024
+    positional_embeddings: null
+  diversity_penalty: 0.0
+  do_sample: false
+  eagle_collator: true
+  early_stopping: false
+  encoder_no_repeat_ngram_size: 0
+  eos_token_id: null
+  exponential_decay_length_penalty: null
+  finetuning_task: null
+  forced_bos_token_id: null
+  forced_eos_token_id: null
+  id2label:
+    0: LABEL_0
+    1: LABEL_1
+  is_decoder: false
+  is_encoder_decoder: false
+  label2id:
+    LABEL_0: 0
+    LABEL_1: 1
+  length_penalty: 1.0
+  load_bf16: false
+  max_length: 20
+  min_length: 0
+  model_name: nvidia/Eagle-Block2A-2B-v2
+  no_repeat_ngram_size: 0
+  num_beam_groups: 1
+  num_beams: 1
+  num_return_sequences: 1
+  output_attentions: false
+  output_hidden_states: false
+  output_scores: false
+  pad_token_id: null
+  prefix: null
+  problem_type: null
+  pruned_heads: {}
+  random_rotation_angle: null
+  remove_invalid_values: false
+  repetition_penalty: 1.0
+  reproject_vision: false
+  return_dict: true
+  return_dict_in_generate: false
+  sep_token_id: null
+  state_dropout_prob: 0.0
+  suppress_tokens: null
+  task_specific_params: null
+  temperature: 1.0
+  tf_legacy_loss: false
+  tie_encoder_decoder: false
+  tie_word_embeddings: true
+  tokenizer_class: null
+  top_k: 50
+  top_p: 1.0
+  torch_dtype: null
+  torchscript: false
+  transformers_version: null
+  tune_diffusion_model: true
+  tune_llm: false
+  tune_projector: true
+  tune_visual: false
+  typical_p: 1.0
+  use_bfloat16: false
+  use_relative_action: true
+training: !!python/object:gr00t.configs.training.training_config.TrainingConfig
+  add_rl_callback: false
+  assert_loss_less_than: null
+  batch_size: null
+  bf16: true
+  dataloader_num_workers: 32
+  ddp_bucket_cap_mb: 100
+  deepspeed_stage: 2
+  enable_open_loop_eval: false
+  enable_profiling: false
+  eval_batch_size: 2
+  eval_bf16: true
+  eval_set_split_ratio: 0.1
+  eval_steps: 500
+  eval_strategy: 'no'
+  experiment_name: null
+  fp16: false
+  global_batch_size: 32
+  gradient_accumulation_steps: 1
+  gradient_checkpointing: false
+  learning_rate: 0.0001
+  logging_steps: 10
+  lr_scheduler_type: cosine
+  max_concurrent_uploads: 2
+  max_grad_norm: 1.0
+  max_retries: 3
+  max_steps: 10000
+  num_gpus: 4
+  open_loop_eval_plot_indices: null
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_traj_ids:
+  - 0
+  optim: adamw_torch
+  output_dir: /workspaces/isaaclab_arena/models
+  remove_unused_columns: false
+  save_best_eval_metric_greater_is_better: true
+  save_best_eval_metric_name: ''
+  save_steps: 200
+  save_total_limit: 5
+  save_vl_model: false
+  start_from_checkpoint: nvidia/GR00T-N1.6-3B
+  tf32: true
+  transformers_access_token: null
+  transformers_cache_dir: null
+  transformers_local_files_only: false
+  transformers_trust_remote_code: true
+  upload_checkpoints: false
+  upload_every: 1000
+  upload_last_n_checkpoints: 5
+  use_ddp: false
+  use_wandb: true
+  wandb_project: finetune-gr00t-n1d6
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  weight_decay: 1.0e-05

experiment_cfg/dataset_statistics.json

@@ -0,0 +1,1046 @@
+{
+  "new_embodiment": {
+    "state": {
+      "single_arm": {
+        "min": [
+          -0.10585178434848785,
+          -1.4820414781570435,
+          -2.8973002433776855,
+          -3.0718090534210205,
+          -0.5827794075012207,
+          0.5105782747268677,
+          -0.08852053433656693
+        ],
+        "max": [
+          2.868842840194702,
+          1.4677705764770508,
+          0.31179291009902954,
+          -1.1280159950256348,
+          0.7319774031639099,
+          3.0557637214660645,
+          1.9780731201171875
+        ],
+        "mean": [
+          0.22983360290527344,
+          -0.688245415687561,
+          -0.30744367837905884,
+          -2.3262205123901367,
+          -0.12742096185684204,
+          2.051501989364624,
+          0.8900724053382874
+        ],
+        "std": [
+          0.2312260866165161,
+          0.36163097620010376,
+          0.3226556181907654,
+          0.4621410071849823,
+          0.13340327143669128,
+          0.6676610708236694,
+          0.331983357667923
+        ],
+        "q01": [
+          -0.07056046731770038,
+          -1.3751703095436096,
+          -1.1829473924636842,
+          -2.876049301624298,
+          -0.46858262926340105,
+          0.7588222640752792,
+          0.30041238009929655
+        ],
+        "q99": [
+          0.7168123406171805,
+          0.08001862108707435,
+          0.2015080998837948,
+          -1.1704320156574248,
+          0.22084160149097443,
+          2.967472722530365,
+          1.6471182358264924
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.0020317176822572947,
+          0.001846641767770052
+        ],
+        "max": [
+          0.03999999910593033,
+          0.03999999910593033
+        ],
+        "mean": [
+          0.029287075623869896,
+          0.029307108372449875
+        ],
+        "std": [
+          0.0159694105386734,
+          0.015958961099386215
+        ],
+        "q01": [
+          0.002783102965913713,
+          0.002801688495092094
+        ],
+        "q99": [
+          0.03999999910593033,
+          0.03999999910593033
+        ]
+      }
+    },
+    "action": {
+      "single_arm": {
+        "min": [
+          -0.3021092414855957,
+          -1.8116390705108643,
+          -4.104869365692139,
+          -3.249983549118042,
+          -0.7601099014282227,
+          0.08673098683357239,
+          -0.37897980213165283
+        ],
+        "max": [
+          4.587083339691162,
+          1.6342649459838867,
+          0.4067772626876831,
+          -0.5402984023094177,
+          1.0283595323562622,
+          3.2774882316589355,
+          2.1884078979492188
+        ],
+        "mean": [
+          0.24527984857559204,
+          -0.6943578720092773,
+          -0.30451175570487976,
+          -2.3437130451202393,
+          -0.13110630214214325,
+          2.0647470951080322,
+          0.9045398831367493
+        ],
+        "std": [
+          0.255961537361145,
+          0.40363916754722595,
+          0.3533377945423126,
+          0.4557691514492035,
+          0.1536775678396225,
+          0.6757283806800842,
+          0.35488656163215637
+        ],
+        "q01": [
+          -0.08808465927839279,
+          -1.4722782373428345,
+          -1.2692882645130157,
+          -2.9259047293663025,
+          -0.49780495792627333,
+          0.7401778894662857,
+          0.2420878741145134
+        ],
+        "q99": [
+          0.86079528450966,
+          0.1582057143747808,
+          0.251977399289608,
+          -1.1593033838272095,
+          0.3293571025133144,
+          3.044574635028839,
+          1.720823279619217
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.0
+        ],
+        "max": [
+          0.03999999910593033
+        ],
+        "mean": [
+          0.028149358928203583
+        ],
+        "std": [
+          0.018264442682266235
+        ],
+        "q01": [
+          0.0
+        ],
+        "q99": [
+          0.03999999910593033
+        ]
+      }
+    },
+    "relative_action": {
+      "single_arm": {
+        "min": [
+          [
+            -1.3799628019332886,
+            -0.846703827381134,
+            -2.5530529022216797,
+            -0.6990278959274292,
+            -0.7236334085464478,
+            -0.7074007987976074,
+            -0.6110000014305115
+          ],
+          [
+            -1.3839592933654785,
+            -0.8901764154434204,
+            -2.5966553688049316,
+            -0.704864501953125,
+            -0.7758755683898926,
+            -0.7596008777618408,
+            -0.6546932458877563
+          ],
+          [
+            -1.3854783773422241,
+            -0.9336205720901489,
+            -2.640198230743408,
+            -0.7093747854232788,
+            -0.8227287530899048,
+            -0.8116786479949951,
+            -0.6927764415740967
+          ],
+          [
+            -1.3836368322372437,
+            -0.977132260799408,
+            -2.683812141418457,
+            -0.7132972478866577,
+            -0.8571776747703552,
+            -0.8638787269592285,
+            -0.724608838558197
+          ],
+          [
+            -1.375198245048523,
+            -1.0207138061523438,
+            -2.727403402328491,
+            -0.7126319408416748,
+            -0.8606671094894409,
+            -0.9160792827606201,
+            -0.7249724864959717
+          ],
+          [
+            -1.3581116199493408,
+            -1.064221978187561,
+            -2.771050453186035,
+            -0.7117702960968018,
+            -0.8995289206504822,
+            -0.9682760238647461,
+            -0.7250082492828369
+          ],
+          [
+            -1.3482881784439087,
+            -1.1072661876678467,
+            -2.8145911693573,
+            -0.7129005193710327,
+            -0.9034186601638794,
+            -1.0203931331634521,
+            -0.7250344157218933
+          ],
+          [
+            -1.3369436264038086,
+            -1.1432496309280396,
+            -2.8581626415252686,
+            -0.7318893671035767,
+            -0.9344077110290527,
+            -1.0660440921783447,
+            -0.7429417371749878
+          ],
+          [
+            -1.3278416395187378,
+            -1.1847753524780273,
+            -2.9017910957336426,
+            -0.7744251489639282,
+            -0.9756299257278442,
+            -1.1081738471984863,
+            -0.7565152645111084
+          ],
+          [
+            -1.3209303617477417,
+            -1.2282776832580566,
+            -2.9453225135803223,
+            -0.8113338947296143,
+            -1.01007878780365,
+            -1.1416666507720947,
+            -0.7700652480125427
+          ],
+          [
+            -1.316198706626892,
+            -1.2717782258987427,
+            -2.981689453125,
+            -0.8384400606155396,
+            -1.0090559720993042,
+            -1.1768505573272705,
+            -0.7804306745529175
+          ],
+          [
+            -1.3137893676757812,
+            -1.3152940273284912,
+            -3.00955867767334,
+            -0.8819290399551392,
+            -1.0146688222885132,
+            -1.2041229009628296,
+            -0.7792332172393799
+          ],
+          [
+            -1.3064651489257812,
+            -1.358786940574646,
+            -3.0330066680908203,
+            -0.9245129823684692,
+            -1.0318551063537598,
+            -1.2158663272857666,
+            -0.778796911239624
+          ],
+          [
+            -1.3069504499435425,
+            -1.402289867401123,
+            -3.0418806076049805,
+            -0.9453524351119995,
+            -1.0474858283996582,
+            -1.239392638206482,
+            -0.7786667943000793
+          ],
+          [
+            -1.2948497533798218,
+            -1.4457826614379883,
+            -3.049644947052002,
+            -0.9705063104629517,
+            -1.0315730571746826,
+            -1.2627869844436646,
+            -0.7782648801803589
+          ],
+          [
+            -1.2816047668457031,
+            -1.4859166145324707,
+            -3.057337760925293,
+            -0.9996410608291626,
+            -1.010919451713562,
+            -1.285316824913025,
+            -0.793167233467102
+          ]
+        ],
+        "max": [
+          [
+            2.567418336868286,
+            0.7940642833709717,
+            1.3483597040176392,
+            0.5038654804229736,
+            0.8064826130867004,
+            0.42534810304641724,
+            0.7363687753677368
+          ],
+          [
+            2.6070313453674316,
+            0.8375340700149536,
+            1.3556522130966187,
+            0.5470445156097412,
+            0.8587763905525208,
+            0.4535168409347534,
+            0.7815820574760437
+          ],
+          [
+            2.6480026245117188,
+            0.8809951543807983,
+            1.3603707551956177,
+            0.5853910446166992,
+            0.9099692702293396,
+            0.4805412292480469,
+            0.8328911066055298
+          ],
+          [
+            2.6904377937316895,
+            0.9244890213012695,
+            1.3617339134216309,
+            0.62050461769104,
+            0.9537554383277893,
+            0.5171980857849121,
+            0.8850930333137512
+          ],
+          [
+            2.733157157897949,
+            0.9679925441741943,
+            1.3563916683197021,
+            0.653130054473877,
+            0.9584329724311829,
+            0.5527119636535645,
+            0.9372936487197876
+          ],
+          [
+            2.776454448699951,
+            1.011548399925232,
+            1.3423844575881958,
+            0.7159806489944458,
+            0.962762176990509,
+            0.5780115127563477,
+            0.9895003437995911
+          ],
+          [
+            2.819460391998291,
+            1.0550143718719482,
+            1.3357713222503662,
+            0.7591596841812134,
+            0.967065691947937,
+            0.6083757877349854,
+            0.9970666170120239
+          ],
+          [
+            2.862929582595825,
+            1.098570466041565,
+            1.3283271789550781,
+            0.8227701187133789,
+            1.0051043033599854,
+            0.6292746067047119,
+            1.0122740268707275
+          ],
+          [
+            2.906933069229126,
+            1.142111897468567,
+            1.3234343528747559,
+            0.9115495681762695,
+            1.0388011932373047,
+            0.6505081653594971,
+            1.0234447717666626
+          ],
+          [
+            2.9509470462799072,
+            1.185598373413086,
+            1.3211033344268799,
+            0.9813358187675476,
+            1.0714678764343262,
+            0.6828610897064209,
+            1.0570249557495117
+          ],
+          [
+            2.9969143867492676,
+            1.2272741794586182,
+            1.3213672637939453,
+            1.039294958114624,
+            1.1007328033447266,
+            0.7069540023803711,
+            1.0594041347503662
+          ],
+          [
+            3.0403480529785156,
+            1.253152847290039,
+            1.3244858980178833,
+            1.1303977966308594,
+            1.1028329133987427,
+            0.7325007915496826,
+            1.0783926248550415
+          ],
+          [
+            3.0788755416870117,
+            1.272753119468689,
+            1.3235565423965454,
+            1.1840159893035889,
+            1.1049879789352417,
+            0.7554268836975098,
+            1.080771803855896
+          ],
+          [
+            3.1056816577911377,
+            1.2986316680908203,
+            1.330661654472351,
+            1.2504518032073975,
+            1.108304738998413,
+            0.7780568599700928,
+            1.0822246074676514
+          ],
+          [
+            3.13144850730896,
+            1.3217449188232422,
+            1.3251062631607056,
+            1.3208149671554565,
+            1.1107087135314941,
+            0.8055588006973267,
+            1.0833207368850708
+          ],
+          [
+            3.156325578689575,
+            1.3329780101776123,
+            1.3181878328323364,
+            1.4193122386932373,
+            1.1303942203521729,
+            0.8331756591796875,
+            1.0827785730361938
+          ]
+        ],
+        "mean": [
+          [
+            0.0162045881152153,
+            -0.01284683309495449,
+            0.003098804736509919,
+            -0.03372170031070709,
+            -0.004505295772105455,
+            0.021783702075481415,
+            0.01521182619035244
+          ],
+          [
+            0.017577122896909714,
+            -0.013530674390494823,
+            0.003645440097898245,
+            -0.03580537065863609,
+            -0.004916673060506582,
+            0.02341781184077263,
+            0.016815518960356712
+          ],
+          [
+            0.018947791308164597,
+            -0.014144103042781353,
+            0.0042039500549435616,
+            -0.037736132740974426,
+            -0.005320144351571798,
+            0.02498255856335163,
+            0.018421174958348274
+          ],
+          [
+            0.02031758800148964,
+            -0.014692251570522785,
+            0.004773355089128017,
+            -0.03952430561184883,
+            -0.005717766005545855,
+            0.026474300771951675,
+            0.02003154158592224
+          ],
+          [
+            0.02169603295624256,
+            -0.015192627906799316,
+            0.005315880756825209,
+            -0.041174083948135376,
+            -0.0061398898251354694,
+            0.027881216257810593,
+            0.021609634160995483
+          ],
+          [
+            0.023072782903909683,
+            -0.015621185302734375,
+            0.0058733210898935795,
+            -0.04267120361328125,
+            -0.006557443644851446,
+            0.02921134978532791,
+            0.023191053420305252
+          ],
+          [
+            0.02445853501558304,
+            -0.016012879088521004,
+            0.006379991304129362,
+            -0.04403478652238846,
+            -0.007005136460065842,
+            0.030454443767666817,
+            0.0247166957706213
+          ],
+          [
+            0.025841616094112396,
+            -0.0163332037627697,
+            0.006896289065480232,
+            -0.04524318873882294,
+            -0.007450786419212818,
+            0.0316125713288784,
+            0.02624223567545414
+          ],
+          [
+            0.027225056663155556,
+            -0.016576094552874565,
+            0.0074104429222643375,
+            -0.04628229886293411,
+            -0.007899835705757141,
+            0.03268665075302124,
+            0.02776043862104416
+          ],
+          [
+            0.02863071672618389,
+            -0.01684136874973774,
+            0.007741305977106094,
+            -0.04724014550447464,
+            -0.008443175815045834,
+            0.03367039933800697,
+            0.029093485325574875
+          ],
+          [
+            0.03001892939209938,
+            -0.01698327250778675,
+            0.008157290518283844,
+            -0.04800064489245415,
+            -0.008937019854784012,
+            0.03457446023821831,
+            0.03048980049788952
+          ],
+          [
+            0.03142230585217476,
+            -0.017429865896701813,
+            0.008554172702133656,
+            -0.04909304529428482,
+            -0.009442312642931938,
+            0.03550037741661072,
+            0.031910303980112076
+          ],
+          [
+            0.032860491424798965,
+            -0.017565155401825905,
+            0.00891648605465889,
+            -0.049608033150434494,
+            -0.009951156564056873,
+            0.03615717589855194,
+            0.03330664336681366
+          ],
+          [
+            0.03432377055287361,
+            -0.017927052453160286,
+            0.00926224235445261,
+            -0.050526637583971024,
+            -0.010460241697728634,
+            0.0369611456990242,
+            0.03471340239048004
+          ],
+          [
+            0.0357922725379467,
+            -0.0182654969394207,
+            0.009601501747965813,
+            -0.051374953240156174,
+            -0.010974744334816933,
+            0.03770560398697853,
+            0.036107003688812256
+          ],
+          [
+            0.03721015155315399,
+            -0.018647940829396248,
+            0.010218545794487,
+            -0.052236221730709076,
+            -0.011343417689204216,
+            0.038379695266485214,
+            0.03776820749044418
+          ]
+        ],
+        "std": [
+          [
+            0.09561787545681,
+            0.1841392070055008,
+            0.1230149194598198,
+            0.1651013195514679,
+            0.08265779167413712,
+            0.20834241807460785,
+            0.11478734016418457
+          ],
+          [
+            0.10193362832069397,
+            0.19995222985744476,
+            0.13262243568897247,
+            0.18077391386032104,
+            0.08958185464143753,
+            0.22795242071151733,
+            0.1250300258398056
+          ],
+          [
+            0.10761761665344238,
+            0.21515797078609467,
+            0.1417541801929474,
+            0.19604629278182983,
+            0.09583412855863571,
+            0.24720466136932373,
+            0.13476639986038208
+          ],
+          [
+            0.11283185333013535,
+            0.22989585995674133,
+            0.1505030393600464,
+            0.2110775262117386,
+            0.10155555605888367,
+            0.26614490151405334,
+            0.14409737288951874
+          ],
+          [
+            0.11775802075862885,
+            0.2443120777606964,
+            0.15899932384490967,
+            0.22597895562648773,
+            0.10690765827894211,
+            0.2848511338233948,
+            0.153147891163826
+          ],
+          [
+            0.12265162914991379,
+            0.2584792673587799,
+            0.16736982762813568,
+            0.2408720850944519,
+            0.11210725456476212,
+            0.3033398687839508,
+            0.16198629140853882
+          ],
+          [
+            0.1275392770767212,
+            0.27235040068626404,
+            0.17564959824085236,
+            0.25575950741767883,
+            0.11716651171445847,
+            0.3216555714607239,
+            0.1706080585718155
+          ],
+          [
+            0.13233336806297302,
+            0.2859433889389038,
+            0.1838448941707611,
+            0.2707005441188812,
+            0.12200450152158737,
+            0.33981001377105713,
+            0.1790284812450409
+          ],
+          [
+            0.13690054416656494,
+            0.2992684543132782,
+            0.1919122040271759,
+            0.28575506806373596,
+            0.12655234336853027,
+            0.3577987253665924,
+            0.18726085126399994
+          ],
+          [
+            0.1411910057067871,
+            0.3122430741786957,
+            0.1998082548379898,
+            0.3008229732513428,
+            0.1307753324508667,
+            0.3756433129310608,
+            0.1953035145998001
+          ],
+          [
+            0.14517177641391754,
+            0.32497692108154297,
+            0.20752082765102386,
+            0.31603574752807617,
+            0.1346835047006607,
+            0.39332452416419983,
+            0.2031596302986145
+          ],
+          [
+            0.1488790065050125,
+            0.33717432618141174,
+            0.21507450938224792,
+            0.33044007420539856,
+            0.13834868371486664,
+            0.41073474287986755,
+            0.21088287234306335
+          ],
+          [
+            0.15247905254364014,
+            0.34908828139305115,
+            0.22253793478012085,
+            0.3444584906101227,
+            0.1418810486793518,
+            0.42792513966560364,
+            0.218512162566185
+          ],
+          [
+            0.1560245156288147,
+            0.36060401797294617,
+            0.22992275655269623,
+            0.3577353358268738,
+            0.14529705047607422,
+            0.44483527541160583,
+            0.22603793442249298
+          ],
+          [
+            0.15951526165008545,
+            0.37189415097236633,
+            0.23725645244121552,
+            0.371038019657135,
+            0.14860472083091736,
+            0.461593896150589,
+            0.23349352180957794
+          ],
+          [
+            0.1629764437675476,
+            0.38295140862464905,
+            0.2445494830608368,
+            0.3841154873371124,
+            0.15181609988212585,
+            0.47810259461402893,
+            0.24087539315223694
+          ]
+        ],
+        "q01": [
+          [
+            -0.19950662806630134,
+            -0.5267880302667618,
+            -0.2804043960571289,
+            -0.49720216274261475,
+            -0.24181631177663804,
+            -0.44567849397659304,
+            -0.25556061893701554
+          ],
+          [
+            -0.21852773502469064,
+            -0.5698291510343552,
+            -0.3015518325567246,
+            -0.5363495886325836,
+            -0.2647059428691864,
+            -0.48813892602920533,
+            -0.275542603135109
+          ],
+          [
+            -0.23660527691245078,
+            -0.6086035186052322,
+            -0.31789174526929853,
+            -0.5604154777526855,
+            -0.28393503457307817,
+            -0.5290071642398835,
+            -0.2913743835687637
+          ],
+          [
+            -0.24862138345837592,
+            -0.649602175951004,
+            -0.3364737406373024,
+            -0.5784401261806488,
+            -0.30646128237247466,
+            -0.5700734174251556,
+            -0.3095038095116615
+          ],
+          [
+            -0.26433300614356997,
+            -0.6882567131519318,
+            -0.352595399916172,
+            -0.5933414793014526,
+            -0.3249793803691864,
+            -0.6097909665107727,
+            -0.32609487652778624
+          ],
+          [
+            -0.27847526580095294,
+            -0.7230063188076019,
+            -0.3685854694247246,
+            -0.6150582671165467,
+            -0.33894465684890746,
+            -0.6511167442798614,
+            -0.33947302132844925
+          ],
+          [
+            -0.2881611752510071,
+            -0.7600984990596771,
+            -0.3851272436976433,
+            -0.6339384734630584,
+            -0.35804505556821825,
+            -0.6873469173908233,
+            -0.3513255763053894
+          ],
+          [
+            -0.29823750019073486,
+            -0.7952856153249741,
+            -0.40130406498909,
+            -0.6537419366836548,
+            -0.36884784877300264,
+            -0.7222859227657318,
+            -0.3663733446598053
+          ],
+          [
+            -0.3080375599861145,
+            -0.8278200089931488,
+            -0.4162969410419464,
+            -0.6719410765171051,
+            -0.3811878323554993,
+            -0.7592104160785675,
+            -0.3828606143593788
+          ],
+          [
+            -0.3192089283466339,
+            -0.8594429874420166,
+            -0.43057493746280673,
+            -0.6891207134723664,
+            -0.3899207302927971,
+            -0.7906776928901672,
+            -0.39594156324863433
+          ],
+          [
+            -0.3260902661085129,
+            -0.890479126572609,
+            -0.4381144392490387,
+            -0.713641117811203,
+            -0.4018076476454735,
+            -0.8239653241634369,
+            -0.40608411192893984
+          ],
+          [
+            -0.32998249888420106,
+            -0.9208912491798401,
+            -0.4485784482955933,
+            -0.7404101181030274,
+            -0.40892592787742615,
+            -0.8567238235473633,
+            -0.41712283045053483
+          ],
+          [
+            -0.337571017742157,
+            -0.9505359423160553,
+            -0.45878848373889924,
+            -0.7647261500358582,
+            -0.41907749682664874,
+            -0.8858703064918518,
+            -0.4311297929286957
+          ],
+          [
+            -0.344514417052269,
+            -0.978490937948227,
+            -0.4731087791919708,
+            -0.7942435812950134,
+            -0.42579659909009937,
+            -0.9157415616512299,
+            -0.4435725921392441
+          ],
+          [
+            -0.3521229174733162,
+            -1.0078155159950257,
+            -0.4874843668937683,
+            -0.8193191123008728,
+            -0.4365257340669632,
+            -0.9454274779558182,
+            -0.4541371762752533
+          ],
+          [
+            -0.35581389248371126,
+            -1.0338427817821503,
+            -0.5004701679944992,
+            -0.8494838178157806,
+            -0.442863190472126,
+            -0.9770041716098785,
+            -0.46435451209545137
+          ]
+        ],
+        "q99": [
+          [
+            0.3228484025597576,
+            0.40104085892438895,
+            0.31244824022054674,
+            0.2954720175266266,
+            0.270292398929597,
+            0.3568690168857576,
+            0.32504342615604404
+          ],
+          [
+            0.347237007021904,
+            0.43345616728067404,
+            0.3387280979752541,
+            0.3265518689155581,
+            0.2923547741770747,
+            0.3882337927818299,
+            0.35589420855045356
+          ],
+          [
+            0.3646373689174653,
+            0.46306476682424547,
+            0.36461793124675757,
+            0.3594550025463107,
+            0.3067959910631185,
+            0.41669229030609134,
+            0.38290659368038193
+          ],
+          [
+            0.3797282963991172,
+            0.49183411687612544,
+            0.3919186636805537,
+            0.3934913563728333,
+            0.32532002121210113,
+            0.44333425402641313,
+            0.41034477829933175
+          ],
+          [
+            0.3966591441631323,
+            0.5208547359704973,
+            0.4143754762411119,
+            0.42633920907974243,
+            0.34047347754240076,
+            0.4709657806158066,
+            0.4371810677647591
+          ],
+          [
+            0.4147697642445564,
+            0.548696836233139,
+            0.4355347976088526,
+            0.459941817522049,
+            0.3566160580515862,
+            0.49860751807689674,
+            0.4619476205110556
+          ],
+          [
+            0.43306236445903784,
+            0.5744219225645065,
+            0.4582824614644051,
+            0.49298112750053413,
+            0.37070908635854743,
+            0.5247364115715029,
+            0.4946159166097642
+          ],
+          [
+            0.4535762205719948,
+            0.5974840182065965,
+            0.4781889137625704,
+            0.52674173116684,
+            0.38307282924652125,
+            0.551970843076706,
+            0.5218617248535158
+          ],
+          [
+            0.4752288824319842,
+            0.6203914856910707,
+            0.4968277835845948,
+            0.563329246044159,
+            0.3946968352794648,
+            0.5789938712120056,
+            0.5464925676584249
+          ],
+          [
+            0.49233787536621126,
+            0.6448995488882068,
+            0.5179557728767397,
+            0.5989310371875763,
+            0.4044640159606941,
+            0.6043756437301636,
+            0.5733339744806294
+          ],
+          [
+            0.5049607670307164,
+            0.6641878437995912,
+            0.5366499340534212,
+            0.6363685309886934,
+            0.4191509392857555,
+            0.6278273916244507,
+            0.5940367233753204
+          ],
+          [
+            0.5195178943872457,
+            0.6806331920623779,
+            0.5539186495542527,
+            0.6678686726093293,
+            0.4275676235556606,
+            0.652043970823288,
+            0.6186116337776186
+          ],
+          [
+            0.5356091660261156,
+            0.698597639799119,
+            0.571190173625946,
+            0.7009571266174317,
+            0.43485867321491284,
+            0.6771110093593598,
+            0.6389012908935552
+          ],
+          [
+            0.5494599288702018,
+            0.714040715694428,
+            0.5911411744356158,
+            0.7320740032196047,
+            0.44157062917947926,
+            0.7026675033569336,
+            0.6588078260421755
+          ],
+          [
+            0.5620473337173466,
+            0.7256977421045304,
+            0.6092195951938636,
+            0.7614367377758029,
+            0.446760193109514,
+            0.7281576597690583,
+            0.6745113229751587
+          ],
+          [
+            0.5714725697040565,
+            0.7397873222827912,
+            0.6257535284757625,
+            0.791651120185852,
+            0.4508011704683311,
+            0.752045361995697,
+            0.6951469552516938
+          ]
+        ]
+      }
+    }
+  }
+}

experiment_cfg/final_model_config.json

@@ -0,0 +1,53 @@
+{
+  "model_type": "Gr00tN1d6",
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "backbone_model_type": "eagle",
+  "model_revision": null,
+  "tune_top_llm_layers": 4,
+  "backbone_embedding_dim": 2048,
+  "tune_llm": false,
+  "tune_visual": false,
+  "select_layer": 16,
+  "reproject_vision": false,
+  "use_flash_attention": true,
+  "load_bf16": true,
+  "collator_overwrite_image_inputs": false,
+  "eagle_collator": true,
+  "backbone_trainable_params_fp32": true,
+  "apply_sincos_state_encoding": true,
+  "use_relative_action": true,
+  "max_state_dim": 128,
+  "max_action_dim": 128,
+  "action_horizon": 50,
+  "hidden_size": 1024,
+  "input_embedding_dim": 1536,
+  "add_pos_embed": true,
+  "attn_dropout": 0.2,
+  "use_vlln": true,
+  "max_seq_len": 1024,
+  "use_alternate_vl_dit": true,
+  "attend_text_every_n_blocks": 2,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "num_inference_timesteps": 4,
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_timestep_buckets": 1000,
+  "tune_projector": true,
+  "tune_diffusion_model": true,
+  "tune_vlln": true,
+  "state_dropout_prob": 0.0,
+  "state_additive_noise_scale": 0.0,
+  "max_num_embodiments": 32
+}

global_step1400/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:589b6f457eac67317b8b7c7fc62422099fbf157331f021eafe9545d1e3143c97
+size 4859923697

global_step1400/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:014002cc4cef79ee278e33254f5b431a1b6d18429a7f4989fa397c994ca61fc2
+size 4859923761

global_step1400/bf16_zero_pp_rank_2_mp_rank_00_optim_states.pt

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

global_step1400/bf16_zero_pp_rank_3_mp_rank_00_optim_states.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64dd4b5b9ea05fcf557ba0d46d4c3b7645c79bbe49911782ec741aeffc378418
+size 4859917873

global_step1400/mp_rank_00_model_states.pt

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

latest

@@ -0,0 +1 @@
+global_step1400

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

processor_config.json

@@ -0,0 +1,456 @@
+{
+  "processor_class": "Gr00tN1d6Processor",
+  "processor_kwargs": {
+    "modality_configs": {
+      "behavior_r1_pro": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "observation.images.rgb.head_256_256",
+            "observation.images.rgb.left_wrist_256_256",
+            "observation.images.rgb.right_wrist_256_256"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "robot_pos",
+            "robot_ori_cos",
+            "robot_ori_sin",
+            "robot_2d_ori",
+            "robot_2d_ori_cos",
+            "robot_2d_ori_sin",
+            "robot_lin_vel",
+            "robot_ang_vel",
+            "arm_left_qpos",
+            "arm_left_qpos_sin",
+            "arm_left_qpos_cos",
+            "eef_left_pos",
+            "eef_left_quat",
+            "gripper_left_qpos",
+            "arm_right_qpos",
+            "arm_right_qpos_sin",
+            "arm_right_qpos_cos",
+            "eef_right_pos",
+            "eef_right_quat",
+            "gripper_right_qpos",
+            "trunk_qpos"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31
+          ],
+          "modality_keys": [
+            "base",
+            "torso",
+            "left_arm",
+            "left_gripper",
+            "right_arm",
+            "right_gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "trunk_qpos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_left_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_right_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "gr1": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "ego_view_bg_crop_pad_res256_freq20"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "task"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "robocasa_panda_omron": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "res256_image_side_0",
+            "res256_image_side_1",
+            "res256_image_wrist_0"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "end_effector_position_relative",
+            "end_effector_rotation_relative",
+            "gripper_qpos",
+            "base_position",
+            "base_rotation"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "end_effector_position",
+            "end_effector_rotation",
+            "gripper_close",
+            "base_motion",
+            "control_mode"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.action.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "new_embodiment": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "ego_view",
+            "left_view",
+            "right_view"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "single_arm",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "single_arm",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "single_arm"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      }
+    },
+    "image_crop_size": null,
+    "image_target_size": null,
+    "use_albumentations": true,
+    "random_rotation_angle": null,
+    "color_jitter_params": {
+      "brightness": 0.3,
+      "contrast": 0.4,
+      "saturation": 0.5,
+      "hue": 0.08
+    },
+    "shortest_image_edge": 256,
+    "crop_fraction": 0.95,
+    "model_name": "nvidia/Eagle-Block2A-2B-v2",
+    "model_type": "eagle",
+    "formalize_language": true,
+    "max_state_dim": 128,
+    "max_action_dim": 128,
+    "max_action_horizon": 50,
+    "use_percentiles": false,
+    "clip_outliers": true,
+    "apply_sincos_state_encoding": true,
+    "use_relative_action": true
+  }
+}

rng_state_0.pth

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

rng_state_1.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56bd7d3a38b53f9caefd474ef17181a618b190df886789a4373ee7200362df48
+size 15429

rng_state_2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:697208fba032129ad87ffe3d2cfd2cf533be72555f62540fee32fc497c2bb7c2
+size 15429

rng_state_3.pth

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

scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01d5a4dee5b66e437551e6fb6cdd5404bad2ffd89d97f7b845865033e04f2cfd
+size 1465

trainer_state.json

@@ -0,0 +1,874 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.14,
+  "eval_steps": 500,
+  "global_step": 1400,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "grad_norm": 2.7235865592956543,
+      "learning_rate": 1.8e-06,
+      "loss": 1.3062,
+      "step": 10
+    },
+    {
+      "grad_norm": 1.728572130203247,
+      "learning_rate": 3.8e-06,
+      "loss": 1.2836,
+      "step": 20
+    },
+    {
+      "grad_norm": 0.9103507995605469,
+      "learning_rate": 5.8e-06,
+      "loss": 1.2381,
+      "step": 30
+    },
+    {
+      "grad_norm": 0.5990565419197083,
+      "learning_rate": 7.8e-06,
+      "loss": 1.1586,
+      "step": 40
+    },
+    {
+      "grad_norm": 0.7476415634155273,
+      "learning_rate": 9.800000000000001e-06,
+      "loss": 1.1645,
+      "step": 50
+    },
+    {
+      "grad_norm": 0.779353141784668,
+      "learning_rate": 1.18e-05,
+      "loss": 1.1062,
+      "step": 60
+    },
+    {
+      "grad_norm": 1.5331050157546997,
+      "learning_rate": 1.3800000000000002e-05,
+      "loss": 1.0745,
+      "step": 70
+    },
+    {
+      "grad_norm": 0.8821016550064087,
+      "learning_rate": 1.58e-05,
+      "loss": 1.0459,
+      "step": 80
+    },
+    {
+      "grad_norm": 0.7069372534751892,
+      "learning_rate": 1.78e-05,
+      "loss": 1.0316,
+      "step": 90
+    },
+    {
+      "grad_norm": 0.8882275819778442,
+      "learning_rate": 1.9800000000000004e-05,
+      "loss": 1.0574,
+      "step": 100
+    },
+    {
+      "grad_norm": 0.7405692934989929,
+      "learning_rate": 2.18e-05,
+      "loss": 1.0404,
+      "step": 110
+    },
+    {
+      "grad_norm": 0.8989152908325195,
+      "learning_rate": 2.38e-05,
+      "loss": 1.0413,
+      "step": 120
+    },
+    {
+      "grad_norm": 0.7994363903999329,
+      "learning_rate": 2.58e-05,
+      "loss": 1.0046,
+      "step": 130
+    },
+    {
+      "grad_norm": 0.7832587361335754,
+      "learning_rate": 2.7800000000000005e-05,
+      "loss": 1.0432,
+      "step": 140
+    },
+    {
+      "grad_norm": 0.7279527187347412,
+      "learning_rate": 2.98e-05,
+      "loss": 1.0247,
+      "step": 150
+    },
+    {
+      "grad_norm": 0.5736011862754822,
+      "learning_rate": 3.18e-05,
+      "loss": 1.0226,
+      "step": 160
+    },
+    {
+      "grad_norm": 0.7875821590423584,
+      "learning_rate": 3.38e-05,
+      "loss": 1.0132,
+      "step": 170
+    },
+    {
+      "grad_norm": 1.9923306703567505,
+      "learning_rate": 3.58e-05,
+      "loss": 1.0068,
+      "step": 180
+    },
+    {
+      "grad_norm": 1.0174721479415894,
+      "learning_rate": 3.7800000000000004e-05,
+      "loss": 0.9882,
+      "step": 190
+    },
+    {
+      "grad_norm": 1.2093017101287842,
+      "learning_rate": 3.9800000000000005e-05,
+      "loss": 0.9596,
+      "step": 200
+    },
+    {
+      "grad_norm": 1.1982413530349731,
+      "learning_rate": 4.18e-05,
+      "loss": 0.9083,
+      "step": 210
+    },
+    {
+      "grad_norm": 1.0665863752365112,
+      "learning_rate": 4.38e-05,
+      "loss": 0.8807,
+      "step": 220
+    },
+    {
+      "grad_norm": 1.2544691562652588,
+      "learning_rate": 4.58e-05,
+      "loss": 0.839,
+      "step": 230
+    },
+    {
+      "grad_norm": 1.3970142602920532,
+      "learning_rate": 4.78e-05,
+      "loss": 0.8039,
+      "step": 240
+    },
+    {
+      "grad_norm": 1.375622034072876,
+      "learning_rate": 4.9800000000000004e-05,
+      "loss": 0.7449,
+      "step": 250
+    },
+    {
+      "grad_norm": 1.4242069721221924,
+      "learning_rate": 5.1800000000000005e-05,
+      "loss": 0.6853,
+      "step": 260
+    },
+    {
+      "grad_norm": 1.7568080425262451,
+      "learning_rate": 5.380000000000001e-05,
+      "loss": 0.6114,
+      "step": 270
+    },
+    {
+      "grad_norm": 2.0625576972961426,
+      "learning_rate": 5.580000000000001e-05,
+      "loss": 0.5084,
+      "step": 280
+    },
+    {
+      "grad_norm": 2.137798547744751,
+      "learning_rate": 5.7799999999999995e-05,
+      "loss": 0.4182,
+      "step": 290
+    },
+    {
+      "grad_norm": 1.4877147674560547,
+      "learning_rate": 5.9800000000000003e-05,
+      "loss": 0.3704,
+      "step": 300
+    },
+    {
+      "grad_norm": 1.568670630455017,
+      "learning_rate": 6.18e-05,
+      "loss": 0.3203,
+      "step": 310
+    },
+    {
+      "grad_norm": 1.2726070880889893,
+      "learning_rate": 6.38e-05,
+      "loss": 0.2442,
+      "step": 320
+    },
+    {
+      "grad_norm": 1.7254693508148193,
+      "learning_rate": 6.58e-05,
+      "loss": 0.226,
+      "step": 330
+    },
+    {
+      "grad_norm": 1.3184534311294556,
+      "learning_rate": 6.780000000000001e-05,
+      "loss": 0.179,
+      "step": 340
+    },
+    {
+      "grad_norm": 1.247791051864624,
+      "learning_rate": 6.98e-05,
+      "loss": 0.1607,
+      "step": 350
+    },
+    {
+      "grad_norm": 1.0366252660751343,
+      "learning_rate": 7.18e-05,
+      "loss": 0.1462,
+      "step": 360
+    },
+    {
+      "grad_norm": 1.5369986295700073,
+      "learning_rate": 7.38e-05,
+      "loss": 0.1731,
+      "step": 370
+    },
+    {
+      "grad_norm": 1.5386507511138916,
+      "learning_rate": 7.58e-05,
+      "loss": 0.169,
+      "step": 380
+    },
+    {
+      "grad_norm": 1.2936265468597412,
+      "learning_rate": 7.780000000000001e-05,
+      "loss": 0.1706,
+      "step": 390
+    },
+    {
+      "grad_norm": 1.420244574546814,
+      "learning_rate": 7.98e-05,
+      "loss": 0.1689,
+      "step": 400
+    },
+    {
+      "grad_norm": 1.0225824117660522,
+      "learning_rate": 8.18e-05,
+      "loss": 0.169,
+      "step": 410
+    },
+    {
+      "grad_norm": 0.8011528253555298,
+      "learning_rate": 8.38e-05,
+      "loss": 0.1387,
+      "step": 420
+    },
+    {
+      "grad_norm": 0.7911453247070312,
+      "learning_rate": 8.58e-05,
+      "loss": 0.1591,
+      "step": 430
+    },
+    {
+      "grad_norm": 0.9051388502120972,
+      "learning_rate": 8.78e-05,
+      "loss": 0.1456,
+      "step": 440
+    },
+    {
+      "grad_norm": 0.7053804397583008,
+      "learning_rate": 8.98e-05,
+      "loss": 0.1255,
+      "step": 450
+    },
+    {
+      "grad_norm": 1.1338269710540771,
+      "learning_rate": 9.180000000000001e-05,
+      "loss": 0.1367,
+      "step": 460
+    },
+    {
+      "grad_norm": 1.299891471862793,
+      "learning_rate": 9.38e-05,
+      "loss": 0.1254,
+      "step": 470
+    },
+    {
+      "grad_norm": 0.9713226556777954,
+      "learning_rate": 9.58e-05,
+      "loss": 0.13,
+      "step": 480
+    },
+    {
+      "grad_norm": 0.6963210701942444,
+      "learning_rate": 9.78e-05,
+      "loss": 0.1523,
+      "step": 490
+    },
+    {
+      "grad_norm": 0.9888042211532593,
+      "learning_rate": 9.98e-05,
+      "loss": 0.1221,
+      "step": 500
+    },
+    {
+      "grad_norm": 0.7608430981636047,
+      "learning_rate": 9.9999778549206e-05,
+      "loss": 0.1271,
+      "step": 510
+    },
+    {
+      "grad_norm": 1.088894248008728,
+      "learning_rate": 9.999901304280685e-05,
+      "loss": 0.1448,
+      "step": 520
+    },
+    {
+      "grad_norm": 0.7456501126289368,
+      "learning_rate": 9.999770075521164e-05,
+      "loss": 0.1395,
+      "step": 530
+    },
+    {
+      "grad_norm": 0.6787962317466736,
+      "learning_rate": 9.99958417007713e-05,
+      "loss": 0.1083,
+      "step": 540
+    },
+    {
+      "grad_norm": 0.7639486193656921,
+      "learning_rate": 9.999343589981615e-05,
+      "loss": 0.151,
+      "step": 550
+    },
+    {
+      "grad_norm": 0.7701183557510376,
+      "learning_rate": 9.999048337865568e-05,
+      "loss": 0.1162,
+      "step": 560
+    },
+    {
+      "grad_norm": 0.7873595356941223,
+      "learning_rate": 9.998698416957815e-05,
+      "loss": 0.1055,
+      "step": 570
+    },
+    {
+      "grad_norm": 0.5775970220565796,
+      "learning_rate": 9.998293831085037e-05,
+      "loss": 0.0935,
+      "step": 580
+    },
+    {
+      "grad_norm": 0.6811988353729248,
+      "learning_rate": 9.997834584671719e-05,
+      "loss": 0.1106,
+      "step": 590
+    },
+    {
+      "grad_norm": 0.8140270113945007,
+      "learning_rate": 9.997320682740107e-05,
+      "loss": 0.1104,
+      "step": 600
+    },
+    {
+      "grad_norm": 0.7332885265350342,
+      "learning_rate": 9.996752130910149e-05,
+      "loss": 0.0995,
+      "step": 610
+    },
+    {
+      "grad_norm": 0.8294130563735962,
+      "learning_rate": 9.99612893539944e-05,
+      "loss": 0.112,
+      "step": 620
+    },
+    {
+      "grad_norm": 0.8760542869567871,
+      "learning_rate": 9.995451103023144e-05,
+      "loss": 0.1088,
+      "step": 630
+    },
+    {
+      "grad_norm": 0.666292130947113,
+      "learning_rate": 9.994718641193928e-05,
+      "loss": 0.1078,
+      "step": 640
+    },
+    {
+      "grad_norm": 0.6352424025535583,
+      "learning_rate": 9.993931557921874e-05,
+      "loss": 0.1002,
+      "step": 650
+    },
+    {
+      "grad_norm": 0.6974917650222778,
+      "learning_rate": 9.993089861814402e-05,
+      "loss": 0.129,
+      "step": 660
+    },
+    {
+      "grad_norm": 0.8117634057998657,
+      "learning_rate": 9.992193562076166e-05,
+      "loss": 0.0894,
+      "step": 670
+    },
+    {
+      "grad_norm": 0.6594702005386353,
+      "learning_rate": 9.991242668508954e-05,
+      "loss": 0.1132,
+      "step": 680
+    },
+    {
+      "grad_norm": 0.6721746921539307,
+      "learning_rate": 9.990237191511587e-05,
+      "loss": 0.0971,
+      "step": 690
+    },
+    {
+      "grad_norm": 0.6974657773971558,
+      "learning_rate": 9.989177142079802e-05,
+      "loss": 0.1038,
+      "step": 700
+    },
+    {
+      "grad_norm": 0.6547107696533203,
+      "learning_rate": 9.988062531806126e-05,
+      "loss": 0.0993,
+      "step": 710
+    },
+    {
+      "grad_norm": 0.5421633720397949,
+      "learning_rate": 9.986893372879762e-05,
+      "loss": 0.1023,
+      "step": 720
+    },
+    {
+      "grad_norm": 0.5468288660049438,
+      "learning_rate": 9.985669678086443e-05,
+      "loss": 0.1001,
+      "step": 730
+    },
+    {
+      "grad_norm": 0.702468991279602,
+      "learning_rate": 9.984391460808298e-05,
+      "loss": 0.0896,
+      "step": 740
+    },
+    {
+      "grad_norm": 0.6491390466690063,
+      "learning_rate": 9.983058735023709e-05,
+      "loss": 0.1026,
+      "step": 750
+    },
+    {
+      "grad_norm": 0.5702087879180908,
+      "learning_rate": 9.98167151530715e-05,
+      "loss": 0.1,
+      "step": 760
+    },
+    {
+      "grad_norm": 0.5395886301994324,
+      "learning_rate": 9.980229816829034e-05,
+      "loss": 0.0921,
+      "step": 770
+    },
+    {
+      "grad_norm": 0.7115892767906189,
+      "learning_rate": 9.978733655355544e-05,
+      "loss": 0.1247,
+      "step": 780
+    },
+    {
+      "grad_norm": 0.5494698286056519,
+      "learning_rate": 9.977183047248464e-05,
+      "loss": 0.0917,
+      "step": 790
+    },
+    {
+      "grad_norm": 0.543135941028595,
+      "learning_rate": 9.975578009464992e-05,
+      "loss": 0.0863,
+      "step": 800
+    },
+    {
+      "grad_norm": 0.5769555568695068,
+      "learning_rate": 9.97391855955757e-05,
+      "loss": 0.0904,
+      "step": 810
+    },
+    {
+      "grad_norm": 0.8262929320335388,
+      "learning_rate": 9.972204715673669e-05,
+      "loss": 0.0851,
+      "step": 820
+    },
+    {
+      "grad_norm": 0.5591674447059631,
+      "learning_rate": 9.970436496555617e-05,
+      "loss": 0.1025,
+      "step": 830
+    },
+    {
+      "grad_norm": 0.5183696746826172,
+      "learning_rate": 9.968613921540373e-05,
+      "loss": 0.0941,
+      "step": 840
+    },
+    {
+      "grad_norm": 0.5633620023727417,
+      "learning_rate": 9.966737010559326e-05,
+      "loss": 0.0963,
+      "step": 850
+    },
+    {
+      "grad_norm": 0.6012111306190491,
+      "learning_rate": 9.964805784138072e-05,
+      "loss": 0.0751,
+      "step": 860
+    },
+    {
+      "grad_norm": 0.7718893885612488,
+      "learning_rate": 9.962820263396195e-05,
+      "loss": 0.1002,
+      "step": 870
+    },
+    {
+      "grad_norm": 0.43630078434944153,
+      "learning_rate": 9.960780470047033e-05,
+      "loss": 0.0739,
+      "step": 880
+    },
+    {
+      "grad_norm": 0.4448564052581787,
+      "learning_rate": 9.958686426397437e-05,
+      "loss": 0.0771,
+      "step": 890
+    },
+    {
+      "grad_norm": 0.47105884552001953,
+      "learning_rate": 9.956538155347534e-05,
+      "loss": 0.0833,
+      "step": 900
+    },
+    {
+      "grad_norm": 0.7328141927719116,
+      "learning_rate": 9.95433568039047e-05,
+      "loss": 0.0841,
+      "step": 910
+    },
+    {
+      "grad_norm": 0.45969247817993164,
+      "learning_rate": 9.952079025612162e-05,
+      "loss": 0.0719,
+      "step": 920
+    },
+    {
+      "grad_norm": 0.5842729210853577,
+      "learning_rate": 9.949768215691022e-05,
+      "loss": 0.0833,
+      "step": 930
+    },
+    {
+      "grad_norm": 0.6713224649429321,
+      "learning_rate": 9.9474032758977e-05,
+      "loss": 0.0975,
+      "step": 940
+    },
+    {
+      "grad_norm": 0.40328073501586914,
+      "learning_rate": 9.944984232094794e-05,
+      "loss": 0.0768,
+      "step": 950
+    },
+    {
+      "grad_norm": 0.4870326519012451,
+      "learning_rate": 9.942511110736584e-05,
+      "loss": 0.0797,
+      "step": 960
+    },
+    {
+      "grad_norm": 0.46975821256637573,
+      "learning_rate": 9.939983938868726e-05,
+      "loss": 0.1042,
+      "step": 970
+    },
+    {
+      "grad_norm": 0.4261704385280609,
+      "learning_rate": 9.93740274412797e-05,
+      "loss": 0.0757,
+      "step": 980
+    },
+    {
+      "grad_norm": 0.49296048283576965,
+      "learning_rate": 9.934767554741846e-05,
+      "loss": 0.07,
+      "step": 990
+    },
+    {
+      "grad_norm": 0.700587809085846,
+      "learning_rate": 9.932078399528361e-05,
+      "loss": 0.0773,
+      "step": 1000
+    },
+    {
+      "grad_norm": 0.48134827613830566,
+      "learning_rate": 9.929335307895689e-05,
+      "loss": 0.0712,
+      "step": 1010
+    },
+    {
+      "grad_norm": 0.48594099283218384,
+      "learning_rate": 9.926538309841839e-05,
+      "loss": 0.0887,
+      "step": 1020
+    },
+    {
+      "grad_norm": 0.6529817581176758,
+      "learning_rate": 9.923687435954334e-05,
+      "loss": 0.071,
+      "step": 1030
+    },
+    {
+      "grad_norm": 0.43099382519721985,
+      "learning_rate": 9.920782717409873e-05,
+      "loss": 0.0685,
+      "step": 1040
+    },
+    {
+      "grad_norm": 0.5061713457107544,
+      "learning_rate": 9.917824185973994e-05,
+      "loss": 0.0813,
+      "step": 1050
+    },
+    {
+      "grad_norm": 0.49007850885391235,
+      "learning_rate": 9.914811874000723e-05,
+      "loss": 0.0863,
+      "step": 1060
+    },
+    {
+      "grad_norm": 0.7527875900268555,
+      "learning_rate": 9.911745814432218e-05,
+      "loss": 0.0942,
+      "step": 1070
+    },
+    {
+      "grad_norm": 0.6349867582321167,
+      "learning_rate": 9.90862604079842e-05,
+      "loss": 0.0966,
+      "step": 1080
+    },
+    {
+      "grad_norm": 0.505797266960144,
+      "learning_rate": 9.90545258721667e-05,
+      "loss": 0.0783,
+      "step": 1090
+    },
+    {
+      "grad_norm": 0.7068039774894714,
+      "learning_rate": 9.90222548839135e-05,
+      "loss": 0.0747,
+      "step": 1100
+    },
+    {
+      "grad_norm": 0.4806417226791382,
+      "learning_rate": 9.898944779613495e-05,
+      "loss": 0.0969,
+      "step": 1110
+    },
+    {
+      "grad_norm": 0.5106626152992249,
+      "learning_rate": 9.89561049676041e-05,
+      "loss": 0.068,
+      "step": 1120
+    },
+    {
+      "grad_norm": 0.4418148994445801,
+      "learning_rate": 9.89222267629528e-05,
+      "loss": 0.0833,
+      "step": 1130
+    },
+    {
+      "grad_norm": 0.8706759810447693,
+      "learning_rate": 9.888781355266763e-05,
+      "loss": 0.107,
+      "step": 1140
+    },
+    {
+      "grad_norm": 0.5463252067565918,
+      "learning_rate": 9.885286571308598e-05,
+      "loss": 0.0902,
+      "step": 1150
+    },
+    {
+      "grad_norm": 0.3997989892959595,
+      "learning_rate": 9.881738362639182e-05,
+      "loss": 0.0878,
+      "step": 1160
+    },
+    {
+      "grad_norm": 0.46273931860923767,
+      "learning_rate": 9.878136768061154e-05,
+      "loss": 0.0846,
+      "step": 1170
+    },
+    {
+      "grad_norm": 0.5125852227210999,
+      "learning_rate": 9.874481826960979e-05,
+      "loss": 0.063,
+      "step": 1180
+    },
+    {
+      "grad_norm": 0.5445611476898193,
+      "learning_rate": 9.870773579308503e-05,
+      "loss": 0.0753,
+      "step": 1190
+    },
+    {
+      "grad_norm": 0.6660953760147095,
+      "learning_rate": 9.867012065656533e-05,
+      "loss": 0.0844,
+      "step": 1200
+    },
+    {
+      "grad_norm": 0.750414252281189,
+      "learning_rate": 9.863197327140376e-05,
+      "loss": 0.0946,
+      "step": 1210
+    },
+    {
+      "grad_norm": 0.6888939738273621,
+      "learning_rate": 9.859329405477403e-05,
+      "loss": 0.0792,
+      "step": 1220
+    },
+    {
+      "grad_norm": 0.47912099957466125,
+      "learning_rate": 9.855408342966585e-05,
+      "loss": 0.0813,
+      "step": 1230
+    },
+    {
+      "grad_norm": 0.36273863911628723,
+      "learning_rate": 9.851434182488033e-05,
+      "loss": 0.0791,
+      "step": 1240
+    },
+    {
+      "grad_norm": 0.6081685423851013,
+      "learning_rate": 9.84740696750253e-05,
+      "loss": 0.0687,
+      "step": 1250
+    },
+    {
+      "grad_norm": 0.4913078248500824,
+      "learning_rate": 9.843326742051055e-05,
+      "loss": 0.0788,
+      "step": 1260
+    },
+    {
+      "grad_norm": 0.5190815925598145,
+      "learning_rate": 9.839193550754297e-05,
+      "loss": 0.0606,
+      "step": 1270
+    },
+    {
+      "grad_norm": 0.35437485575675964,
+      "learning_rate": 9.835007438812177e-05,
+      "loss": 0.0743,
+      "step": 1280
+    },
+    {
+      "grad_norm": 0.39313238859176636,
+      "learning_rate": 9.830768452003341e-05,
+      "loss": 0.076,
+      "step": 1290
+    },
+    {
+      "grad_norm": 0.41886866092681885,
+      "learning_rate": 9.826476636684671e-05,
+      "loss": 0.0759,
+      "step": 1300
+    },
+    {
+      "grad_norm": 0.541922926902771,
+      "learning_rate": 9.822132039790773e-05,
+      "loss": 0.074,
+      "step": 1310
+    },
+    {
+      "grad_norm": 0.595855176448822,
+      "learning_rate": 9.817734708833461e-05,
+      "loss": 0.0819,
+      "step": 1320
+    },
+    {
+      "grad_norm": 0.4398232698440552,
+      "learning_rate": 9.813284691901243e-05,
+      "loss": 0.0591,
+      "step": 1330
+    },
+    {
+      "grad_norm": 0.4390011727809906,
+      "learning_rate": 9.808782037658792e-05,
+      "loss": 0.0582,
+      "step": 1340
+    },
+    {
+      "grad_norm": 0.6022183299064636,
+      "learning_rate": 9.804226795346411e-05,
+      "loss": 0.0868,
+      "step": 1350
+    },
+    {
+      "grad_norm": 0.6155015826225281,
+      "learning_rate": 9.799619014779503e-05,
+      "loss": 0.0823,
+      "step": 1360
+    },
+    {
+      "grad_norm": 0.6430284976959229,
+      "learning_rate": 9.794958746348013e-05,
+      "loss": 0.0624,
+      "step": 1370
+    },
+    {
+      "grad_norm": 0.46535441279411316,
+      "learning_rate": 9.790246041015896e-05,
+      "loss": 0.0621,
+      "step": 1380
+    },
+    {
+      "grad_norm": 0.41016414761543274,
+      "learning_rate": 9.785480950320538e-05,
+      "loss": 0.0729,
+      "step": 1390
+    },
+    {
+      "grad_norm": 0.48556017875671387,
+      "learning_rate": 9.78066352637221e-05,
+      "loss": 0.0603,
+      "step": 1400
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 10000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 9223372036854775807,
+  "save_steps": 200,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 0.0,
+  "train_batch_size": 8,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:782285ad74aaef2aad4e3e5daa440fd31ceb61526953366bc68d03be12e58f0b
+size 7633

wandb_config.json

@@ -0,0 +1 @@
+{"project": "finetune-gr00t-n1d6", "run_id": "models"}

zero_to_fp32.py

@@ -0,0 +1,760 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+#   python zero_to_fp32.py . output_dir/
+#   or
+#   python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device, weights_only=False)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+    total_files = len(files)
+    state_dicts = []
+    for f in tqdm(files, desc='Loading checkpoint shards'):
+        state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage <= 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+    """
+    A pseudo tensor that collects partitioned weights.
+    It is more memory efficient when there are multiple groups.
+    """
+
+    def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+        self.flat_groups = flat_groups
+        self.flat_groups_offset = flat_groups_offset
+        self.offset = offset
+        self.partitioned_numel = partitioned_numel
+        self.shape = shape
+        self.dtype = self.flat_groups[0][0].dtype
+
+    def contiguous(self):
+        """
+        Merge partitioned weights from flat_groups into a single tensor.
+        """
+        end_idx = self.offset + self.partitioned_numel
+        world_size = len(self.flat_groups)
+        pad_flat_param_chunks = []
+
+        for rank_i in range(world_size):
+            # for each rank, we need to collect weights from related group/groups
+            flat_groups_at_rank_i = self.flat_groups[rank_i]
+            start_group_id = None
+            end_group_id = None
+            for group_id in range(len(self.flat_groups_offset)):
+                if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+                    start_group_id = group_id
+                if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+                    end_group_id = group_id
+                    break
+            # collect weights from related group/groups
+            for group_id in range(start_group_id, end_group_id + 1):
+                flat_tensor = flat_groups_at_rank_i[group_id]
+                start_offset = self.offset - self.flat_groups_offset[group_id]
+                end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+                pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+        # collect weights from all ranks
+        pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+        param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+        return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+    for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # memory efficient tensor
+        tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+        state_dict[name] = tensor
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+    """
+    Convert state_dict of GatheredTensor to torch tensor
+    """
+    torch_state_dict = {}
+    converted_tensors = {}
+    for name, tensor in state_dict.items():
+        tensor_id = id(tensor)
+        if tensor_id in converted_tensors:  # shared tensors
+            shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+            torch_state_dict[name] = shared_tensor
+        else:
+            converted_tensors[tensor_id] = name
+            if return_empty_tensor:
+                torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+            else:
+                torch_state_dict[name] = tensor.contiguous()
+    return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                             tag=None,
+                                             exclude_frozen_parameters=False,
+                                             lazy_mode=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+        - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+          Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+    You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint. Or you can load state_dict in lazy mode ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+        for name, lazy_tensor in state_dict.item():
+            tensor = lazy_tensor.contiguous()  # to cpu
+            print(name, tensor)
+            # del tensor to release memory if it no longer in use
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+    if lazy_mode:
+        return state_dict
+    else:
+        return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+                                               output_dir,
+                                               max_shard_size="5GB",
+                                               safe_serialization=False,
+                                               tag=None,
+                                               exclude_frozen_parameters=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_dir``: directory to the pytorch fp32 state_dict output files
+        - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+        - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+    """
+
+    # Dependency pre-check
+    if safe_serialization:
+        try:
+            from safetensors.torch import save_file
+        except ImportError:
+            print('If you want to use `safe_serialization`, please `pip install safetensors`')
+            raise
+    if max_shard_size is not None:
+        try:
+            from huggingface_hub import split_torch_state_dict_into_shards
+        except ImportError:
+            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+            raise
+
+    # Convert zero checkpoint to state_dict
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                                          tag,
+                                                          exclude_frozen_parameters,
+                                                          lazy_mode=True)
+
+    # Shard the model if it is too big.
+    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+    if max_shard_size is not None:
+        filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+        # an memory-efficient approach for sharding
+        empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+        state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+                                                              filename_pattern=filename_pattern,
+                                                              max_shard_size=max_shard_size)
+    else:
+        from collections import namedtuple
+        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+        state_dict_split = StateDictSplit(is_sharded=False,
+                                          filename_to_tensors={weights_name: list(state_dict.keys())})
+
+    # Save the model by shard
+    os.makedirs(output_dir, exist_ok=True)
+    filename_to_tensors = state_dict_split.filename_to_tensors.items()
+    for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+        shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+        shard_state_dict = to_torch_tensor(shard_state_dict)
+        output_path = os.path.join(output_dir, shard_file)
+        if safe_serialization:
+            save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+        else:
+            torch.save(shard_state_dict, output_path)
+        # release the memory of current shard
+        for tensor_name in list(shard_state_dict.keys()):
+            del state_dict[tensor_name]
+            del shard_state_dict[tensor_name]
+        del shard_state_dict
+        gc.collect()
+
+    # Save index if sharded
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+        save_index_file = os.path.join(output_dir, save_index_file)
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info("Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info("Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument("output_dir",
+                        type=str,
+                        help="directory to the pytorch fp32 state_dict output files"
+                        "(e.g. path/checkpoint-12-output/)")
+    parser.add_argument(
+        "--max_shard_size",
+        type=str,
+        default="5GB",
+        help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+        "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+        "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+        "without CPU OOM issues.")
+    parser.add_argument(
+        "--safe_serialization",
+        default=False,
+        action='store_true',
+        help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+    parser.add_argument("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+    parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+                                               args.output_dir,
+                                               max_shard_size=args.max_shard_size,
+                                               safe_serialization=args.safe_serialization,
+                                               tag=args.tag,
+                                               exclude_frozen_parameters=args.exclude_frozen_parameters)