add a loss scaler

Author: binary-husky

ajet/backbone/verl/dp_actor.py

@@ -131,10 +131,13 @@ def update_policy(self, data: DataProto):
 
                     calculate_entropy = self.config.calculate_entropy or (entropy_coeff != 0)
 
-                    if self.config.use_dynamic_bsz:
+                    if self.config.override_ppo_mini_batch_num > 0:
+                        loss_scale_factor = response_mask.shape[0] / mini_batch_split_size
+                    elif self.config.use_dynamic_bsz:
                         loss_scale_factor = response_mask.shape[0] / self.config.ppo_mini_batch_size
                     else:
                         loss_scale_factor = 1 / self.gradient_accumulation
+                    loss_scale_factor *= self.config.loss_extra_scale_ratio  # [AJET] Extra scaling for loss if needed
 
                     # all return: (bsz, response_length)
                     outputs = self._forward_micro_batch(

ajet/default_config/ajet_default.yaml

@@ -288,6 +288,9 @@ ajet:
     kl_loss_coef: 0.002
     kl_loss_type: low_var_kl
 
+    # loss = loss * loss_extra_scale_ratio
+    loss_extra_scale_ratio: 1.0
+
     # Ulysses specific configs
     ulysses_sequence_parallel_size: 1
 

ajet/default_config/verl/config_auto_convertion_verl.jsonc

@@ -15,6 +15,7 @@
     "ajet.trainer_common.kl_loss_coef": "actor_rollout_ref.actor.kl_loss_coef",
     "ajet.trainer_common.kl_loss_type": "actor_rollout_ref.actor.kl_loss_type",
     "ajet.trainer_common.ulysses_sequence_parallel_size": "actor_rollout_ref.actor.ulysses_sequence_parallel_size",
+    "ajet.trainer_common.loss_extra_scale_ratio": "actor_rollout_ref.actor.loss_extra_scale_ratio",
 
     "ajet.trainer_common.save_freq": "trainer.save_freq",
     "ajet.trainer_common.test_freq": "trainer.test_freq",
@@ -30,6 +31,8 @@
         "actor_rollout_ref.ref.log_prob_max_token_len_per_gpu"
     ],
 
+    "ajet.rollout.max_num_seqs": "actor_rollout_ref.rollout.max_num_seqs",
+    "ajet.rollout.temperature": "actor_rollout_ref.rollout.temperature",
     "ajet.rollout.multi_turn": "actor_rollout_ref.rollout.multi_turn",
     "ajet.rollout.val_kwargs": "actor_rollout_ref.rollout.val_kwargs",
     "ajet.rollout.num_repeat": [

ajet/default_config/verl/verl_default.yaml

@@ -70,6 +70,7 @@ actor_rollout_ref:
     rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}
     strategy: fsdp
     ppo_mini_batch_size: 256
+    loss_extra_scale_ratio: 1.0
     override_ppo_mini_batch_num: 1 # special in agentjet
     ppo_micro_batch_size: null
     ppo_micro_batch_size_per_gpu: null

tests/bench/README.md

@@ -25,9 +25,13 @@ VERL_PYTHON="./.venv/bin/python" python -m pytest -s tests/bench/benchmark_count
 VERL_PYTHON="./.venv/bin/python" python -m pytest -s tests/bench/benchmark_learn2ask/execute_benchmark_learn2ask.py::TestBenchmarkLearnToAsk::test_01_begin_verl
 VERL_PYTHON="./.venv/bin/python" python -m pytest -s tests/bench/benchmark_frozenlake/execute_benchmark_frozenlake.py::TestBenchmarkFrozenLake::test_01_begin_verl
 
+python -m ajet.launcher --conf tests/bench/benchmark_math/benchmark_math.yaml --autokill --db="UPP"
 
 export APPWORLD_PATH="/dev/shm/pack_all_in_one"
 export APPWORLD_SCRIPT="bash EnvService/env_sandbox/appworld.sh"
 python -m ajet.launcher --conf tests/bench/benchmark_appworld/benchmark_appworld.yaml --with-appworld --backbone=debug --autokill
 python -m ajet.launcher --conf tests/bench/benchmark_appworld/benchmark_appworld_oai_sdk.yaml --with-appworld --autokill --db="EXT"
 ```
+
+
+VERL_PYTHON="./.venv/bin/python" python -m pytest -s tests/bench/benchmark_math/execute_benchmark_math.py::TestBenchmarkMath::test_01_begin_verl

tests/bench/benchmark_appworld/benchmark_appworld.yaml

@@ -47,12 +47,15 @@ ajet:
     max_prompt_length: 3000
     max_response_length: 15000
 
+  # trainer common configurations
   trainer_common:
     save_freq: 99999
     test_freq: 99999
     total_epochs: 99999
     nnodes: 1
     n_gpus_per_node: 8
+    # loss = loss * loss_extra_scale_ratio
+    loss_extra_scale_ratio: 10.0
 
   execute_test: True # DO NOT EDIT, THIS IS FOR TEST ROBOT
   execute_testing_lambda: "tests/bench/benchmark_appworld/benchmark_appworld.py->TestProbe" #

tests/bench/benchmark_countdown/benchmark_countdown.yaml

@@ -116,7 +116,8 @@ ajet:
     kl_loss_coef: 0.002
     kl_loss_type: low_var_kl
     ulysses_sequence_parallel_size: 1
-
+    # loss = loss * loss_extra_scale_ratio
+    loss_extra_scale_ratio: 10.0
 
   # DO NOT EDIT, FOR ROBOT TESTING PURPOSE ONLY. NOT FOR HUMAN.
   execute_test: True # FOR ROBOT TESTING PURPOSE ONLY. NOT FOR HUMAN.

tests/bench/benchmark_frozenlake/benchmark_frozenlake.yaml

@@ -69,7 +69,8 @@ ajet:
     nnodes: 1
     n_gpus_per_node: 8
     logger: swanlab
-
+    # loss = loss * loss_extra_scale_ratio
+    loss_extra_scale_ratio: 10.0
 
   execute_test: True
   execute_testing_lambda: "tests/bench/benchmark_frozenlake/benchmark_frozenlake.py->TestProbe"

tests/bench/benchmark_learn2ask/benchmark_learn2ask.yaml

@@ -45,7 +45,8 @@ ajet:
     test_freq: 100
     total_epochs: 100
     logger: swanlab
-
+    # loss = loss * loss_extra_scale_ratio
+    loss_extra_scale_ratio: 10.0
 
   execute_test: True # DO NOT EDIT, THIS IS FOR TEST ROBOT
   execute_testing_lambda: "tests/bench/benchmark_learn2ask/benchmark_learn2ask.py->TestProbe" # DO NOT EDIT, THIS IS FOR TEST ROBOT

tests/bench/benchmark_math/benchmark_math.yaml

@@ -21,7 +21,7 @@ ajet:
     user_workflow: "tutorial.example_math_agent.math_agent->ExampleMathLearn" # ✨✨✨✨ 编写并选择Agent
     temperature: 1.0
     max_env_worker: 64
-    max_num_seqs: 256
+    max_num_seqs: 10
     num_repeat: 6
     agent_madness_reward: 0.0
     tensor_model_parallel_size: 1
@@ -49,7 +49,8 @@ ajet:
     logger: swanlab
     nnodes: 1
     n_gpus_per_node: 4
-
+    # loss = loss * loss_extra_scale_ratio
+    loss_extra_scale_ratio: 40.0
 
 
   execute_test: True # DO NOT EDIT, THIS IS FOR TEST ROBOT