Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests

[BUG]: chatgpt ppo training hangs when using gemini
### 🐛 Describe the bug

## Problem

Running `ChatGPT/examples/train_prompts.py`, I found sometimes the training hangs when using Gemini.

This occurs when randomly while changing batch size.

## Possible reason

I found the [padding policy](https://huggingface.co/docs/transformers/main_classes/tokenizer#transformers.PreTrainedTokenizer.__call__.padding) is to pad to the longest sequence in the batch.
In DDP scheme, different process may have different input lengths due to random sampling. That is to say, they may have different generation steps.

When using Gemini, which need communication during forward, different forward steps leads to different number of communication calls. And this asymmetric communication leads to hang.

## Possible solution

Change padding policy to `'max_length'`, see [huggingface tokenizer doc](https://huggingface.co/docs/transformers/main_classes/tokenizer#transformers.PreTrainedTokenizer.__call__.padding) for more details.

In addition, when enabling early stopping, we should also consider DDP and ensure the number of generation steps of each i process is the same.



### Environment

_No response_


Discussion:

Resolving patch:
diff --git a/applications/ChatGPT/chatgpt/models/generation.py b/applications/ChatGPT/chatgpt/models/generation.py
--- a/applications/ChatGPT/chatgpt/models/generation.py
+++ b/applications/ChatGPT/chatgpt/models/generation.py
@@ -1,6 +1,7 @@
 from typing import Any, Callable, Optional
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 
 try:
@@ -27,6 +28,14 @@ def prepare_logits_processor(top_k: Optional[int] = None,
     return processor_list
 
 
+def _is_sequence_finished(unfinished_sequences: torch.Tensor) -> bool:
+    if dist.is_initialized() and dist.get_world_size() > 1:
+        # consider DP
+        unfinished_sequences = unfinished_sequences.clone()
+        dist.all_reduce(unfinished_sequences)
+    return unfinished_sequences.max() == 0
+
+
 def sample(model: nn.Module,
            input_ids: torch.Tensor,
            max_length: int,
@@ -74,7 +83,7 @@ def sample(model: nn.Module,
             unfinished_sequences = unfinished_sequences.mul((next_tokens != eos_token_id).long())
 
         # stop when each sentence is finished if early_stopping=True
-        if early_stopping and unfinished_sequences.max() == 0:
+        if early_stopping and _is_sequence_finished(unfinished_sequences):
             break
 
     return input_ids
diff --git a/applications/ChatGPT/examples/train_prompts.py b/applications/ChatGPT/examples/train_prompts.py
--- a/applications/ChatGPT/examples/train_prompts.py
+++ b/applications/ChatGPT/examples/train_prompts.py
@@ -46,7 +46,6 @@ def main(args):
         initial_model = deepcopy(actor)
         reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).to(torch.cuda.current_device())
 
-
     # configure optimizer
     if args.strategy.startswith('colossalai'):
         actor_optim = HybridAdam(actor.parameters(), lr=5e-6)
@@ -70,7 +69,9 @@ def main(args):
     dataset = pd.read_csv(args.prompt_path)['prompt']
 
     def tokenize_fn(texts):
-        batch = tokenizer(texts, return_tensors='pt', max_length=96, padding=True, truncation=True)
+        # MUST padding to max length to ensure inputs of all ranks have the same length
+        # Different length may lead to hang when using gemini, as different generation steps
+        batch = tokenizer(texts, return_tensors='pt', max_length=96, padding='max_length', truncation=True)
         return {k: v.cuda() for k, v in batch.items()}
 
     (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare(
@@ -101,7 +102,7 @@ def tokenize_fn(texts):
                 num_episodes=args.num_episodes,
                 max_timesteps=args.max_timesteps,
                 update_timesteps=args.update_timesteps)
-    # save model checkpoint after fitting 
+    # save model checkpoint after fitting
     strategy.save_model(actor, args.save_path, only_rank0=True)
     # save optimizer checkpoint on all ranks
     if args.need_optim_ckpt:


Test patch: