Add more training models and RLHF algorithms #6368
Conversation
for more information, see https://pre-commit.ci
| # [minibatch_size x num_of_generation] | ||
| loss_mask = torch.ones(action_mask.size(0), device=action_mask.device).bool() | ||
| # [minibatch_size x num_of_generation] | ||
| loss_mask = torch.ones(action_mask.size(0), device=action_mask.device).bool() | ||
|
|
There was a problem hiding this comment.
may be better to move the common calculations outside of the if statements for conciseness
| # [minibatch_size x num_generations] | ||
| advantages = ((reward - reward_mean)).unsqueeze(dim=-1) | ||
|
|
||
| advantages_mean = advantages.mean(dim=0) |
There was a problem hiding this comment.
Isn't the advantages_mean always 0 as advantage is already zero-centered in the previous step?
| advantages_std = advantages.std(dim=0) | ||
|
|
||
| advantages = (advantages - advantages_mean) / (advantages_std + 1e-4) | ||
|
|
There was a problem hiding this comment.
maybe consider double-checking the reinforce++ baseline advantage calculation. In reinforce ++, each sample's advantage is calculated by subtracting the mean reward of all generation in the global batch, not per prompt mean
There was a problem hiding this comment.
For reinforce++, we should calculate norm adv using batch level mean and std.
| @@ -0,0 +1,2 @@ | |||
| 4.51.0: qwen2.5 + grpo, qwen3 + grpo, cannot: llama2, llama3.2 | |||
| 4.47.0: | |||
| @@ -227,13 +227,13 @@ | |||
| os.environ["TOKENIZERS_PARALLELISM"] = "false" # Disable tokenizers parallelism to avoid deadlock | |||
|
|
|||
| inference_model_config = dict(path=args.model) | |||
| train_model_config = dict(path=args.model, use_flash_attention_2=True, use_cache=False) | |||
| train_model_config = dict(path=args.model, use_flash_attention_2=False, use_cache=False) | |||
There was a problem hiding this comment.
why is flash attention not supported?
| generate_config = dict(top_k=args.top_k, top_p=args.top_p, temperature=args.temperature) | ||
|
|
||
| if args.backend == "transformers": | ||
| inference_model_config.update( | ||
| dict( | ||
| use_flash_attention_2=True, | ||
| use_flash_attention_2=False, |
There was a problem hiding this comment.
probably also consider force num_generation to 1 for reinforce++
| advantages_std = advantages.std(dim=0) | ||
|
|
||
| advantages = (advantages - advantages_mean) / (advantages_std + 1e-4) | ||
|
|
There was a problem hiding this comment.
For reinforce++, we should calculate norm adv using batch level mean and std.
| @@ -0,0 +1,2 @@ | |||
| 4.51.0: qwen2.5 + grpo, qwen3 + grpo, cannot: llama2, llama3.2 | |||
| 4.47.0: | |||
Closes #6367
📌 Checklist before creating the PR
[doc/gemini/tensor/...]: A concise descriptionpip install pre-commit && pre-commit install🚨 Issue number
📝 What does this PR do?
Add more training models (LLaMA3, Qwen3) and RLHF algorithms (REINFORCE++, RLOO).
💥 Checklist before requesting a review
⭐️ Do you enjoy contributing to Colossal-AI?
Tell us more if you don't enjoy contributing to Colossal-AI.