[hot-fix] Fix memory leakage bug, support TP+PP (#6258)

* 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

---------

Co-authored-by: Tong Li <tong.li35271158@gmail.com>
This commit is contained in:
YeAnbang
2025-04-10 10:52:18 +08:00
committed by GitHub
parent ed43a4be04
commit 9467c10690
3 changed files with 12 additions and 8 deletions

View File

@@ -121,7 +121,7 @@ if __name__ == "__main__":
# plugin_config={}, # for zero
plugin_config={
"pp_size": 2,
"tp_size": 1,
"tp_size": 2,
"microbatch_size": args.train_microbatch_size // 2,
"zero_stage": 0,
"max_norm": 1.0,