[ColossalChat] Add PP support (#6001)

* support pp training

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

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

* update rm

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

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

* refactor

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

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

* update test case

* fix

* change to 4

* fix eval

* test

* add pp

* hotfix

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

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

* support pp training

* update rm

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

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

* refactor

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

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

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

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

* update test case

* fix

* change to 4

* fix eval

* test

* add pp

* hotfix

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

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

* update

* skip pp eval

* update all reduce

* update sft

* update ignore

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

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

* update no cache

* add eval

* remove fi

* remove debug

* remove parentheses to avoid warning

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

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

* Revert "add eval"

This reverts commit 3ab2f6fa32.

* add all reduce

---------

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

applications/ColossalChat/examples/training_scripts/train_dpo.py

@@ -267,6 +267,7 @@ def train(args):
         ref_model=ref_model,
         booster=booster,
         actor_optim=optim,
+        plugin=plugin,
         actor_lr_scheduler=lr_scheduler,
         tokenizer=tokenizer,
         max_epochs=args.max_epochs,

applications/ColossalChat/examples/training_scripts/train_kto.py

@@ -286,6 +286,7 @@ def train(args):
         ref_model=ref_model,
         booster=booster,
         actor_optim=optim,
+        plugin=plugin,
         actor_lr_scheduler=lr_scheduler,
         tokenizer=tokenizer,
         max_epochs=args.max_epochs,

applications/ColossalChat/examples/training_scripts/train_orpo.py

@@ -250,6 +250,7 @@ def train(args):
         actor=model,
         booster=booster,
         actor_optim=optim,
+        plugin=plugin,
         actor_lr_scheduler=lr_scheduler,
         tokenizer=tokenizer,
         max_epochs=args.max_epochs,

applications/ColossalChat/examples/training_scripts/train_rm.py

@@ -262,6 +262,7 @@ def train(args):
         model,
         booster,
         optim,
+        plugin,
         lr_scheduler,
         tokenizer,
         loss_fn=loss_fn,

applications/ColossalChat/examples/training_scripts/train_sft.py

@@ -114,7 +114,7 @@ def train(args):
             parallel_output=False,
             max_norm=args.grad_clip,
             precision=args.mixed_precision,
-            microbatch_size=args.batch_size,
+            microbatch_size=args.microbatch_size,
         )
     else:
         raise ValueError(f"Unknown plugin {args.plugin}")
@@ -269,6 +269,7 @@ def train(args):
         model=model,
         booster=booster,
         optim=optim,
+        plugin=plugin,
         lr_scheduler=lr_scheduler,
         max_epochs=args.max_epochs,
         accumulation_steps=args.accumulation_steps,
@@ -344,6 +345,7 @@ if __name__ == "__main__":
     parser.add_argument("--use_wandb", default=False, action="store_true")
     parser.add_argument("--grad_checkpoint", default=False, action="store_true")
     parser.add_argument("--use_flash_attn", default=False, action="store_true")
+    parser.add_argument("--microbatch_size", type=int, default=1)
     args = parser.parse_args()
     if args.config_file is not None:
         os.makedirs(os.path.dirname(args.config_file), exist_ok=True)