github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-31 08:34:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

clean code

Browse Source
This commit is contained in:
YeAnbang
2025-04-29 16:57:46 +08:00
parent 8c5dd4131d
commit 3381881908
3 changed files with 12 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
applications/ColossalChat/coati/distributed/grpo_consumer.py
View File

@@ -151,19 +151,6 @@ class GRPOConsumer(BaseConsumer):
eta_min=0.1 * grpo_config.get("lr", 1e-6),
)
def get_device_mesh_mapping(self):
return {
"rank": self.rank,
"tp_rank": self.tp_rank,
"tp_size": self.tp_size,
"dp_size": self.dp_size,
"dp_rank": self.dp_rank,
"pp_size": self.booster.plugin.stage_manager.num_stages,
"pp_stage": self.booster.plugin.stage_manager.stage,
"is_last_stage": self.booster.plugin.stage_manager.is_last_stage(),
"world_size": self.world_size,
}
def get_model_state_dict_keys(self):
return self.orig_state_dict_key

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

@@ -116,11 +116,10 @@ def launch_distributed(
consumer_procs.append(consumer)
# setup the consumer procs first
ray.get([p.setup.remote() for p in consumer_procs])
# get the device mesh mapping from consumer procs
consumer_device_mesh_mapping = ray.get([p.get_device_mesh_mapping.remote() for p in consumer_procs])
# get state dict key for checking syncing integrity
model_state_dict_keys = ray.get(consumer_procs[0].get_model_state_dict_keys.remote())
# setup the producer procs
ray.get([p.setup.remote(consumer_device_mesh_mapping, model_state_dict_keys) for p in producer_procs])
ray.get([p.setup.remote(model_state_dict_keys) for p in producer_procs])
# loop
procs = producer_procs + consumer_procs
ray.get([p.loop.remote() for p in procs])

30
applications/ColossalChat/coati/distributed/producer.py
View File

@@ -80,18 +80,12 @@ class BaseProducer:
else:
raise ValueError(f"Unexpected backend {backend}")
def setup(self, consumer_device_mesh_mapping: Dict[str, Any] = None, model_state_dict_keys: List = None) -> None:
self.consumer_device_mesh_mapping = consumer_device_mesh_mapping
def setup(self, model_state_dict_keys: List = None) -> None:
self.model_state_dict_keys = model_state_dict_keys
cc.init_collective_group(1 + self.num_consumer_procs, 0, group_name=f"sync_data_{self.producer_idx}")
# cc.init_collective_group(self.num_producers + 1, self.producer_idx, group_name="sync_model_pp_stage_0")
for i in range(self.num_consumer_procs):
device_mesh_mapping = self.consumer_device_mesh_mapping[i]
device_mesh_mapping["rank"]
# TODO: support ep, sp
if device_mesh_mapping["dp_rank"] == 0 and device_mesh_mapping["tp_rank"] == 0:
group_name = f"sync_model_pp_stage_{device_mesh_mapping['pp_stage']}"
cc.init_collective_group(self.num_producers + 1, self.producer_idx, group_name=group_name)
for pp_stage in range(self.consumer_plugin_config.get("pp_size", 1)):
group_name = f"sync_model_pp_stage_{pp_stage}"
cc.init_collective_group(self.num_producers + 1, self.producer_idx, group_name=group_name)
def rollout(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]:
raise NotImplementedError
@@ -141,17 +135,13 @@ class BaseProducer:
print(
f"[P{self.producer_idx}] Sync model episode {episode} step {(i + 1) // self.num_microbatches - 1}"
)
for consumer_rank_id in range(self.num_consumer_procs):
device_mesh_mapping = self.consumer_device_mesh_mapping[consumer_rank_id]
device_mesh_mapping["rank"]
# TODO: support ep, sp
if device_mesh_mapping["dp_rank"] == 0 and device_mesh_mapping["tp_rank"] == 0:
group_name = f"sync_model_pp_stage_{device_mesh_mapping['pp_stage']}"
state_dict.update(
ray_broadcast_tensor_dict(
None, src=self.num_producers, device=self.device, group_name=group_name
)
for pp_stage in range(self.consumer_plugin_config.get("pp_size", 1)):
group_name = f"sync_model_pp_stage_{pp_stage}"
state_dict.update(
ray_broadcast_tensor_dict(
None, src=self.num_producers, device=self.device, group_name=group_name
)
)
# check model sync integrity
assert len(state_dict) == len(
self.model_state_dict_keys