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

[Distributed RLHF] Integration of PP (#6257)

Browse Source 
* update help information

* update style

* fix

* minor fix

* support PP training

* add pp support

* remove unused code

* address conversation

---------

Co-authored-by: Tong Li <tong.li35271158@gmail.com>
This commit is contained in:
YeAnbang
2025-04-09 13:23:24 +08:00
committed by GitHub
parent 50153005b4
commit ed43a4be04
7 changed files with 263 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
applications/ColossalChat/coati/distributed/consumer.py
View File

@@ -54,7 +54,6 @@ class BaseConsumer:
self.model_config = model_config
self.plugin_config = plugin_config
assert self.plugin_config.get("pp_size", 1) == 1, "pp_size > 1 is not supported now"
self.device = get_current_device()
self.lr_scheduler = None
@@ -95,7 +94,6 @@ class BaseConsumer:
i = 0
for _ in range(self.num_recv_per_update):
# receive data from producers
for r in range(self.num_producers):
print(f"[T{dist.get_rank()}] Recv data episode {episode} step {step} from {r}")
self.buffer.extend(