sync to internal 1.6.0.dev3 (#4)

Co-authored-by: 乙划 <[email protected]>

Author: skydoorkai

.isort.cfg

@@ -1,5 +1,5 @@
 [settings]
 multi_line_output=3
 line_length=120
-known_third_party = GPy,accelerate,agd,apex,data,datasets,deepspeed,distutils,dlrover,einops,evaluate,example_utils,fairscale,flash_attn,google,grpc,instruction_dataset_utils,matplotlib,model,modeling,networkx,numpy,packaging,pandas,peft,psutil,pymoo,pyomo,pytest,redis,safetensors,scipy,sklearn,tiktoken,torch,torch_npu,torchvision,tqdm,transformers,triton,typing_extensions,utils,yaml
+known_third_party = accelerate,agd,apex,datasets,deepspeed,distutils,dlrover,einops,evaluate,example_utils,fairscale,flash_attn,google,grpc,instruction_dataset_utils,matplotlib,megatron,model,model_registry,moe_modules,networkx,numpy,packaging,pandas,peft,psutil,pytest,redis,safetensors,scipy,seaborn,sklearn,tiktoken,torch,torch_npu,torchvision,tqdm,transformers,triton,typing_extensions,utils,wrapt,yaml
 include_trailing_comma=True

atorch/auto/accelerate.py

@@ -43,25 +43,46 @@ def model_transform(
     assert not strategy.is_tunable()
     record_user_defined_half_precision_dtype(strategy)
     cpu_offload = False
+    has_fsdp2 = False
+    param_init_by_user_fn_already = False
     for opt in strategy:
         opt_name = opt[0]
         opt_config = opt[1]
         model_context = opt_lib[opt_name].transform(model_context, opt_config)
         if opt_name == "fsdp" and opt_config is not None and opt_config.get("cpu_offload", False) is True:
             cpu_offload = True
+        if opt_name == "fsdp2":
+            has_fsdp2 = True
+            model_device = next(model_context.model.parameters()).device
+            if "param_init_fn" in opt_config and model_device is torch.device("cuda"):
+                param_init_by_user_fn_already = True
+
     model_context.adjust_wrappers()
     if apply_wrapper:
         model_context.apply_wrappers(is_pre_wrapper=True)
     if create_dataloader:
         model_context.update_dataloader()
-    if create_optim:
-        model_context.update_optim()
     if use_sample_batch:
         model_context.update_sample_batch()
-    if apply_wrapper:
-        model_context.apply_wrappers(is_pre_wrapper=False)
-    if torch.cuda.is_available() and not model_context.gpu_used and not cpu_offload:
-        reload_meta_module(model_context.model, torch.device(type="cuda", index=local_rank()), False)
+    if not has_fsdp2:
+        if create_optim:
+            model_context.update_optim()
+        if apply_wrapper:
+            model_context.apply_wrappers(is_pre_wrapper=False)
+
+    if (
+        torch.cuda.is_available()
+        and not model_context.gpu_used
+        and not cpu_offload
+        and not param_init_by_user_fn_already
+    ):
+        reload_meta_module(model_context.model, torch.device(type="cuda", index=local_rank()), False, True, has_fsdp2)
+
+    if has_fsdp2:
+        if create_optim:
+            model_context.update_optim()
+        if apply_wrapper:
+            model_context.apply_wrappers(is_pre_wrapper=False)
     return model_context
 
 

atorch/auto/clip_grad_norm.py

@@ -29,7 +29,7 @@ def clip_grad_norm(model, max_norm, norm_type=2, optimizer=None, process_group_n
     Returns:
         Total norm of the parameters (viewed as a single vector) or None if using ds zero optimizer.
     """
-    if isinstance(optimizer, DeepSpeedZeroOptimizer):
+    if DeepSpeedZeroOptimizer is not None and isinstance(optimizer, DeepSpeedZeroOptimizer):
         assert norm_type == 2, "deep speed zero optimizer only supports L2 norm"
         optimizer.clip_grad = max_norm
         return None

atorch/auto/engine/sg_algo/hebo/acq_optimizers/evolution_optimizer.py

@@ -1,16 +1,20 @@
 import numpy as np
 import pandas as pd
-from pymoo.config import Config
-from pymoo.core.problem import Problem
-from pymoo.factory import get_algorithm, get_crossover, get_mutation
-from pymoo.operators.mixed_variable_operator import MixedVariableCrossover, MixedVariableMutation
-from pymoo.optimize import minimize
+
+try:
+    from pymoo.config import Config
+    from pymoo.core.problem import Problem
+    from pymoo.factory import get_algorithm, get_crossover, get_mutation
+    from pymoo.operators.mixed_variable_operator import MixedVariableCrossover, MixedVariableMutation
+    from pymoo.optimize import minimize
+
+    Config.show_compile_hint = False
+except (ImportError, ModuleNotFoundError):
+    Problem = object
 
 from atorch.auto.engine.sg_algo.hebo.acquisitions.acq import Acquisition
 from atorch.auto.engine.sg_algo.hebo.design_space.design_space import DesignSpace
 
-Config.show_compile_hint = False
-
 
 class BOProblem(Problem):
     def __init__(
@@ -26,6 +30,9 @@ def __init__(
         self.space = space
         self.fix = fix  # NOTE: use self.fix to enable contextual BO
 
+        if Problem == object:
+            print("Install pymoo==0.5.0 to support evolution optimizer.")
+
         super().__init__(len(lb), xl=lb, xu=ub, n_obj=acq.num_obj, n_constr=acq.num_constr)
 
     def _evaluate(self, x: np.ndarray, out: dict, *args, **kwargs):

atorch/auto/engine/sg_algo/hebo/models/gauss_process/gpy_wgp.py

@@ -1,10 +1,15 @@
 import logging
 import warnings
 
-import GPy
 import numpy as np
-from GPy.models import InputWarpedGP
-from GPy.util.input_warping_functions import KumarWarping
+
+try:
+    import GPy
+    from GPy.models import InputWarpedGP
+    from GPy.util.input_warping_functions import KumarWarping
+except (ImportError, ModuleNotFoundError):
+    print("Install GPy package to support auto training optimization.")
+
 from sklearn.preprocessing import MinMaxScaler, StandardScaler
 
 from atorch.auto.engine.sg_algo.hebo.models.base_model import BaseModel

atorch/auto/model_context.py

@@ -5,7 +5,7 @@
 import types
 
 try:
-    from collections import abs as collections_abc  # type: ignore[attr-defined]
+    from collections import abc as collections_abc  # type: ignore[attr-defined]
 except ImportError:
     import collections as collections_abc  # type: ignore[no-redef]
 
@@ -28,6 +28,7 @@
 from atorch.distributed.distributed import (
     get_data_partition_rank_and_size,
     local_rank,
+    parallel_group,
     parallel_group_and_ranks,
     parallel_group_size,
     rank,
@@ -396,6 +397,9 @@ def create_optim(self):
                 src = ranks[0]
             torch.distributed._broadcast_coalesced(process_group, module_states, int(250 * 1024 * 1024), src)
 
+        if "fsdp2" in self.pre_wrappers and parallel_group("expert") is not None:
+            self.optim_args["foreach"] = False
+
         if not self.check_pipe_model():
             if not self.optim_param_func:
                 optim = self.optim_func(self.model.parameters(), **self.optim_args)
@@ -416,6 +420,7 @@ def create_optim(self):
             and "ds_zero" not in self.post_wrappers
             and "zero2" not in self.post_wrappers
             and "fsdp" not in self.pre_wrappers
+            and "fsdp2" not in self.pre_wrappers
             and "ds_3d_parallel" not in self.post_wrappers
         ):
             is_cuda = next(self.model.parameters()).is_cuda
@@ -497,6 +502,8 @@ def adjust_wrappers(self):
                 self.pre_wrappers.pop("zero2")
             if "fsdp" in self.pre_wrappers:
                 self.pre_wrappers.pop("fsdp")
+            if "fsdp2" in self.pre_wrappers:
+                self.pre_wrappers.pop("fsdp2")
 
             # DDP is supported and handled internally by PiPPy.
             if "ddp" in self.post_wrappers:
@@ -572,13 +579,18 @@ def adjust_wrappers(self):
         ds_3d_parallel_wrapper_exist = "ds_3d_parallel" in self.post_wrappers
         fairscale_zero2_wrapper_exist = "zero2" in self.post_wrappers
         fsdp_wrapper_exist = "fsdp" in self.pre_wrappers or "zero2" in self.pre_wrappers
+        fsdp2_wrapper_exist = "fsdp2" in self.pre_wrappers
         tensor_parallel_wrapper_exist = "tp" in self.pre_wrappers
         ckpt_wrapper_exist = "checkpoint" in self.post_wrappers
         native_dynamo_wrapper_exist = "native_dynamo" in self.pre_wrappers
 
         # remove ddp wrapper when using zero2
         if ddp_wrapper_exist and (
-            fairscale_zero2_wrapper_exist or fsdp_wrapper_exist or ds_zero_wrapper_exist or ds_3d_parallel_wrapper_exist
+            fairscale_zero2_wrapper_exist
+            or fsdp_wrapper_exist
+            or ds_zero_wrapper_exist
+            or ds_3d_parallel_wrapper_exist
+            or fsdp2_wrapper_exist
         ):
             logger.info("Found Zero, ds_3d_parallel, or pipe wrapper, remove ddp wrapper.")
             self.post_wrappers.pop("ddp")
@@ -587,21 +599,28 @@ def adjust_wrappers(self):
             logger.info("Found fsdp and amp_native wrapper, turn on mixed_precision in FSDP")
             _, amp_native_config = self.post_wrappers["amp_native"]
             fp16_dtype = amp_native_config.get("dtype", torch.float16)
-            mixed_precision_param = (
-                MixedPrecision(param_dtype=fp16_dtype, reduce_dtype=fp16_dtype, buffer_dtype=fp16_dtype)
-                if MixedPrecision
-                else True
-            )
+            mixed_precision_param = {"param_dtype": fp16_dtype, "reduce_dtype": fp16_dtype, "buffer_dtype": fp16_dtype}
             config = self.pre_wrappers["fsdp"][1] or {}
             config["mixed_precision"] = mixed_precision_param
             self.pre_wrappers["fsdp"] = (
                 self.pre_wrappers["fsdp"][0],
                 config,
             )
+        elif fsdp2_wrapper_exist and "amp_native" in self.post_wrappers:
+            logger.info("Found fsdp2 and amp_native wrapper, turn on mixed_precision in FSDP")
+            _, amp_native_config = self.post_wrappers["amp_native"]
+            fp16_dtype = amp_native_config.get("dtype", torch.float16)
+            mixed_precision_param = {"param_dtype": fp16_dtype, "reduce_dtype": fp16_dtype, "buffer_dtype": fp16_dtype}
+            config = self.pre_wrappers["fsdp2"][1] or {}
+            config["mixed_precision"] = mixed_precision_param
+            self.pre_wrappers["fsdp2"] = (
+                self.pre_wrappers["fsdp2"][0],
+                config,
+            )
 
         # move dynamo_native wrapper behind ddp or fsdp (fsdp will adjusted later)
         # Note that dynamo_native wrapper and fsdp wrapper are pre-wrappers while ddp wrapper is a post-wrapper.
-        if native_dynamo_wrapper_exist and ddp_wrapper_exist and not fsdp_wrapper_exist:
+        if native_dynamo_wrapper_exist and ddp_wrapper_exist and not fsdp_wrapper_exist and not fsdp2_wrapper_exist:
             # ddp wrapper is a post-wrapper. Popping dynamo_native wrapper from pre-wrappers
             # then insert it after ddp wrapper.
             post_wrappers_list = []
@@ -616,8 +635,13 @@ def adjust_wrappers(self):
 
         if tensor_parallel_wrapper_exist:
             wrap_cls = None
+            fsdp_wrapper = None
             if fsdp_wrapper_exist and torch_version() >= (1, 12, 0):
                 fsdp_wrapper = self.pre_wrappers["fsdp"]
+            elif fsdp2_wrapper_exist and torch_version() >= (1, 12, 0):
+                fsdp_wrapper = self.pre_wrappers["fsdp2"]
+
+            if fsdp_wrapper is not None:
                 fsdp_wrapper = list(fsdp_wrapper)
                 if fsdp_wrapper[1] is None:
                     fsdp_wrapper[1] = dict()
@@ -644,15 +668,19 @@ def adjust_wrappers(self):
             leaf_modules = _propose_leaf_modules(wrap_cls)
             auto_wrap_cls = _propose_wrap_cls(leaf_modules)
 
-            if fsdp_wrapper_exist and torch_version() >= (1, 12, 0):
+            if (fsdp_wrapper_exist or fsdp2_wrapper_exist) and torch_version() >= (1, 12, 0):
                 if "atorch_wrap_cls" in fsdp_config:
                     if auto_wrap_cls is not None:
                         fsdp_config["atorch_wrap_cls"] = auto_wrap_cls
                     else:
                         fsdp_config.pop("atorch_wrap_cls")
 
                     fsdp_wrapper[1] = fsdp_config
-                    self.pre_wrappers["fsdp"] = tuple(fsdp_wrapper)
+
+                    if fsdp_wrapper_exist:
+                        self.pre_wrappers["fsdp"] = tuple(fsdp_wrapper)
+                    elif fsdp2_wrapper_exist:
+                        self.pre_wrappers["fsdp2"] = tuple(fsdp_wrapper)
 
             if ckpt_wrapper_exist:
                 if auto_wrap_cls is not None:
@@ -671,7 +699,7 @@ def adjust_wrappers(self):
             tensor_parallel_wrapper_item = list(tensor_parallel_wrapper_item)
             tensor_parallel_wrapper_item[1] = list(tensor_parallel_wrapper_item[1])
             tensor_parallel_wrapper_item[1][1]["leaf_modules"] = leaf_modules
-            if fsdp_wrapper_exist or pipe_wrapper_exist:
+            if fsdp_wrapper_exist or fsdp2_wrapper_exist or pipe_wrapper_exist:
                 tensor_parallel_wrapper_item[1][1]["defer_init"] = True
             tensor_parallel_wrapper_item[1] = tuple(tensor_parallel_wrapper_item[1])
             tensor_parallel_wrapper_item = tuple(tensor_parallel_wrapper_item)
@@ -687,7 +715,7 @@ def adjust_wrappers(self):
                 _insert_amp_config_for_tp_ckpt(amp_config)
 
         # adjust pre_wrapper order
-        order_wrapper_name = ["half", "module_replace", "sequence_parallel", "fp8", "fsdp", "native_dynamo"]
+        order_wrapper_name = ["half", "module_replace", "sequence_parallel", "fp8", "fsdp", "fsdp2", "native_dynamo"]
         match_names = []
         for name in self.pre_wrappers:
             if name in order_wrapper_name: