github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 18:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

add prompt template (#6273)

Browse Source 
Co-authored-by: Tong Li <tong.li35271158@gmail.com>
This commit is contained in:
Tong Li
2025-04-22 10:39:47 +08:00
committed by GitHub
parent 9467c10690
commit 03f4b1dde3
2 changed files with 16 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
applications/ColossalChat/rl_example.py
View File

@@ -49,6 +49,7 @@ if __name__ == "__main__":
)
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("-s", "--system-prompt", type=str, default=None, help="System prompt for data construction.")
args = parser.parse_args()
assert args.train_minibatch_size > 0, "Train mini batch size must be greater than 0"
@@ -112,20 +113,20 @@ 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},
dataset_config={"path": args.dataset, "max_length": 300, "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={}, # for zero
plugin_config={
"pp_size": 2,
"tp_size": 2,
"microbatch_size": args.train_microbatch_size // 2,
"zero_stage": 0,
"max_norm": 1.0,
}, # for pp
plugin_config={}, # Default setting: zero.
# plugin_config={
# "pp_size": 2,
# "tp_size": 2,
# "microbatch_size": args.train_microbatch_size // 2,
# "zero_stage": 0,
# "max_norm": 1.0,
# }, # for pp
inference_backend=args.backend,
master_addr="localhost",
master_port=29506,