github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-22 01:48:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

add simple grpo

Browse Source
This commit is contained in:
Tong Li
2025-02-23 22:54:26 +08:00
parent 8e6c9a4ab3
commit ffd3878a1e
8 changed files with 253 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
applications/ColossalChat/coati/distributed/launch.py
View File

@@ -2,7 +2,7 @@ from typing import Any, Dict, Optional
import ray
from .consumer import SimpleConsumer
from .grpo_consumer import GRPOConsumer
from .producer import SimpleProducer
@@ -68,7 +68,7 @@ def launch_distributed(
)
procs.append(producer)
for i in range(num_consumer_procs):
consumer = SimpleConsumer.options(num_gpus=1).remote(
consumer = GRPOConsumer.options(num_gpus=1).remote(
num_producers=num_producers,
num_episodes=num_episodes,
rank=i,