Code

R1-Style distributed GRPO#2326

Code

Merged

ebsmothers merged 117 commits intopytorch:mainpytorch/torchtune:mainfrom RedTachyon:mainRedTachyon/torchtune:mainCopy head branch name to clipboardFeb 21, 2025

+2,425-2

Commits

Commits on Jan 25, 2025

• RL starter code

  RedTachyon

  committed
  1c29d67

  Copy full SHA for 1c29d67

Commits on Jan 27, 2025

• Add gsm8k

  RedTachyon

  committed
  ed544da

  Copy full SHA for ed544da

• Notebooks

  RedTachyon

  committed
  6ca2c38

  Copy full SHA for 6ca2c38

• Distributed dev progress

  RedTachyon

  committed
  7e39b69

  Copy full SHA for 7e39b69

Commits on Jan 29, 2025

• Decent progress on R1 RL

  RedTachyon

  committed
  55b1c65

  Copy full SHA for 55b1c65

Commits on Jan 30, 2025

• PoC training loop, but not working - code checkpoint

  RedTachyon

  committed
  aa34954

  Copy full SHA for aa34954

Commits on Jan 31, 2025

• 8B recipe, sorta running?

  RedTachyon

  committed
  0cc7795

  Copy full SHA for 0cc7795

• Merge branch 'main' into grpo

  RedTachyon

  authored
  b11e742

  Copy full SHA for b11e742

• Merge pull request #1 from RedTachyon/grpo

  Show description for ce7e8ce
  RedTachyon

  authored
  ce7e8ce

  Copy full SHA for ce7e8ce

Commits on Feb 1, 2025

• Some updates, some progress

  RedTachyon

  committed
  736b14f

  Copy full SHA for 736b14f

• Multi-node training, new reward shaping, success metric

  RedTachyon

  committed
  eb3d5b9

  Copy full SHA for eb3d5b9

Commits on Feb 2, 2025

• Synchronize metrics, fix parsing, more memory management

  RedTachyon

  committed
  9df4b38

  Copy full SHA for 9df4b38

Commits on Feb 3, 2025

• Sync rewards and successes between processes

  Show description for 3d39ffc
  RedTachyon

  committed
  3d39ffc

  Copy full SHA for 3d39ffc

• Add filter kwargs and partition for easier dataset filtering

  Show description for 30fa65c
  RedTachyon

  committed
  30fa65c

  Copy full SHA for 30fa65c

• Merge pull request #2 from RedTachyon/grpo

  Show description for e8f19f2
  RedTachyon

  authored
  e8f19f2

  Copy full SHA for e8f19f2

Commits on Feb 4, 2025

• Reorganize methods

  Show description for 499c013
  RedTachyon

  committed
  499c013

  Copy full SHA for 499c013

Commits on Feb 5, 2025

• Batched logit to logprob conversion

  Show description for 3ca4f49
  RedTachyon

  committed
  3ca4f49

  Copy full SHA for 3ca4f49

Commits on Feb 6, 2025

• Remove old notes

  RedTachyon

  committed
  300c117

  Copy full SHA for 300c117

• Revert config changes

  RedTachyon

  committed
  3cf169c

  Copy full SHA for 3cf169c

• More config cleanup

  RedTachyon

  committed
  7d9d37a

  Copy full SHA for 7d9d37a

• Recipe cleanup

  RedTachyon

  committed
  654eb56

  Copy full SHA for 654eb56

• Remove unnecessary PPO change

  RedTachyon

  committed
  c75cb7a

  Copy full SHA for c75cb7a

• Cleanup

  RedTachyon

  committed
  1f6be85

  Copy full SHA for 1f6be85

• Merge branch 'pytorch:main' into main

  RedTachyon

  authored
  cf770d3

  Copy full SHA for cf770d3

• Remove redundant code

  RedTachyon

  committed
  1fd86c8

  Copy full SHA for 1fd86c8

• More redundant code

  RedTachyon

  committed
  e198825

  Copy full SHA for e198825

• Reorganize recipes

  RedTachyon

  committed
  474c8dc

  Copy full SHA for 474c8dc

