Code

Enable FSDP2 cpu offloading#624

Code

Merged

mori360 merged 17 commits intopytorch:mainpytorch/torchtitan:mainfrom mori360:cpu_offloadmori360/torchtitan:cpu_offloadCopy head branch name to clipboardOct 28, 2024

+52-10

Commits

Commits on Oct 17, 2024

• enable FSDP2 cpuoffload

  mori360

  committed
  67407e6

  Copy full SHA for 67407e6

• rename config

  mori360

  committed
  9a8c90d

  Copy full SHA for 9a8c90d

• lint fix

  mori360

  committed
  530683d

  Copy full SHA for 530683d

Commits on Oct 21, 2024

• update cpu offload config

  mori360

  committed
  01b0662

  Copy full SHA for 01b0662

Commits on Oct 23, 2024

• manage freqs_cis as nn.Parameter

  mori360

  committed
  9569e57

  Copy full SHA for 9569e57

• init freqs_cis at meta device

  mori360

  committed
  30c204f

  Copy full SHA for 30c204f

• remove memory snapshot

  mori360

  committed
  7d36b87

  Copy full SHA for 7d36b87

• Merge branch 'main' into cpu_offload

  mori360

  authored
  82f4fec

  Copy full SHA for 82f4fec

Commits on Oct 24, 2024

• change default device, not set nn.Parameters

  mori360

  committed
  c691a6d

  Copy full SHA for c691a6d

• typo

  mori360

  committed
  3dfca1b

  Copy full SHA for 3dfca1b

• use a context manager for model.init_weights

  mori360

  committed
  6fef39a

  Copy full SHA for 6fef39a

• correct condition for non_cpu_offloading case

  mori360

  committed
  edfe0a6

  Copy full SHA for edfe0a6

• add buffer_device as optional input

  mori360

  committed
  f6f9393

  Copy full SHA for f6f9393

• add ngpu for 1D test

  mori360

  committed
  d6840dd

  Copy full SHA for d6840dd

Commits on Oct 25, 2024

• modify config help, update condition logic

  mori360

  committed
  ab1e258

  Copy full SHA for ab1e258

Commits on Oct 28, 2024

• test cpu offload with pp

  mori360

  committed
  2ca9882

  Copy full SHA for 2ca9882

• move init_device and buffer_device outside pp condition

  mori360

  committed
  7af331c

  Copy full SHA for 7af331c