moved distrib opt to own file.

28b63ec7 · Lawrence McAfee · 55695f81 · 28b63ec7 · 28b63ec7 · 28b63ec7
Commit 28b63ec7 authored Mar 09, 2022 by Lawrence McAfee
3 changed files
--- a/megatron/optimizer/__init__.py
+++ b/megatron/optimizer/__init__.py
@@ -25,11 +25,8 @@ from lutil import pax, tp
 # <<<
 from .grad_scaler import ConstantGradScaler, DynamicGradScaler
 # >>>
-from .optimizer import (
+from .optimizer import Float16OptimizerWithFloat16Params, FP32Optimizer
-    Float16OptimizerWithFloat16Params,
+from .distrib_optimizer import DistributedOptimizer
-    Float16DistributedOptimizer,
-    FP32Optimizer,
-)
 # <<<
 def get_param_groups(modules,

--- a/megatron/optimizer/distrib_optimizer.py
+++ b/megatron/optimizer/distrib_optimizer.py
--- a/megatron/optimizer/optimizer.py
+++ b/megatron/optimizer/optimizer.py