Merge branch 'grpo-latest-model-sync-5-16' into grpo-latest-model-sync-5-20

Author: duanjunwen

applications/ColossalChat/coati/distributed/consumer.py

@@ -18,7 +18,6 @@
 from .comm import ray_broadcast_tensor_dict
 from .utils import bind_batch, post_recv, unbind_batch
 
-
 class BaseConsumer:
     def __init__(
         self,
@@ -33,6 +32,7 @@ def __init__(
         batch_size: int,
         model_config: Dict[str, Any],
         plugin_config: Dict[str, Any],
+        generate_config: Dict[str, Any],
         minibatch_size: int = 1,
         save_interval: int = 100,
         save_dir: str = "./model",
@@ -55,8 +55,11 @@ def __init__(
         self.model_config = model_config
         self.plugin_config = plugin_config
 
-        self.device = get_current_device()
+        # self.device = get_current_device()
+        self.device = 'npu'
+        # self.device = torch.device(f"npu:{torch.npu.current_device()}")
         self.lr_scheduler = None
+        self.generate_config = generate_config
 
     def setup(self) -> None:
         launch(self.rank, self.world_size, self.master_addr, self.master_port, local_rank=0)
@@ -73,24 +76,26 @@ def setup(self) -> None:
         self.booster = Booster(plugin=self.plugin)
         self.dp_rank = dist.get_rank(self.plugin.dp_group)
         self.tp_rank = dist.get_rank(self.plugin.tp_group)
+        self.sp_rank = dist.get_rank(self.plugin.sp_group)
         self.pp_rank = dist.get_rank(self.plugin.pp_group)
 
         self.dp_size = dist.get_world_size(self.plugin.dp_group)
         self.tp_size = dist.get_world_size(self.plugin.tp_group)
+        self.sp_size = dist.get_world_size(self.plugin.sp_group)
         self.pp_size = dist.get_world_size(self.plugin.pp_group)
 
         # Init Hybrid ray process group
         for i in range(self.num_producers):
-            cc.init_collective_group(self.world_size + 1, self.rank + 1, group_name=f"sync_data_{i}")
+            cc.init_collective_group(self.world_size + 1, self.rank + 1, backend='hccl',group_name=f"sync_data_{i}")
         if self.pp_size > 1:
             # use hybrid tp + pp
             if self.tp_rank == 0 and self.dp_rank == 0:
                 cc.init_collective_group(
-                    self.num_producers + 1, self.num_producers, group_name=f"sync_model_{self.pp_rank}"
+                    self.num_producers + 1, self.num_producers, backend='hccl', group_name=f"sync_model_{self.pp_rank}"
                 )
         else:
             if self.rank == 0:
-                cc.init_collective_group(self.num_producers + 1, self.num_producers, group_name="sync_model")
+                cc.init_collective_group(self.num_producers + 1, self.num_producers, backend='hccl', group_name="sync_model")
 
         self.buffer = []
         self.recv_cnt = 0
@@ -156,7 +161,7 @@ def loop(self) -> None:
                                 f"[T{dist.get_rank()}] Sync model PP stage {self.pp_rank} episode {episode} step {step}"
                             )
                         else:
-                            print(f"[T{dist.get_rank()}] Sync model episode {episode} step {step}")
+                            print(f"[T{dist.get_rank()}] Sync model episode {episode} step {step}")  
                         torch.cuda.empty_cache()
                         state_dict = self.state_dict()
                         if self.pp_size > 1:

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -341,6 +341,7 @@ def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]:
                                 num_action,
                                 self.plugin.shard_config,
                             )
+                            del reference_model_logits
                         else:
                             # Dummy reference logprobs for data iterator.
                             reference_action_log_probs = None
@@ -420,6 +421,7 @@ def _criterion(outputs, inputs):
                         num_action,
                         self.plugin.shard_config,
                     )
+                    del policy_model_logits
 
                     if self.policy_loss_fn.beta > 0:
                         with torch.no_grad():
@@ -433,6 +435,7 @@ def _criterion(outputs, inputs):
                             num_action,
                             self.plugin.shard_config,
                         )
