[feat] Support DAPO (#6263)

* update help information

* update style

* fix

* minor fix

* support PP training

* add pp support

* remove unused code

* address conversation

* fix memory leakage support tp+pp

* move empty cache

* move empty cache

* add DAPO support

* remove format reward

* fix filtering, still buggy

* small fix

* add DAPO support

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tested multi-node training; fix bind_batch bug

* fix conversation; support sleep mode

* support reusing excessive samples

* add dynamic batching control flag

* add dynamic batching control flag

* refactored

* fix logging

---------

Co-authored-by: Tong Li <tong.li35271158@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

applications/ColossalChat/rl_example.py

@@ -1,4 +1,5 @@
 import argparse
+import os
 
 import ray
 import torch
@@ -8,15 +9,17 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("-m", "--model", type=str, default="Qwen/Qwen2.5-7B")
     parser.add_argument("-d", "--dataset", type=str, default="data.jsonl")
+    parser.add_argument("-p", "--project", type=str, default="GRPO", help="Project name.")
+
+    # Distributed training parameters
     parser.add_argument("-t", "--num-trainers", type=int, default=2)
     parser.add_argument("-i", "--num-inferencer", type=int, default=2)
     parser.add_argument("-g", "--num-generations", type=int, default=8, help="Number of generations per prompt.")
-    parser.add_argument("-p", "--project", type=str, default="GRPO", help="Project name.")
     parser.add_argument(
         "-ibs",
         "--inference-batch-size",
         type=int,
-        default=64,
+        default=None,
         help="Number of prompts to generate per inference step. It should be divisible by tbs, and the weights on the inference backend will be synced every ibs/tbs training steps of the policy model.",
     )
     parser.add_argument(
@@ -37,7 +40,7 @@ if __name__ == "__main__":
         "-tMbs",
         "--train-minibatch-size",
         type=int,
-        default=1,
+        default=None,
         help="Number of unique prompts in each training batch per dp group. The inference backend must generate tMbs * g * dp_size samples before forwarding. Satisfy tMbs * g >= tmbs",
     )
     parser.add_argument(
@@ -47,22 +50,78 @@ if __name__ == "__main__":
         default=2,
         help="Effective batch size per dp group for forwarding and backwarding. Please select based on the availiable memory.",
     )
+    parser.add_argument(
+        "--ray_dir", type=str, default=None, help="Custom temperary directory for storing ray cluster data, Optional"
+    )
+    parser.add_argument(
+        "--master_address", type=str, default=None, help="Master address for multi-node distributed training, Optional"
+    )
+    parser.add_argument(
+        "--master_port", type=int, default=29506, help="Master port for multi-node distributed training, Optional"
+    )
+
+    # Sampling parameters
     parser.add_argument("-b", "--backend", type=str, default="transformers", choices=["transformers", "vllm"])
-    parser.add_argument("-a", "--algo", type=str, default="GRPO", choices=["Simple", "GRPO", "EvalGRPO"])
+    parser.add_argument("-temp", "--temperature", type=float, default=1.0, help="Temperature for sampling.")
+    parser.add_argument(
+        "-topk",
+        "--top-k",
+        type=int,
+        default=None,
+        help="Top k for sampling. Please check the generation arguments documentation for your backend.",
+    )
+    parser.add_argument(
+        "-topp",
+        "--top-p",
+        type=float,
+        default=1.0,
+        help="Top p for sampling. Please check the generation arguments documentation for your backend.",
+    )
     parser.add_argument("-s", "--system-prompt", type=str, default=None, help="System prompt for data construction.")
+    parser.add_argument("-mnt", "--max-new-tokens", type=int, default=1024 * 4 - 512, help="Max length for generation.")
+    parser.add_argument("-mpt", "--max-prompt-tokens", type=int, default=512, help="Max length for prompt.")
+
+    # GRPO parameters
+    parser.add_argument("-a", "--algo", type=str, default="GRPO", choices=["DAPO", "GRPO"])
+    parser.add_argument("-lr", "--learning-rate", type=float, default=1e-6, help="Learning rate for GRPO.")
+    parser.add_argument("-kl", "--kl-coeff", type=float, default=0.01, help="KL penalty coefficient for GRPO.")
+
+    # Logging/Checkpointing parameters
+    parser.add_argument("-si", "--save-interval", type=int, default=100, help="Interval for saving checkpoints.")
+    parser.add_argument("-sd", "--save-dir", type=str, default="./model", help="Directory for saving checkpoints.")
+
     args = parser.parse_args()
 
-    assert args.train_minibatch_size > 0, "Train mini batch size must be greater than 0"
+    if args.train_minibatch_size is None:
+        # Default settings: Using train batch size as mini batch size
+        args.train_minibatch_size = args.train_batch_size
+    if args.inference_batch_size is None:
+        # Default settings: Using train batch size as inference batch size, sync every inference model every train step
+        args.inference_batch_size = args.train_batch_size
     assert (
         args.train_minibatch_size * args.num_generations >= args.train_microbatch_size
         and args.train_microbatch_size > 0
     ), "Train micro batch size must be greater than 0 less than train mini batch size * num generations"
+    assert (
+        args.train_minibatch_size <= args.train_batch_size
+    ), "Train mini batch size must be less than or equals to train batch size"
 
-    ray.init(address="local", namespace="ray-example")
+    if args.master_address is None:
+        # Default settings: Using single machine
+        ray.init(address="local", namespace="ray-example")
+    else:
+        # For ray distributed multi-machine training, Please change _node_ip_address to your IP address of your master node
+        ray.init(_node_ip_address=args.master_address, namespace="ray-example", _temp_dir=args.ray_dir)
+
+    if args.top_k is None:
+        if args.backend == "transformers":
+            args.top_k = 50
+        elif args.backend == "vllm":
+            args.top_k = -1
 
     inference_model_config = dict(path=args.model)
     train_model_config = dict(path=args.model, use_flash_attention_2=True, use_cache=False)
-    generate_config = dict(top_k=50, top_p=0.75, temperature=0.9)
+    generate_config = dict(top_k=args.top_k, top_p=args.top_p, temperature=args.temperature)
 
     if args.backend == "transformers":
         inference_model_config.update(
@@ -73,7 +132,7 @@ if __name__ == "__main__":
         )
         generate_config.update(
             dict(
-                max_length=1024 + 512,
+                max_length=args.max_new_tokens + args.max_prompt_tokens,
                 do_sample=True,
                 max_new_tokens=None,
                 early_stopping=False,
@@ -81,31 +140,57 @@ if __name__ == "__main__":
             )
         )
     elif args.backend == "vllm":
-        inference_model_config.update(dict(gpu_memory_utilization=0.7, enforce_eager=True, enable_chunked_prefill=True))
+        inference_model_config.update(
+            dict(
+                gpu_memory_utilization=0.7,
+                enforce_eager=True,
+                enable_chunked_prefill=True,
+                max_model_len=args.max_new_tokens + args.max_prompt_tokens,
+                tensor_parallel_size=1,
+            )
+        )
         generate_config.update(
             dict(
-                max_tokens=2048,
+                max_tokens=args.max_new_tokens,  # max new tokens
                 ignore_eos=True,
                 include_stop_str_in_output=True,
                 stop=["</answer>"],
             )
         )
     else:
-        inference_model_config.update(
-            dict(
-                mem_fraction_static=0.6,
-            )
-        )
-        generate_config.update(
-            dict(
-                max_new_tokens=256,
-                ignore_eos=True,
-            )
-        )
+        raise ValueError(f"Unsupported backend: {args.backend}")
+
+    if args.algo == "GRPO":
+        # Default Settings
+        grpo_config = {
+            "lr": args.learning_rate,
+            "train_microbatch_size": args.train_microbatch_size,
+            "beta": args.kl_coeff,  # KL penalty coefficient
+            "loss_variation": "sample_level",
+        }
+    elif args.algo == "DAPO":
+        # DAPO variant settings
+        grpo_config = {
+            "filter_range": [0.01, 0.99],  # only filter out all zero batch and all one batch
+            "lr": args.learning_rate,
+            "train_microbatch_size": args.train_microbatch_size,
+            "dynamic_batching": True,
+            "clip_eps_low": 0.2,
+            "clip_eps_high": 0.28,
+            "skip_threshold": 20.0,
+            "beta": 0,  # no KL penalty for DAPO
+            "loss_variation": "token_level",
+            "soft_over_length_punishment": True,
+            "max_length": args.max_new_tokens + args.max_prompt_tokens,
+            "cache_length": min(1024, int(args.max_new_tokens / 4)),
+            "filter_truncated_response": True,
+        }
+    else:
+        raise ValueError(f"Unsupported algorithm: {args.algo}")
 
     launch_distributed(
         num_producers=args.num_inferencer,
-        num_proc_per_producer=1,
+        num_proc_per_producer=inference_model_config.get("tensor_parallel_size", 1),
         num_consumer_procs=args.num_trainers,
         num_episodes=1,
         inference_batch_size=args.inference_batch_size,
@@ -113,15 +198,22 @@ if __name__ == "__main__":
         train_batch_size=args.train_batch_size,
         train_minibatch_size=args.train_minibatch_size,
         train_microbatch_size=args.train_microbatch_size,
-        dataset_config={"path": args.dataset, "max_length": 300, "system_prompt": args.system_prompt},
+        dataset_config={
+            "path": args.dataset,
+            "max_length": args.max_prompt_tokens,
+            "system_prompt": args.system_prompt,
+        },
         dataloaders_config={},
         inference_model_config=inference_model_config,
         generate_config=generate_config,
         num_generations=args.num_generations,
         train_model_config=train_model_config,
-        plugin_config={},  # Default setting: zero.
+        grpo_config=grpo_config,
+        plugin_config={
+            "zero_stage": 2,
+        },  # for zero
+        # currently not support tp/pp
         # plugin_config={
-        #     "pp_size": 2,
         #     "tp_size": 2,
         #     "microbatch_size": args.train_microbatch_size // 2,
         #     "zero_stage": 0,
@@ -129,7 +221,9 @@ if __name__ == "__main__":
         # },  # for pp
         inference_backend=args.backend,
         master_addr="localhost",
-        master_port=29506,
+        master_port=args.master_port,
         core_algo=args.algo,
         project_name=args.project,
+        save_interval=args.save_interval,
+        save_dir=os.path.join(args.save_dir, args.project.replace(" ", "_")),
     )