Initial commit

arguments.json

@@ -0,0 +1,87 @@
+{
+    "actor_model_name_or_path": "output/sft_opt_1b",
+    "reward_model_name_or_path": "output/rm_opt",
+    "cost_model_name_or_path": "output/cm_opt",
+    "reward_critic_model_name_or_path": "output/rm_opt",
+    "cost_critic_model_name_or_path": "output/cm_opt",
+    "max_length": 512,
+    "trust_remote_code": true,
+    "train_datasets": [
+        [
+            "PKU-SafeRLHF/train",
+            {
+                "proportion": 1.0
+            }
+        ]
+    ],
+    "ptx_datasets": [
+        [
+            "alpaca",
+            {
+                "proportion": 1.0
+            }
+        ]
+    ],
+    "eval_datasets": null,
+    "threshold": 0.0,
+    "lambda_init": 1.0,
+    "lambda_lr": 0.1,
+    "lambda_max": 5.0,
+    "lambda_update_delay_steps": 0,
+    "episode_cost_window_size": 128,
+    "kl_coeff": 0.01,
+    "clip_range_ratio": 0.2,
+    "clip_range_score": 50.0,
+    "clip_range_value": 5.0,
+    "ptx_coeff": 16.0,
+    "epochs": 1,
+    "update_iters": 1,
+    "per_device_prompt_batch_size": 16,
+    "per_device_train_batch_size": 16,
+    "per_device_eval_batch_size": 16,
+    "gradient_accumulation_steps": 2,
+    "actor_lr": 1e-05,
+    "actor_weight_decay": 0.01,
+    "actor_lr_scheduler_type": "cosine",
+    "actor_lr_warmup_ratio": 0.03,
+    "actor_gradient_checkpointing": true,
+    "critic_lr": 5e-06,
+    "critic_weight_decay": 0.0,
+    "critic_lr_scheduler_type": "constant",
+    "critic_lr_warmup_ratio": 0.03,
+    "critic_gradient_checkpointing": true,
+    "normalize_reward": false,
+    "normalize_cost": false,
+    "seed": 42,
+    "fp16": true,
+    "bf16": false,
+    "tf32": null,
+    "temperature": 1.0,
+    "top_p": 1.0,
+    "num_return_sequences": 1,
+    "repetition_penalty": 1.0,
+    "eval_strategy": "epoch",
+    "eval_interval": 1000000,
+    "need_eval": false,
+    "eval_split_ratio": null,
+    "output_dir": "/home/rwliang/Reinforcement-Learning/FINAL_PROJECT/rl-final/safe-rlhf/output/ppo-lag",
+    "log_type": "wandb",
+    "log_dir": "/home/rwliang/Reinforcement-Learning/FINAL_PROJECT/rl-final/safe-rlhf/output/ppo-lag",
+    "log_project": "Safe-RLHF-PPO",
+    "log_run_name": "ppo_lag-2024-12-04-22-29-24",
+    "save_16bit": false,
+    "save_interval": 1000000,
+    "local_rank": 0,
+    "zero_stage": 1,
+    "offload": "none",
+    "deepspeed": false,
+    "deepspeed_config": null,
+    "deepscale": false,
+    "deepscale_config": null,
+    "global_rank": 0,
+    "device": {
+        "type": "torch.device",
+        "repr": "device(type='cuda', index=0)"
+    },
+    "total_training_steps": 806
+}

arguments.pkl

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

config.json

