Merge pull request #65 from foundation-model-stack/meta_init

switch to new meta device init method

fms_fsdp/policies/__init__.py

@@ -1,3 +1,4 @@
 from .ac_handler import apply_fsdp_checkpointing
 from .mixed_precision import *
+from .param_init import param_init_function
 from .wrapping import get_llama_wrapper

fms_fsdp/policies/param_init.py

@@ -0,0 +1,18 @@
+import torch
+from fms.modules.attention import MultiHeadAttention
+from fms.modules.embedding import WordEmbedding
+from fms.modules.feedforward import GatedLinearUnit
+from fms.modules.layernorm import LayerNormParameterized
+
+
+# for details, read https://github.com/foundation-model-stack/fms-fsdp/issues/64
+def param_init_function(module):
+    if (
+        isinstance(module, MultiHeadAttention)
+        or isinstance(module, WordEmbedding)
+        or isinstance(module, GatedLinearUnit)
+        or isinstance(module, LayerNormParameterized)
+    ):
+        module.to_empty(device=torch.cuda.current_device())
+        with torch.no_grad():
+            module.reset_parameters()

fms_fsdp/utils/train_utils.py

@@ -180,17 +180,16 @@ def setup_environ_flags():
 
 
 def get_policies(cfg, rank):
-    """Get the policies for mixed precision and fsdp wrapping and sharding strategy"""
+    """Get policies for mixed precision, FSDP wrapping, sharding strategy and param init function."""
 
+    # mixed precision
     verify_bfloat_support = (
         torch.version.cuda
         and torch.cuda.is_bf16_supported()
         and packaging.version.parse(torch.version.cuda).release >= (11, 0)
         and dist.is_nccl_available()
         and nccl.version() >= (2, 10)
     )
-
-    # mixed precision
     if cfg.mixed_precision:
         bf16_ready = verify_bfloat_support
         if bf16_ready:
@@ -219,7 +218,13 @@ def get_policies(cfg, rank):
     if rank == 0:
         print(f"Sharding strategy = {cfg.sharding_strategy}")
 
-    return mixed_precision_policy, wrapping_policy, sharding_strategy
+    # param init function
+    if cfg.low_cpu_fsdp:
+        param_init_fn = param_init_function
+    else:
+        param_init_fn = None
+
+    return mixed_precision_policy, wrapping_policy, sharding_strategy, param_init_fn
 
 
 def get_profiler(cfg, rank):

main_training.py

@@ -46,20 +46,18 @@ def main(**kwargs):
     setup_environ_flags()
 
     # get policy
-    mixed_precision_policy, wrapping_policy, sharding_strategy_policy = get_policies(
-        cfg, rank
-    )
+    (
+        mixed_precision_policy,
+        wrapping_policy,
+        sharding_strategy_policy,
+        param_init_fn,
+    ) = get_policies(cfg, rank)
 
     # get fms model
     llama_config = get_model_config(cfg.model_variant)
-
     if cfg.low_cpu_fsdp:
-        if rank == 0:
+        with torch.device("meta"):
             model = LLaMA(llama_config)
-            model.reset_parameters()
-        else:
-            with torch.device("meta"):
-                model = LLaMA(llama_config)
     else:
         model = LLaMA(llama_config)
         model.reset_parameters()
@@ -87,12 +85,7 @@ def main(**kwargs):
         use_orig_params=cfg.use_torch_compile,
         device_id=torch.cuda.current_device(),
         limit_all_gathers=True,
-        sync_module_states=cfg.low_cpu_fsdp,
-        param_init_fn=lambda module: (
-            module.to_empty(device=torch.device("cuda"), recurse=False)
-            if cfg.low_cpu_fsdp
-            else None
-        ),
+        param_init_fn=param_init_fn,
     )
     # we need this post-fsdp call to avoid graph break with torch.compile, until we figure out a better solution.
     model.rot_emb.compute_freqs_cis(