Code

Migrate distributed state dict API#2138

Code

Merged

joecummings merged 22 commits intopytorch:mainpytorch/torchtune:mainfrom mori360:state_dictmori360/torchtune:state_dictCopy head branch name to clipboardJan 8, 2025

+250-196

Commits

Commits on Dec 20, 2024

• change load_from_full_model_state_dict

  mori360

  committed
  de72b9c

  Copy full SHA for de72b9c

• lint

  mori360

  committed
  4b959f0

  Copy full SHA for 4b959f0

• correct input parameter for gather_cpu_state_dict

  mori360

  committed
  2ad64d2

  Copy full SHA for 2ad64d2

• correct test_distribtued

  mori360

  committed
  19cf431

  Copy full SHA for 19cf431

• correct NF4 tensor process, add todo

  mori360

  committed
  7e11c4f

  Copy full SHA for 7e11c4f

• correct load_from_full_model_state_dict witj nf4 tensor

  mori360

  committed
  b22c2e8

  Copy full SHA for b22c2e8

• correct import error

  mori360

  committed
  851fcf8

  Copy full SHA for 851fcf8

• correct nf4 tensor process in load_from_full_model_state_dict

  mori360

  committed
  9409ee8

  Copy full SHA for 9409ee8

• remove print

  mori360

  committed
  581a65c

  Copy full SHA for 581a65c

• modify load_from_full_model_state_dict to optimize memory cost as before

  mori360

  committed
  1b0ba78

  Copy full SHA for 1b0ba78

• change load_from_full_model_state_dict

  mori360

  committed
  ca234c7

  Copy full SHA for ca234c7

• add trainable_only to gather_cpu_state_dict

  mori360

  committed
  9cd28a9

  Copy full SHA for 9cd28a9

• add trainable_only to gather_cpu_state_dict

  mori360

  committed
  bfc7668

  Copy full SHA for bfc7668

• adjust dcp api using

  mori360

  committed
  f368b96

  Copy full SHA for f368b96

• adjust dcp api using

  mori360

  committed
  eda1b9f

  Copy full SHA for eda1b9f

• add cpu_offload at load_from_full_optimizer_state_dict

  mori360

  committed
  db0eeed

  Copy full SHA for db0eeed

• remove is_rank_zero from load_from_full_model_state_dict, add torch version check

  mori360

  committed
  91e9818

  Copy full SHA for 91e9818

• import init_optim

  mori360

  committed
  11e24f1

  Copy full SHA for 11e24f1

• update _USE_DISTRIBUTED_STATE_DICT_API version check

  mori360

  committed
  3d0d26f

  Copy full SHA for 3d0d26f

Commits on Jan 6, 2025

• restructure load_from_full_model_state_dict, rewrite comments on version check, change weights_only to adapter_weights_only

  mori360

  committed
  e4730de

  Copy full SHA for e4730de

Commits on Jan 8, 2025

• Merge branch 'pytorch:main' into state_dict

  mori360

  authored
  ad08f42

  Copy full SHA for ad08f42

• change version check name, add more comments, load_from_full_model_state_dict return

  mori360

  committed
  4dfef98

  Copy full SHA for 4dfef98