File size: 7,724 Bytes
2bf51a1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
{
    "exp": "Wan+JEPA",
    "seed": 1039,
    "auto_tag_run": false,
    "eval": false,
    "debug": false,
    "timestamp": "\"2606150056\"",
    "log": {
        "logging_dir": "logs",
        "report_to": "wandb",
        "log_freq": 100
    },
    "wandb": {
        "project": "psi",
        "entity": "keycharon",
        "group": "finetune",
        "id": "wtg4c4ql",
        "name": "Wan+JEPA.wholeb.acvjepa.nullstate.cosine.lr1.0e-04.b256.gpus8.\"2606150056\"",
        "resume": "allow"
    },
    "train": {
        "num_workers": 8,
        "overfit_single_batch": false,
        "name": "finetune",
        "resume_from_checkpoint": null,
        "skip_resumed_steps": false,
        "hf_token": ".hf_token",
        "lora": false,
        "output_dir": "/mnt/data0/runs",
        "gradient_accumulation_steps": 1,
        "mixed_precision": "bf16",
        "max_grad_norm": 1.0,
        "optimizer_foreach": null,
        "train_batch_size": 32,
        "val_batch_size": 16,
        "val_num_batches": 20,
        "checkpointing_steps": 100,
        "max_checkpoints_to_keep": 5,
        "validation_steps": 0,
        "learning_rate": 0.0001,
        "lr_scheduler_type": "cosine",
        "lr_scheduler_kwargs": {
            "betas": [
                0.95,
                0.999
            ],
            "weight_decay": 1e-6,
            "eps": 1e-8
        },
        "scheduler_specific_kwargs": {},
        "data_parallel": "deepspeed",
        "sharding_strategy": "full-shard",
        "deepspeed_config": "/mnt/data0/Psi0/scripts/deepspeed/zero3.json",
        "enable_gradient_checkpointing": true,
        "enable_mixed_precision_training": true,
        "reduce_in_full_precision": true,
        "max_training_steps": 100000,
        "num_train_epochs": null,
        "warmup_steps": 1000,
        "warmup_ratio": null
    },
    "data": {
        "transform": {
            "repack": {
                "dataset_name": "wholebody",
                "stage": "postpre",
                "action_chunk_size": 25,
                "pad_action_dim": 66,
                "pad_state_dim": 45,
                "num_past_frames": 1,
                "use_delta_actions": false
            },
            "model": {
                "resize": {
                    "size": [
                        270,
                        480
                    ]
                },
                "center_crop": {
                    "size": [
                        270,
                        480
                    ]
                },
                "jepa_resize": {
                    "size": [
                        256,
                        256
                    ]
                },
                "jepa_center_crop": {
                    "size": [
                        256,
                        256
                    ]
                },
                "color_jitter": {
                    "brightness": 0.2,
                    "contrast": [
                        0.8,
                        1.2
                    ],
                    "saturation": [
                        0.8,
                        1.2
                    ],
                    "hue": 0.05
                },
                "gaussian_noise": {
                    "mean": 0.0,
                    "std": 3.0,
                    "prob_skip": 0.1
                },
                "img_aug": false,
                "adaptive_resize": false,
                "img_sizes": {
                    "egodex": [
                        270,
                        480
                    ],
                    "he": [
                        240,
                        320
                    ]
                }
            },
            "field": {
                "stat_path": "/home/baai/wam_ckpt/stats/wam_stat_no_human.json",
                "state_stat_path": "/home/baai/wam_ckpt/stats/wam_stat_state_no_human.json",
                "stat_action_key": "wam_smpl",
                "stat_state_key": "wam_smpl",
                "action_norm_type": "std",
                "normalize_state": true,
                "state_norm_type": "bounds_q99",
                "clip_value": null,
                "eps": 1e-6,
                "low": null,
                "high": null
            }
        },
        "root_dir": "/mnt/data0/posttrain_data_new",
        "annotation_dir": null,
        "ego_video_dir": null,
        "exo_video_dir": null,
        "require_exo_video": false,
        "upsample_rate": 1,
        "chunk_size": 25,
        "num_past_frames": 0,
        "img_history_size": 1,
        "use_delta_actions": false,
        "data_downsample": 1,
        "force_rewrite_cache": false
    },
    "model": {
        "resnet_store_path": null,
        "pretrained_action_header_path": null,
        "rtc": true,
        "max_delay": 8,
        "action_dim": 66,
        "action_chunk_size": 25,
        "action_exec_horizon": 25,
        "observation_horizon": 1,
        "img_chunk": 1,
        "n_cams": 1,
        "use_obs": "add_token",
        "dropout": 0.1,
        "noise_scheduler": "flow",
        "train_diffusion_steps": 1000,
        "eval_diffusion_steps": 10,
        "share_cam_features": false,
        "early_fusion": false,
        "odim": 66,
        "n_conditions": 1,
        "token_fusion": "concat",
        "loss_w": [
            0.1,
            0.2,
            0.1
        ],
        "time_dim": 256,
        "hidden_dim": 1536,
        "num_blocks": 6,
        "dim_feedforward": 2048,
        "nhead": 24,
        "activation": "gelu",
        "view_feature_dim": 1920,
        "use_film": false,
        "combined_temb": false,
        "use_dit": false,
        "weight_decay": 0.01,
        "model_name_or_path": "/home/baai/wam_ckpt/omni_cross_38900",
        "vlm_ckpt_step": null,
        "tune_vlm": false,
        "tune_mm_llm": false,
        "tune_mm_vision": false,
        "tune_mm_mlp": false,
        "gradient_checkpointing": true,
        "lang_backbone_lr": 0.00001,
        "mm_projector_lr": 0.00001,
        "vision_tower_lr": 1e-6,
        "optim": "adamw_torch",
        "model_max_length": 4096,
        "data_flatten": true,
        "data_packing": true,
        "max_pixels": 451584,
        "min_pixels": 12544,
        "posttrain_mode": "ac_vjepa",
        "ac_vjepa_jepa_loss_weight": 1.0,
        "ac_vjepa_action_loss_weight": 1.0,
        "ac_vjepa_abs_action_loss_weight": 0.0,
        "ac_vjepa_use_action_tokenizer": false,
        "ac_vjepa_action_tokenizer_type": "fast",
        "ac_vjepa_action_tokenizer_ckpt": null,
        "ac_vjepa_action_tokenizer_bins": 2048,
        "ac_vjepa_action_token_loss_weight": 0.0,
        "ac_vjepa_action_token_use_absolute": true,
        "ac_vjepa_hidden_dim": 768,
        "ac_vjepa_depth": 12,
        "ac_vjepa_num_heads": 16,
        "ac_vjepa_latent_dim": 768,
        "ac_vjepa_state_dim": null,
        "ac_vjepa_use_state": false,
        "ac_vjepa_force_null_state": true,
        "ac_vjepa_state_dropout_prob": 0.1,
        "ac_vjepa_state_dropout_prob_end": 0.8,
        "ac_vjepa_state_dropout_anneal_steps": 30000,
        "ac_vjepa_image_size": 256,
        "ac_vjepa_patch_size": 16,
        "ac_vjepa_encoder_hidden_dim": 1408,
        "ac_vjepa_encoder_depth": 40,
        "ac_vjepa_encoder_heads": 16,
        "ac_vjepa_encoder_mlp_ratio": 4.363636363636363,
        "ac_vjepa_encoder_use_rope": true,
        "ac_vjepa_encoder_ckpt_key": "target_encoder",
        "ac_vjepa_predictor_ckpt_key": "predictor",
        "ac_vjepa_load_predictor_from_ckpt": false,
        "ac_vjepa_freeze_frame_encoder": true,
        "ac_vjepa_force_predictor_trainable": true
    }
}