@@ -0,0 +1,31 @@
+{
+  "_name_or_path": "output/sft_opt_1b",
+  "_remove_final_layer_norm": false,
+  "activation_dropout": 0.0,
+  "activation_function": "relu",
+  "architectures": [
+    "OPTForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 2,
+  "do_layer_norm_before": true,
+  "dropout": 0.1,
+  "enable_bias": true,
+  "eos_token_id": 2,
+  "ffn_dim": 8192,
+  "hidden_size": 2048,
+  "init_std": 0.02,
+  "layer_norm_elementwise_affine": true,
+  "layerdrop": 0.0,
+  "max_position_embeddings": 2048,
+  "model_type": "opt",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 24,
+  "pad_token_id": 1,
+  "prefix": "</s>",
+  "torch_dtype": "float16",
+  "transformers_version": "4.45.2",
+  "use_cache": true,
+  "vocab_size": 50272,
+  "word_embed_proj_dim": 2048
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 2,
+  "eos_token_id": 2,
+  "pad_token_id": 1,
+  "transformers_version": "4.45.2"
+}

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,395 @@
+{
+  "metadata": {
+    "total_size": 5056032768
+  },
+  "weight_map": {
+    "model.decoder.embed_positions.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.0.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.1.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.10.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.11.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.12.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.13.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.14.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.15.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.16.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.17.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.18.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.19.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.2.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.20.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.21.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.22.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.decoder.layers.23.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.decoder.layers.23.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.decoder.layers.23.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.decoder.layers.23.final_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "model.decoder.layers.23.final_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.23.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.3.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.4.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.5.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.6.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.7.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.8.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.fc1.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.fc1.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.fc2.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.fc2.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.final_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "model.decoder.layers.9.self_attn_layer_norm.weight": "model-00001-of-00002.safetensors"
+  }
+}

script.sh

@@ -0,0 +1,199 @@
+#!/usr/bin/env bash
+#
+# Copyright 2023-2024 PKU-Alignment Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+if [ -z "${BASH_VERSION}" ]; then
+	echo "Please use bash to run this script." >&2
+	exit 1
+fi
+
+set -x
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" &>/dev/null && pwd)"
+ROOT_DIR="$(dirname "${SCRIPT_DIR}")"
+export PYTHONPATH="${ROOT_DIR}${PYTHONPATH:+:${PYTHONPATH}}"
+export LOGLEVEL="${LOGLEVEL:-WARNING}"
+
+ACTOR_MODEL_NAME_OR_PATH="PKU-Alignment/alpaca-7b-reproduced"
+REWARD_MODEL_NAME_OR_PATH="${ROOT_DIR}/output/rm"
+COST_MODEL_NAME_OR_PATH="${ROOT_DIR}/output/cm"
+unset {REWARD,COST}_CRITIC_MODEL_NAME_OR_PATH
+OUTPUT_DIR="${ROOT_DIR}/output/ppo-lag"
+unset HOSTFILE
+ZERO_STAGE=1
+OFFLOAD="none"
+while [[ "$#" -gt 0 ]]; do
+	arg="$1"
+	shift
+	case "${arg}" in
+		--actor_model_name_or_path)
+			ACTOR_MODEL_NAME_OR_PATH="$1"
+			shift
+			;;
+		--actor_model_name_or_path=*)
+			ACTOR_MODEL_NAME_OR_PATH="${arg#*=}"
+			;;
+		--reward_model_name_or_path)
+			REWARD_MODEL_NAME_OR_PATH="$1"
+			shift
+			;;
+		--reward_model_name_or_path=*)
+			REWARD_MODEL_NAME_OR_PATH="${arg#*=}"
+			;;
+		--reward_critic_model_name_or_path)
+			REWARD_CRITIC_MODEL_NAME_OR_PATH="$1"
+			shift
+			;;
+		--reward_critic_model_name_or_path=*)
+			REWARD_CRITIC_MODEL_NAME_OR_PATH="${arg#*=}"
+			;;
+		--cost_model_name_or_path)
+			COST_MODEL_NAME_OR_PATH="$1"
+			shift
+			;;
+		--cost_model_name_or_path=*)
+			COST_MODEL_NAME_OR_PATH="${arg#*=}"
+			;;
+		--cost_critic_model_name_or_path)
+			COST_CRITIC_MODEL_NAME_OR_PATH="$1"
+			shift
+			;;
+		--cost_critic_model_name_or_path=*)
+			COST_CRITIC_MODEL_NAME_OR_PATH="${arg#*=}"
+			;;
+		--output_dir)
+			OUTPUT_DIR="$1"
+			shift
+			;;
+		--output_dir=*)
+			OUTPUT_DIR="${arg#*=}"
+			;;
+		--hostfile)
+			HOSTFILE="$1"
+			shift
+			;;
+		--hostfile=*)
+			HOSTFILE="${arg#*=}"
+			;;
+		--zero_stage)
+			ZERO_STAGE="$1"
+			shift
+			;;
+		--zero_stage=*)
+			ZERO_STAGE="${arg#*=}"
+			;;
+		--offload)
+			OFFLOAD="$1"
+			shift
+			;;
+		--offload=*)
+			OFFLOAD="${arg#*=}"
+			;;
+		*)
+			echo "Unknown parameter passed: '${arg}'" >&2
+			exit 1
+			;;
+	esac
+done
+
+if [[ -z "${REWARD_CRITIC_MODEL_NAME_OR_PATH+x}" ]]; then
+	REWARD_CRITIC_MODEL_NAME_OR_PATH="${REWARD_MODEL_NAME_OR_PATH}"
+fi
+if [[ -z "${COST_CRITIC_MODEL_NAME_OR_PATH+x}" ]]; then
+	COST_CRITIC_MODEL_NAME_OR_PATH="${COST_MODEL_NAME_OR_PATH}"
+fi
+
+mkdir -p "${OUTPUT_DIR}"
+OUTPUT_DIR="$(cd "${OUTPUT_DIR}" &>/dev/null && pwd)"
+if [[ ! -f "${OUTPUT_DIR}/.gitignore" ]]; then
+	echo '*' >"${OUTPUT_DIR}/.gitignore"
+fi
+
+cp -f "$0" "${OUTPUT_DIR}/script.sh"
+
+if [[ -z "${WANDB_API_KEY}" ]]; then
+	export WANDB_MODE="offline"
+fi
+
+MASTER_PORT_START=10000
+MASTER_PORT_END=65535
+MASTER_PORT="$(
+	comm -23 \
+		<(seq "${MASTER_PORT_START}" "${MASTER_PORT_END}" | sort) \
+		<(ss -Htan | awk '{ print $4 }' | awk -F ':' '{ print $NF }' | sort -u) |
+		shuf | head -n 1
+)"
+
+DEEPSPEED_ARGS=()
+if [[ -n "${HOSTFILE+x}" ]]; then
+	DEEPSPEED_ARGS+=("--hostfile" "${HOSTFILE}")
+fi
+DEEPSPEED_ARGS+=("--master_port" "${MASTER_PORT}")
+
+exec 1> >(tee "${OUTPUT_DIR}/stdout.log" >&1) 2> >(tee "${OUTPUT_DIR}/stderr.log" >&2)
+
+deepspeed "${DEEPSPEED_ARGS[@]}" \
+	--master_port "${MASTER_PORT}" \
+	--module safe_rlhf.algorithms.ppo_lag \
+	--train_datasets PKU-SafeRLHF/train \
+	--ptx_datasets alpaca \
+	--actor_model_name_or_path "${ACTOR_MODEL_NAME_OR_PATH}" \
+	--reward_model_name_or_path "${REWARD_MODEL_NAME_OR_PATH}" \
+	--reward_critic_model_name_or_path "${REWARD_CRITIC_MODEL_NAME_OR_PATH}" \
+	--cost_model_name_or_path "${COST_MODEL_NAME_OR_PATH}" \
+	--cost_critic_model_name_or_path "${COST_CRITIC_MODEL_NAME_OR_PATH}" \
+	--max_length 512 \
+	--temperature 1.0 \
+	--num_return_sequences 1 \
+	--repetition_penalty 1.0 \
+	--trust_remote_code True \
+	--epochs 1 \
+	--update_iters 1 \
+	--per_device_prompt_batch_size 16 \
+	--per_device_train_batch_size 16 \
+	--gradient_accumulation_steps 1 \
+	--actor_lr 1e-5 \
+	--actor_weight_decay 0.01 \
+	--actor_lr_scheduler_type cosine \
+	--actor_lr_warmup_ratio 0.03 \
+	--actor_gradient_checkpointing \
+	--critic_lr 5e-6 \
+	--critic_weight_decay 0.0 \
+	--critic_lr_scheduler_type constant \
+	--critic_lr_warmup_ratio 0.03 \
+	--critic_gradient_checkpointing \
+	--normalize_reward False \
+	--normalize_cost False \
+	--seed 42 \
+	--threshold 0.0 \
+	--lambda_init 1.0 \
+	--lambda_lr 0.1 \
+	--lambda_max 5.0 \
+	--lambda_update_delay_steps 0 \
+	--episode_cost_window_size 128 \
+	--kl_coeff 0.01 \
+	--clip_range_ratio 0.2 \
+	--clip_range_score 50.0 \
+	--clip_range_value 5.0 \
+	--ptx_coeff 16.0 \
+	--output_dir "${OUTPUT_DIR}" \
+	--log_type wandb \
+	--log_project Safe-RLHF-PPO \
+	--zero_stage "${ZERO_STAGE}" \
+	--offload "${OFFLOAD}" \
+	--fp16 True
+	# --bf16 True \
+	# --tf32 True

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,35 @@
+{
+  "add_bos_token": true,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "</s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "errors": "replace",
+  "max_length": 512,
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "stride": 0,
+  "tokenizer_class": "GPT2Tokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "</s>"
+}