+                        del reference_model_logits
                         per_token_kl = (
                             torch.exp(reference_action_log_probs - action_log_probs)
                             - (reference_action_log_probs - action_log_probs)

applications/ColossalChat/coati/distributed/inference_backend.py

@@ -210,6 +210,7 @@ def __init__(
         self.model_config = model_config
         self.tokenizer = tokenizer
         self.num_generations = num_generations
+        self.max_length = generate_config['max_tokens']
 
     @torch.no_grad()
     def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]:

applications/ColossalChat/coati/distributed/launch.py

@@ -1,4 +1,5 @@
 import copy
+import os
 import uuid
 from typing import Any, Dict, Optional
 
@@ -64,24 +65,119 @@ def launch_distributed(
         core_consumer = ALGO_MAP.get(core_algo, SimpleConsumer)
 
     train_dp_size = get_dp_size_fast(num_consumer_procs, plugin_config)
+    print(f"inference_batch_size {inference_batch_size} num_producers {num_producers} train_batch_size {train_batch_size} train_dp_size {train_dp_size}")
     assert (inference_batch_size * num_producers) % (train_batch_size * train_dp_size) == 0
 
     dataset_path = train_dataset_config["path"]
     num_samples = get_jsonl_size_fast(dataset_path)
     global_inference_batch_size = inference_batch_size * num_producers
     num_update_per_episode = num_samples // global_inference_batch_size
     num_recv_per_update = inference_batch_size // inference_microbatch_size
-
+    
     run_name = f"{inference_backend}_bs_{train_batch_size * train_dp_size}_temp_{generate_config['temperature']:.01f}_top_p_{generate_config['top_p']:.02f}"
     wandb_group_name = str(uuid.uuid4())
     rollout_log_file = os.path.join(
         rollout_save_dir,
         f"{project_name.replace(' ','_')}_run_{wandb_group_name}.jsonl",
     )
 
-    procs = []
+
+    # ###########################################
+    # # Old version, may lead colossalai init stuck in multinodes
+    # ############################################
+    # procs = []
+    # for i in range(num_producers):
+    #     # producer = SimpleProducer.options(num_gpus=num_proc_per_producer).remote(
+    #     producer = SimpleProducer.options(num_cpus=1, resources={"NPU":num_proc_per_producer}).remote(
+    #         producer_idx=i,
+    #         num_producers=num_producers,
+    #         num_consumer_procs=num_consumer_procs,
+    #         num_episodes=num_episodes,
+    #         batch_size=inference_batch_size,
+    #         dataset_config=dataset_config,
+    #         dataloaders_config=dataloaders_config,
+    #         model_config=inference_model_config,
+    #         generate_config=generate_config,
+    #         tokenizer_config=tokenizer_config,
+    #         microbatch_size=inference_microbatch_size,
+    #         backend=inference_backend,
+    #         num_generations=num_generations,
+    #         consumer_plugin_config=plugin_config,
+    #     )
+    #     procs.append(producer)
+    # generate_config_consumer = copy.deepcopy(generate_config)
+    # generate_config_consumer.update(
+    #     dict(
+    #         backend=inference_backend,
+    #     )
+    # )
+    # for i in range(num_consumer_procs):
+    #     # consumer = core_consumer.options(num_gpus=1).remote(
+    #     consumer = core_consumer.options(num_cpus=1, resources={"NPU":1}).remote(
+    #         num_producers=num_producers,
+    #         num_episodes=num_episodes,
+    #         rank=i,
+    #         world_size=num_consumer_procs,
+    #         master_addr=master_addr,
+    #         master_port=master_port,
+    #         num_update_per_episode=num_update_per_episode,
+    #         num_recv_per_update=num_recv_per_update,
+    #         batch_size=train_batch_size,
+    #         model_config=train_model_config,
+    #         plugin_config=plugin_config,
+    #         minibatch_size=train_minibatch_size,
+    #         generate_config=generate_config_consumer,
+    #         grpo_config=grpo_config,
+    #         num_generations=num_generations,
+    #         project_name=project_name,
+    #         save_interval=save_interval,
+    #         save_dir=save_dir,
+    #     )
+    #     procs.append(consumer)
+    # ray.get([p.setup.remote() for p in procs])
+    # ray.get([p.loop.remote() for p in procs])
+    
+    ###########################################
+    # New version, assign master ip for colossalai & vllm respectively
+    ###########################################
+    nodes = ray.nodes()
+    node_info = {
+        node["NodeID"]: {
+            # "num_gpus": node["Resources"].get("GPU", 0),
+            "num_gpus": node["Resources"].get("NPU", 0),
+            "address": node["NodeManagerAddress"],
+        }  # Default to 0 if no GPUs are available
+        for node in nodes
+    }
+    print(f"node_info {node_info}")
+    gpu_to_node_id = []
+    gpu_to_ip_address = []
+    for node_id in node_info:
+        for idx in range(int(node_info[node_id]["num_gpus"])): # use num_gpus instead of num_npus
+            gpu_to_node_id.append(node_id)
+            gpu_to_ip_address.append(node_info[node_id]["address"])
+    print(f"node_info {node_info} \n gpu_to_node_id {gpu_to_node_id} \n gpu_to_ip_address {gpu_to_ip_address} \n")
+
+    producer_procs = []
+    
     for i in range(num_producers):
-        producer = SimpleProducer.options(num_gpus=num_proc_per_producer).remote(
+        node_id = gpu_to_node_id[0]
+        producer_ip_address = gpu_to_ip_address[0]
+        for _ in range(num_proc_per_producer):
+            gpu_to_node_id.pop(0)
+            gpu_to_ip_address.pop(0)
+        print(f"Schedual Producer P[{i}] which requires {num_proc_per_producer} GPUs on node {producer_ip_address}")
+        
+        producer = SimpleProducer.options(
+            # num_cpus=1,
+            # num_cpus=num_proc_per_producer, 
+            num_gpus=0,
+            resources={"NPU":num_proc_per_producer},
+            scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+                node_id=node_id,
+                soft=False,
+            ),
+        ).remote(
             producer_idx=i,
             num_producers=num_producers,
             num_consumer_procs=num_consumer_procs,
@@ -107,20 +203,36 @@ def launch_distributed(
             log_rollout_interval=log_rollout_interval,
             rollout_log_file=rollout_log_file,
         )
-        procs.append(producer)
+        producer_procs.append(producer)
+    ray.get([p.setup.remote() for p in producer_procs])
     generate_config_consumer = copy.deepcopy(generate_config)
     generate_config_consumer.update(
         dict(
             backend=inference_backend,
         )
     )
+    consumer_master_ip_address = gpu_to_ip_address[0]
+    print(f"Use {consumer_master_ip_address} as master address for torch DDP.")
+    consumer_procs = []
     for i in range(num_consumer_procs):
-        consumer = core_consumer.options(num_gpus=1).remote(
+        node_id = gpu_to_node_id[0]
+        consumer_ip_address = gpu_to_ip_address[0]
+        gpu_to_node_id.pop(0)
+        gpu_to_ip_address.pop(0)
+        print(f"Schedual Consumer T[{i}] which requires 1 GPUs on node {consumer_ip_address}")
+        consumer = core_consumer.options(
+            resources={"NPU":1},
+            scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+                node_id=node_id,
+                soft=False,
+            ),
+        ).remote(
             num_producers=num_producers,
             num_episodes=num_episodes,
             rank=i,
             world_size=num_consumer_procs,
-            master_addr=master_addr,
+            # master_addr=master_addr,
+            master_addr=consumer_master_ip_address,
             master_port=master_port,
             num_update_per_episode=num_update_per_episode,
             num_recv_per_update=num_recv_per_update,
@@ -137,6 +249,6 @@ def launch_distributed(
             run_name=run_name,
             wandb_group_name=wandb_group_name,
         )
-        procs.append(consumer)
-    ray.get([p.setup.remote() for p in procs])
-    ray.get([p.loop.remote() for p in procs])
+        consumer_procs.append(consumer)
+    ray.get([p.setup.remote() for p in consumer_procs])
+    ray.get([p.loop.remote() for p in (producer_procs + consumer_procs)])

applications/ColossalChat/coati/distributed/producer.py

@@ -151,6 +151,9 @@ def __init__(
         else:
             raise ValueError("eval_dataset_config is not defined")
         self.device = get_current_device()
+        # self.device = get_current_device()
+        self.device = 'npu'
+        # self.device = torch.device(f"npu:{torch.npu.current_device()}")
 
         # init backend
         if backend in BACKEND_MAP:
@@ -161,18 +164,12 @@ def __init__(
         self.consumer_pp_size = consumer_plugin_config.get("pp_size", 1)  # consumer pp size
 
     def setup(self) -> None:
-        cc.init_collective_group(
-            world_size=self.num_producers,
-            rank=self.producer_idx,
-            backend=Backend.NCCL,
-            group_name="producer_group",
-        )
-        cc.init_collective_group(1 + self.num_consumer_procs, 0, group_name=f"sync_data_{self.producer_idx}")
+        cc.init_collective_group(1 + self.num_consumer_procs, 0, backend='hccl', group_name=f"sync_data_{self.producer_idx}")
         if self.consumer_pp_size > 1:
             for i in range(self.consumer_pp_size):
-                cc.init_collective_group(self.num_producers + 1, self.producer_idx, group_name=f"sync_model_{i}")
+                cc.init_collective_group(self.num_producers + 1, self.producer_idx, backend='hccl', group_name=f"sync_model_{i}")
         else:
-            cc.init_collective_group(self.num_producers + 1, self.producer_idx, group_name="sync_model")
+            cc.init_collective_group(self.num_producers + 1, self.producer_idx, backend='hccl', group_name="sync_model")
 
     def rollout(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]:
         raise NotImplementedError
@@ -250,7 +247,7 @@ def loop(self) -> None:
                 outputs["temperature"] = torch.tensor(
                     [self.model.generate_config["temperature"]] * outputs["input_ids"].size(0)
                 ).to(outputs["input_ids"].device)
-                outputs = pre_send(outputs)
+                # outputs = pre_send(outputs)
                 ray_broadcast_tensor_dict(
                     outputs, src=0, device=self.device, group_name=f"sync_data_{self.producer_idx}"
                 )

applications/ColossalChat/coati/distributed/utils.py

@@ -3,6 +3,7 @@
 from typing import Any, Dict, List
 
 import torch
+import math
 from filelock import FileLock
 
 from colossalai.shardformer.layer.loss import dist_log_prob

applications/ColossalChat/fusion_result.json

@@ -0,0 +1 @@
+null