Re-enable FSDP+TP w/ strided sharding

ghstack-source-id: 13d0d4c
Pull Request resolved: #507

Author: XilunWu

torchtitan/parallelisms/parallelize_llama.py

@@ -29,6 +29,7 @@
 from torchtitan.config_manager import JobConfig, TORCH_DTYPE_MAP
 from torchtitan.logging import logger
 from torchtitan.parallelisms.parallel_dims import ParallelDims
+from torchtitan.parallelisms.utils import check_strided_sharding_enabled
 
 
 def parallelize_llama(
@@ -313,17 +314,8 @@ def apply_fsdp(
         )
     fully_shard(model, **fsdp_config, reshard_after_forward=not pp_enabled)
 
-    if pp_enabled:
-        # TODO
-        # This PR https://github.com/pytorch/pytorch/pull/129519 added a safety check to avoid using 2D/3D DCP since
-        # without strided sharding, DCP can not safely support resharding for 2D/3D.  However, for PP to work, even
-        # without resharding, we load a seed-checkpoint and need to disable the safety mechanism.  This hack should be
-        # removed after strided sharding is landed in DCP.
-        for module in model.modules():
-            assert len(module._load_state_dict_pre_hooks) <= 1
-            module._load_state_dict_pre_hooks.clear()
-            assert len(module._state_dict_pre_hooks) <= 1
-            module._state_dict_pre_hooks.clear()
+    # check if strided sharding is enabled, which is necessary for 2D/3D DCP
+    check_strided_sharding_enabled()
 
     logger.info("Applied FSDP to the model")
 

torchtitan/parallelisms/utils.py

@@ -0,0 +1,30 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+
+from torchtitan.logging import logger
+
+
+def check_strided_sharding_enabled() -> None:
+    # Correct 2D/3D DCP usage requires DTensor's strided sharding in PR
+    # https://github.com/pytorch/pytorch/pull/130760. This function checks if users'
+    # PyTorch nightly-build version is newer than 2024-08-09 to make sure this PR is
+    # included when 2D/3D DCP is used.
+    if "git" in torch.__version__:  # pytorch is built from source
+        # notify users to check if the commit hash is newer than 2024-08-09
+        logger.warning(
+            "detected that the pytorch is built from source. Please make sure the PR "
+            "(https://github.com/pytorch/pytorch/pull/130760) is included in pytorch "
+            "for correct 2D/3D DCP usage."
+        )
+    elif torch.__version__ < "2.5.0.dev20240809":
+        # the nightly build pytorch was built before 2024-08-09
+        logger.warning(
+            f"detected that the pytorch version {torch.__version__} is older than "
+            "2.5.0.dev20240809. Please upgrade a newer version to include the change "
+            "made in https://github.com/pytorch/pytorch/pull/130760 for correct 2D/3D "
+            "DCP usage."
+        )