github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-13 13:45:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

merge

Browse Source
This commit is contained in:
YeAnbang 2025-04-29 16:09:44 +08:00
commit 8c5dd4131d
1 changed files with 11 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
applications/ColossalChat/rl_example.py
View File

@ -219,19 +219,18 @@ if __name__ == "__main__":
num_generations=args.num_generations,
train_model_config=train_model_config,
grpo_config=grpo_config,
# plugin_config={
# "zero_stage": 2,
# }, # for zero
# currently not support tp/pp
plugin_config={
"tp_size": 2,
"pp_size": 2,
"microbatch_size": max(
1, args.train_microbatch_size // 2
), # microbatch size should be set to train_microbatch_size // pp_size
"zero_stage": 0,
"max_norm": 1.0,
}, # for pp
"zero_stage": 2,
}, # for zero
# plugin_config={
# "tp_size": 2,
# "pp_size": 2,
# "microbatch_size": max(
# 1, args.train_microbatch_size // 2
# ), # microbatch size should be set to train_microbatch_size // pp_size
# "zero_stage": 0,
# "max_norm": 1.0,
# }, # for pp, tp
inference_backend=args.backend,
master_addr="localhost",
master_port=args.master_port,