Commits on Feb 7, 2025

• Some cleanup of RL dataset and GSM8k

  RedTachyon

  committed
  60d7cd9

  Copy full SHA for 60d7cd9

• Pre-commit cleanup

  RedTachyon

  committed
  8d018ba

  Copy full SHA for 8d018ba

• Remove MATH dataset for now

  RedTachyon

  committed
  c9a01cb

  Copy full SHA for c9a01cb

• Properly remove MATH dataset

  RedTachyon

  committed
  c17bfc0

  Copy full SHA for c17bfc0

• Docstrings

  RedTachyon

  committed
  b312af8

  Copy full SHA for b312af8

• Reorganize some recipes, add sbatch for SFT

  RedTachyon

  committed
  f7c0929

  Copy full SHA for f7c0929

• Remove unused 8B configs, add another sbatch

  RedTachyon

  committed
  f6c7e53

  Copy full SHA for f6c7e53

• Recipes leanup

  RedTachyon

  committed
  72cb4bd

  Copy full SHA for 72cb4bd

• GRPO recipe cleanup

  RedTachyon

  committed
  a9112b8

  Copy full SHA for a9112b8

• Final MVP bugfixes

  RedTachyon

  committed
  b089ff2

  Copy full SHA for b089ff2

• Remove old unused test

  RedTachyon

  committed
  4870213

  Copy full SHA for 4870213

• Pre-commit

  RedTachyon

  committed
  2dd6f60

  Copy full SHA for 2dd6f60

• Stop token handling for both single and multi device

  RedTachyon

  committed
  7b58712

  Copy full SHA for 7b58712

Commits on Feb 8, 2025

• Update recipes/configs/dev/grpo/3B_full_rl.yaml

  Show description for b010ef6
  

  RedTachyon

  and

  felipemello1

  authored
  b010ef6

  Copy full SHA for b010ef6

• Remove redundant comment

  RedTachyon

  committed
  603d16a

  Copy full SHA for 603d16a

• Delete recipes/configs/llama3_2/3B_full_rl_single_device_mps.yaml

  RedTachyon

  authored
  d7269e4

  Copy full SHA for d7269e4

• Fix function arguments in optimizer setup

  RedTachyon

  committed
  9978a44

  Copy full SHA for 9978a44

• Update recipes/configs/dev/grpo/3B_full_rl.yaml

  Show description for d710200
  

  RedTachyon

  and

  ebsmothers

  authored
  d710200

  Copy full SHA for d710200

• Update torchtune/rlhf/loss/grpo.py

  Show description for 352c4fb
  

  RedTachyon

  and

  ebsmothers

  authored
  352c4fb

  Copy full SHA for 352c4fb

• PPO -> GRPO

  RedTachyon

  committed
  8f60178

  Copy full SHA for 8f60178

