From 367c61581806fbaa7e85f47830dd3f69b3e60bd9 Mon Sep 17 00:00:00 2001 From: ver217 Date: Fri, 12 Aug 2022 18:01:02 +0800 Subject: [PATCH] fix nvme docstring (#1450) --- colossalai/nn/optimizer/cpu_adam.py | 4 ++-- colossalai/nn/optimizer/hybrid_adam.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/colossalai/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py index cf6a93ac0..ea08ff723 100644 --- a/colossalai/nn/optimizer/cpu_adam.py +++ b/colossalai/nn/optimizer/cpu_adam.py @@ -46,8 +46,8 @@ class CPUAdam(NVMeOptimizer): True for decoupled weight decay(also known as AdamW) (default: True) simd_log (boolean, optional): whether to show if you are using SIMD to accelerate. (default: False) - nvme_offload_fraction (float, optional): Fraction of params to be offloaded to NVMe. Defaults to 0.0. - offload_dir (Optional[str], optional): Directory to save NVMe offload files. + nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0. + nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files. If it's ``None``, a random temporary directory will be used. Defaults to None. .. _Adam\: A Method for Stochastic Optimization: diff --git a/colossalai/nn/optimizer/hybrid_adam.py b/colossalai/nn/optimizer/hybrid_adam.py index a82c390f1..761843aab 100644 --- a/colossalai/nn/optimizer/hybrid_adam.py +++ b/colossalai/nn/optimizer/hybrid_adam.py @@ -45,8 +45,8 @@ class HybridAdam(NVMeOptimizer): True for decoupled weight decay(also known as AdamW) (default: True) simd_log (boolean, optional): whether to show if you are using SIMD to accelerate. (default: False) - nvme_offload_fraction (float, optional): Fraction of params to be offloaded to NVMe. Defaults to 0.0. - offload_dir (Optional[str], optional): Directory to save NVMe offload files. + nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0. + nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files. If it's ``None``, a random temporary directory will be used. Defaults to None. .. _Adam\: A Method for Stochastic Optimization: