Code

FSDP memory usage improvements#12

Code

Merged

epwalsh merged 6 commits intomainallenai/OLMo-core:mainfrom epwalsh/fsdp-profileallenai/OLMo-core:epwalsh/fsdp-profileCopy head branch name to clipboardApr 18, 2024

+40-25

Commits

Commits on Apr 18, 2024

• set to none

  epwalsh

  committed
  bc47475

  Copy full SHA for bc47475

• avoid allocating another sharded tensor

  epwalsh

  committed
  dbcd7f1

  Copy full SHA for dbcd7f1

• try allocating one fewer

  epwalsh

  committed
  96a34d4

  Copy full SHA for 96a34d4

• clone when needed

  epwalsh

  committed
  8c7f3a9

  Copy full SHA for 8c7f3a9

• post reduce-scatter

  epwalsh

  committed
  ddabdb2

  Copy full SHA for ddabdb2

• show memory usage when profiling

  epwalsh

  committed
  ed25dca

  Copy full SHA for ed25dca