[zero] add constant placement policy #1705

1SAA · 2022-10-14T03:32:23Z

What does this pr do?

fixes memory leak when paramter is in fp16 in ZeroDDP init.
bans chunk releasement in CUDA. Only when a chunk is about to offload, it is allowed to release.
adds a constant placement policy. With it, users can allocate a reserved caching memory space for parameters.

1SAA added 17 commits September 21, 2022 18:18

[zero] use ChunkV2 for zero training

41121cb

[zero] fix cuda, cpu memory leak

d5f83d7

[zero] add zero optimizer

04082a4

[zero] add unit test for ZeROOptimV2

7f31381

[zero] make fp32_param in pin_memory

c01c62e

[zero] add state_dict, load_state_dict, unit test for ZeRODDP

9f5195b

[hotfix] fix the init device for each chunk

5feb621

[hotfix] fix error for gathered chunks

0f60d90

[zero] add support for different paramter groups

81f33bb

[zero] add state_dict for zero optimizer

b7fd676

[polish] polish zero code and add docstrings

d36ff1c

[refactor] refactor zero and its test files

662e38d

[polish] fix all unit tests with zero

dac2388

Merge remote-tracking branch 'origin/demo_moe' into experiment

a02071f

[zero] baned chunk release before moving

a7f495e

[zero] add const placement policy

8936eae

merged from upstream/main

a00e0ca

1SAA requested review from ver217 and feifeibear October 14, 2022 03:39

ver217 approved these changes Oct 14, 2022

View reviewed changes

1SAA added the Run Build and Test label Oct 14, 2022

fix colo init

2f217cb

1SAA merged commit 1468e4b into hpcaitech:main Oct 14, 2022