Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

RLOO trainer: fix calculations of steps, episodes and epochs #2516

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

RLOO trainer: fix calculations of steps, episodes and epochs #2516

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c0b9ca3
Fix calculations of steps, episodes and epochs
dawidm Dec 23, 2024
d4c2dc9
fix: don't multiply steps by (num_mini_batches * num_ppo_epochs)
dawidm Jan 8, 2025
4eb4818
update documentation: logged episodes are not global steps
dawidm Jan 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/source/rloo_trainer.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ The logged metrics are as follows. Here is an example [tracked run at Weights an
* `val/ratio_var`: The variance of the `val/ratio`, indicating the variability in policy changes.
* `val/num_eos_tokens`: The number of end-of-sequence (EOS) tokens generated, which can indicate the number of complete responses.
* `lr`: lr: The current learning rate used by the optimizer.
* `episode`: episode: The current global step or episode count in the training process.
* `episode`: episode: The current episode count in the training process.


## Cookbook
Expand Down
8 changes: 4 additions & 4 deletions trl/trainer/rloo_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,12 +120,12 @@ def __init__(
# calculate various batch sizes
#########
if args.total_episodes is None: # allow the users to define episodes in terms of epochs.
args.total_episodes = int(args.num_train_epochs * self.train_dataset_len)
args.total_episodes = int(args.num_train_epochs * self.train_dataset_len * args.rloo_k)
accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps)
self.accelerator = accelerator
args.world_size = accelerator.num_processes
args.local_batch_size = (
args.per_device_train_batch_size * args.gradient_accumulation_steps * args.num_mini_batches
args.per_device_train_batch_size * args.gradient_accumulation_steps
)
args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size)
args.batch_size = int(args.local_batch_size * args.world_size)
Expand Down Expand Up @@ -275,8 +275,8 @@ def repeat_generator():
# trainer state initialization
self.state.global_step = 0
self.state.episode = 0
self.state.max_steps = (args.num_total_batches * args.num_mini_batches) // 2
self.state.num_train_epochs = args.total_episodes / self.train_dataset_len
self.state.max_steps = args.num_total_batches
self.state.num_train_epochs = (args.total_episodes / args.rloo_k) / self.train_dataset_len
# Compute absolute values for logging, eval, and save if given as ratio
if args.logging_steps is not None:
if args.logging_steps < 1:
Expand Down