Merge branch 'grpo-latest-rebase-main' of https://github.com/hpcaitech/ColossalAI into grpo-latest-rebase-main

Author: YeAnbang

.github/workflows/run_chatgpt_examples.yml

@@ -21,7 +21,7 @@ jobs:
     container:
       image: image-cloud.luchentech.com/hpcaitech/pytorch-cuda:2.2.2-12.1.0
       options: --gpus all --rm -v /data/scratch/examples-data:/data/scratch/examples-data --shm-size=10.24gb
-    timeout-minutes: 60
+    timeout-minutes: 180
     defaults:
       run:
         shell: bash
@@ -34,6 +34,10 @@ jobs:
           pip install --no-cache-dir -v -e .
 
       - name: Install ChatGPT
+        env:
+          CFLAGS: "-O1"
+          CXXFLAGS: "-O1"
+          MAX_JOBS: 4
         run: |
           pip install flash-attn --no-build-isolation
           cd applications/ColossalChat

.github/workflows/run_chatgpt_unit_tests.yml

@@ -21,7 +21,7 @@ jobs:
     container:
       image: image-cloud.luchentech.com/hpcaitech/pytorch-cuda:2.2.2-12.1.0
       options: --gpus all --rm -v /data/scratch/examples-data:/data/scratch/examples-data
-    timeout-minutes: 30
+    timeout-minutes: 180
     defaults:
       run:
         shell: bash
@@ -30,6 +30,10 @@ jobs:
         uses: actions/checkout@v2
 
       - name: Install ChatGPT
+        env:
+          CFLAGS: "-O1"
+          CXXFLAGS: "-O1"
+          MAX_JOBS: 4
         run: |
           pip install flash-attn --no-build-isolation
           cd applications/ColossalChat

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -530,4 +530,4 @@ def state_dict(self):
         model = self.policy_model.unwrap()
         state_dict = model.state_dict()
         state_dict["consumer_global_step"] = torch.tensor([self.global_step], device=self.device)
-        return state_dict
+        return state_dict

colossalai/shardformer/modeling/qwen3.py

@@ -273,7 +273,7 @@ def qwen3_for_causal_lm_forward(
         hidden_states: Optional[torch.FloatTensor] = None,
         stage_index: Optional[List[int]] = None,
         shard_config: ShardConfig = None,
-        **kwargs
+        **kwargs,
     ):
         r"""
         Args: