This commit is contained in:
YeAnbang 2025-04-29 16:09:44 +08:00
commit 8c5dd4131d

View File

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