• RL -> GRPO (recipe

  RedTachyon

  committed
  bb8b97d

  Copy full SHA for bb8b97d

• Make sure we're not training with float16

  RedTachyon

  committed
  44aeb85

  Copy full SHA for 44aeb85

• Remove chunked loss logic

  RedTachyon

  committed
  86de099

  Copy full SHA for 86de099

• Save rank and world_size in the recipe

  RedTachyon

  committed
  120b1d2

  Copy full SHA for 120b1d2

• Rename R1Trajectory to GRPOTrajectory, other cleanup

  RedTachyon

  committed
  f3e9b77

  Copy full SHA for f3e9b77

• | None -> Optional

  RedTachyon

  committed
  b856695

  Copy full SHA for b856695

• Fix return_logits==False edge case

  RedTachyon

  committed
  e61894b

  Copy full SHA for e61894b

• Undo an accidental change in generation

  RedTachyon

  committed
  96dc9c9

  Copy full SHA for 96dc9c9

• Update generate_trajectory docstring

  RedTachyon

  committed
  b1cfc3b

  Copy full SHA for b1cfc3b

Commits on Feb 10, 2025

• Reenable reference network

  Show description for 8a62d1e
  RedTachyon

  committed
  8a62d1e

  Copy full SHA for 8a62d1e

• Remove optimizer_in_bwd

  RedTachyon

  committed
  7394cc7

  Copy full SHA for 7394cc7

• Remove activation offloading

  RedTachyon

  committed
  24cd238

  Copy full SHA for 24cd238

• (docstring) Reward model -> reward function

  RedTachyon

  committed
  67458f2

  Copy full SHA for 67458f2

• Move reward function to a new file

  RedTachyon

  committed
  d9c37b8

  Copy full SHA for d9c37b8

• Remove dead code

  RedTachyon

  committed
  18baeae

  Copy full SHA for 18baeae

• Remove redundant logging

  RedTachyon

  committed
  646a2ce

  Copy full SHA for 646a2ce

• Remove question from the reward function

  RedTachyon

  committed
  f641161

  Copy full SHA for f641161

• Update recipes/dev/grpo_full_finetune_distributed.py

  Show description for 9f4fb88
  

  RedTachyon

  and

  SalmanMohammadi

  authored
  9f4fb88

  Copy full SHA for 9f4fb88

• _grpo_step -> grpo_step

  RedTachyon

  committed
  123883a

  Copy full SHA for 123883a

• Remove breakpoints and comments

  RedTachyon

  committed
  d9cf499

  Copy full SHA for d9cf499

• Remove breakpoints and comments

  RedTachyon

  committed
  98df9a7

  Copy full SHA for 98df9a7

• Docstring for the reward function

  RedTachyon

  committed
  b068132

  Copy full SHA for b068132

• Handle reference checkpoint separately

  Show description for 738dc2d
  RedTachyon

  committed
  738dc2d

  Copy full SHA for 738dc2d

• Remove mentions of activation offloading

  RedTachyon

  committed
  759ff8b

  Copy full SHA for 759ff8b

• Fix messed up loss

  RedTachyon

  committed
  239d382

  Copy full SHA for 239d382

• Fix messed up loss, barriers to keep things in sync

  RedTachyon

  committed
  957671c

  Copy full SHA for 957671c

Commits on Feb 11, 2025

• Delete max_steps_per_epoch and gradient accumulation, simplify inner training loop

  RedTachyon

  committed
  357aa43

  Copy full SHA for 357aa43

• Pre-commit

  RedTachyon

  committed
  3c54d01

  Copy full SHA for 3c54d01

• Reorganize recipes

  RedTachyon

  committed
  0d77799

  Copy full SHA for 0d77799

• Remove dead settings from the GRPO config

  RedTachyon

  committed
  0f8df07

  Copy full SHA for 0f8df07

• Use DiskLogger in GRPO

  RedTachyon

  committed
  184aeef

  Copy full SHA for 184aeef

• Use DiskLogger in SFT for GRPO

  RedTachyon

  committed
  3018eeb

  Copy full SHA for 3018eeb

• Recipe name in logging

  RedTachyon

  committed
  d29c455

  Copy full SHA for d29c455

• Remove redundant logging

  RedTachyon

  committed
  b9a56b9

  Copy full SHA for b9a56b9

Commits on Feb 12, 2025

• Cleaned up official configs

  RedTachyon

  committed
  b143b00

  Copy full SHA for b143b00

• Merge branch 'main' into main

  RedTachyon

  authored
  05cf10d

  Copy full SHA for 05cf10d

• Docstrings for GRPO types

  RedTachyon

  committed
  be4d96d

  Copy full SHA for be4d96d

• Merge remote-tracking branch 'arielpublic/main'

  RedTachyon

  committed
  1303695

  Copy full SHA for 1303695

• Fix checkpointing

  RedTachyon

  authored
  9d5bfa5

  Copy full SHA for 9d5bfa5

• Update recipes/dev/grpo_full_finetune_distributed.py

  Show description for 6cbb10a
  

  RedTachyon

  and

  SalmanMohammadi

  authored
  6cbb10a

  Copy full SHA for 6cbb10a

• Update recipes/dev/grpo_full_finetune_distributed.py

  Show description for 97686e2
  

  RedTachyon

  and

  SalmanMohammadi

  authored
  97686e2

  Copy full SHA for 97686e2

• Update recipes/dev/grpo_full_finetune_distributed.py

  Show description for aeb69cf
  

  RedTachyon

  and

  SalmanMohammadi

  authored
  aeb69cf

  Copy full SHA for aeb69cf

• Remove mention of ac

  Show description for feeb042
  RedTachyon

  committed
  feeb042

  Copy full SHA for feeb042

• Pre-commit

  RedTachyon

  committed
  e7cb937

  Copy full SHA for e7cb937

• Revert generation changes, pre-commit

  RedTachyon

  committed
  2a75224

  Copy full SHA for 2a75224

• Fix RL collate function

  RedTachyon

  committed
  5010f1d

  Copy full SHA for 5010f1d

• Remove resharding from config

  RedTachyon

  committed
  711ee6b

  Copy full SHA for 711ee6b

Commits on Feb 13, 2025

• | -> Optional in rl dataset

  RedTachyon

  committed
  aad81fb

  Copy full SHA for aad81fb

• | -> Optional in GSM8k dataset

  RedTachyon

  committed
  53cd129

  Copy full SHA for 53cd129

Commits on Feb 16, 2025

• Merge branch 'pytorch:main' into main

  RedTachyon

  authored
  4dbe231

  Copy full SHA for 4dbe231

Commits on Feb 17, 2025

• Experimental stuff

  RedTachyon

  committed
  3a8c6de

  Copy full SHA for 3a8c6de

• Additional experimental stuff

  RedTachyon

  committed
  9677df1

  Copy full SHA for 9677df1

• Move experimental code to /dev

  RedTachyon

  committed
  9340f89

  Copy full SHA for 9340f89

• Properly move experimental code to /dev

  RedTachyon

  committed
  53ee98f

  Copy full SHA for 53ee98f

• Remove recursive_reshard from the public API

  RedTachyon

  committed
  682ebef

  Copy full SHA for 682ebef

• Separate optimized generation function, small fixes

  RedTachyon

  committed
  8dd7546

  Copy full SHA for 8dd7546

• Undo generation changes in the main function

  RedTachyon

  committed
  0939c94

  Copy full SHA for 0939c94

• Fix custom generate

  RedTachyon

  committed
  81a7765

  Copy full SHA for 81a7765

• Pre-commit

  RedTachyon

  committed
  e41a520

  Copy full SHA for e41a520

Commits on Feb 18, 2025

• Fix SFT dataset path

  RedTachyon

  committed
  45084e0

  Copy full SHA for 45084e0

Commits on Feb 19, 2025

• Merge branch 'pytorch:main' into main

  RedTachyon

  authored
  8283230

  Copy full SHA for 8283230

• Merge branch 'experiments' into main

  RedTachyon

  authored
  987e971

  Copy full SHA for 987e971

• Revert "Merge branch 'experiments' into main"

  Show description for 7a03139
  RedTachyon

  committed
  7a03139

  Copy full SHA for 7a03139

Commits on Feb 21, 2025

• Update recipes/configs/dev/3B_full_grpo.yaml

  Show description for 3fc43d7
  

  RedTachyon

  and

  ebsmothers

  authored
  3fc43d7

  Copy full SHA for 3fc43d7

• Update recipes/configs/dev/3B_sft_for_grpo.yaml

  Show description for dde5fd8
  

  RedTachyon

  and

  ebsmothers

  authored
  dde5fd8

  Copy full SHA for dde5fd8

• Update recipes/configs/dev/3B_full_grpo.yaml

  Show description for 2fddf9c
  

  RedTachyon

  and

  ebsmothers

  authored
  2fddf9c

  Copy full SHA for 2fddf9c

• Remove redundant async checkpointing code

  RedTachyon

  committed
  aead54e

  Copy full SHA for aead54e

• Remove some redundant clones

  RedTachyon

  committed
  16ad525

  Copy full SHA for 16ad525

• Add a generation comment

  RedTachyon

  committed
  70886b2

  Copy full SHA for 70886b2

• Pre-commit

  RedTachyon

  committed
  2ba4a97

  Copy full SHA for 2ba4a97