github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-22 05:29:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

misc: add offload warning

Browse Source
This commit is contained in:
Wenhao Chen 2024-03-27 18:13:10 +08:00 committed by アマデウス
parent e614aa34f3
commit a1ab2d374e
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/zero/low_level/low_level_optim.py
View File

@ -1,5 +1,6 @@
# this code is inspired by the DeepSpeed library and implemented with our own design from scratch # this code is inspired by the DeepSpeed library and implemented with our own design from scratch
import copy import copy
import warnings
from contextlib import contextmanager from contextlib import contextmanager
from functools import partial from functools import partial
from typing import Dict, Iterator, List, Optional, Tuple from typing import Dict, Iterator, List, Optional, Tuple
@ -87,6 +88,11 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
self._partition_grads = partition_grad self._partition_grads = partition_grad
self._cpu_offload = cpu_offload self._cpu_offload = cpu_offload
if cpu_offload:
warnings.warn(
"Please enlarge OMP_NUM_THREADS to speed up the CPU computation. "
"Use default OMP_NUM_THREADS=1 would significantly slow down the training."
)
# grad accumulation # grad accumulation
self.require_grad_sync = True self.require_grad_sync = True