Code

🔧 Optimize GRPO VRAM Usage#2669

Code

Open

andyl98 wants to merge 45 commits intohuggingface:mainhuggingface/trl:mainfrom andyl98:grpo-vram-optimizationandyl98/trl:grpo-vram-optimizationCopy head branch name to clipboard

Commits

Commits on Jan 27, 2025

• Update utils.py

  andyl98

  committed
  851e04a

  Copy full SHA for 851e04a

• Update grpo_trainer.py

  andyl98

  committed
  427b512

  Copy full SHA for 427b512

• Update utils.py

  andyl98

  committed
  e3a3056

  Copy full SHA for e3a3056

• Update utils.py

  andyl98

  committed
  4cc06b8

  Copy full SHA for 4cc06b8

Commits on Jan 28, 2025

• Add torch.no_grad() for generate

  andyl98

  committed
  fede783

  Copy full SHA for fede783

• fixes

  andyl98

  committed
  ebcc49d

  Copy full SHA for ebcc49d

• Merge branch 'grpo-vram-optimization' of https://github.com/andyl98/trl into grpo-vram-optimization

  andyl98

  committed
  b5286a0

  Copy full SHA for b5286a0

• Delete test.ipynb

  andyl98

  committed
  e003c6a

  Copy full SHA for e003c6a

• improve comments

  andyl98

  committed
  0d298b5

  Copy full SHA for 0d298b5

• Merge branch 'main' into grpo-vram-optimization

  andyl98

  authored
  87421be

  Copy full SHA for 87421be

Commits on Jan 29, 2025

• Update utils.py

  andyl98

  committed
  b31c1a9

  Copy full SHA for b31c1a9

• Update utils.py

  andyl98

  committed
  69dbf38

  Copy full SHA for 69dbf38

• Merge branch 'main' into grpo-vram-optimization

  andyl98

  committed
  a99b6ee

  Copy full SHA for a99b6ee

• ruff

  andyl98

  committed
  31e02be

  Copy full SHA for 31e02be

• Ok need to revert no_grad on prompt tokens

  andyl98

  committed
  a6c07a5

  Copy full SHA for a6c07a5

• Oops I shouldn't unwrap here

  andyl98

  committed
  ae806f6

  Copy full SHA for ae806f6

Commits on Jan 31, 2025

• Add compatibility with gradient checkpointing

  andyl98

  committed
  b067cf3

  Copy full SHA for b067cf3

• Wait I made a big mistake in logits_to_keep for completion tokens

  andyl98

  committed
  2e4fd45

  Copy full SHA for 2e4fd45

• fix original attn mask

  andyl98

  committed
  1d4fe3c

  Copy full SHA for 1d4fe3c

Commits on Feb 1, 2025

• Seems logits_to_keep is not always trust-worthy

  andyl98

  committed
  78e9d9f

  Copy full SHA for 78e9d9f

• Same as above

  andyl98

  committed
  98773d9

  Copy full SHA for 98773d9

• Revert "Same as above"

  Show description for 7a88ed7
  andyl98

  committed
  7a88ed7

  Copy full SHA for 7a88ed7

• Update utils.py

  andyl98

  committed
  b903ebb

  Copy full SHA for b903ebb

Commits on Feb 4, 2025

• Process completion logits by mini batch to save extra memory

  andyl98

  committed
  2329709

  Copy full SHA for 2329709

• Merge remote-tracking branch 'origin' into grpo-vram-optimization

  andyl98

  committed
  86a7884

  Copy full SHA for 86a7884

• From what I've found, logits_to_keep is not very stable and I want to make sure it is respected

  andyl98

  committed
  d3e0f16

  Copy full SHA for d3e0f16

• Add back origin comment

  andyl98

  committed
  d4b15d1

  Copy full SHA for d4b15d1

• Update grpo_config.py

  andyl98

  committed
  6889bc3

  Copy full SHA for 6889bc3

• Update grpo_trainer.py

  andyl98

  committed
  03c938d

  Copy full SHA for 03c938d

• Allow mini batch processing on completion logits!

  andyl98

  committed
  adf4dc4

  Copy full SHA for adf4dc4

• break down

  andyl98

  committed
  dc75d6e

  Copy full SHA for dc75d6e

• log time

  andyl98

  committed
  4a05526

  Copy full SHA for 4a05526

• Update utils.py

  andyl98

  committed
  ccc9547

  Copy full SHA for ccc9547

Commits on Feb 5, 2025

• flag control

  andyl98

  committed
  14d4b6c

  Copy full SHA for 14d4b6c

• use num_logits_to_keep instead

  andyl98

  committed
  20667c7

  Copy full SHA for 20667c7

• let's use prompt_cache all the time

  andyl98

  committed
  c55cf10

  Copy full SHA for c55cf10

• Update grpo_trainer.py

  andyl98

  committed
  7aa950c

  Copy full SHA for 7aa950c

• Update utils.py

  andyl98

  committed
  ce54bb8

  Copy full SHA for ce54bb8

Commits on Feb 6, 2025

• don't delete these

  andyl98

  committed
  c45630a

  Copy full SHA for c45630a

• Use the awesome logexp trick

  andyl98

  committed
  da5e169

  Copy full SHA for da5e169

• Add some shape comments

  andyl98

  committed
  6e528ce

  Copy full SHA for 6e528ce

Commits on Feb 26, 2025

• Merge remote-tracking branch 'upstream/main' into grpo-vram-optimization

  andyl98

  committed
  32ea101

  Copy full SHA for 32ea101

• resolve conflicts

  andyl98

  committed
  f833384

  Copy full SHA for f833384

• Update grpo_config.py

  andyl98

  committed
  11b4c7e

  Copy full SHA for 11b4c7e

• Update grpo_trainer.py

  andyl98

  committed
  dd4f053

  Copy full SHA for dd4f053