diff --git a/colossalai/zero/sharded_optim/sharded_adam.py b/colossalai/zero/sharded_optim/sharded_optim_v2.py similarity index 99% rename from colossalai/zero/sharded_optim/sharded_adam.py rename to colossalai/zero/sharded_optim/sharded_optim_v2.py index 52b0bf5f7..dc8b48db8 100644 --- a/colossalai/zero/sharded_optim/sharded_adam.py +++ b/colossalai/zero/sharded_optim/sharded_optim_v2.py @@ -22,7 +22,7 @@ class OptimState(Enum): UNSCALED = 2 -class ShardedAdam(ColossalaiOptimizer): +class ShardedOptimizerV2(ColossalaiOptimizer): def __init__(self, adam_optim: Optimizer,