comments

Author: xmfan

torchtitan/components/optimizer.py

@@ -340,11 +340,6 @@ def build_optimizers_with_moe_load_balancing(
         ft_manager=ft_manager,
     )
 
-    def should_manual_allreduce(tokens_per_expert_by_layer):
-        return not isinstance(
-            tokens_per_expert_by_layer, torch.distributed.tensor.DTensor
-        )
-
     def _should_register_moe_balancing_hook(model_parts: list[nn.Module]) -> bool:
         for model_part in model_parts:
             for transformer_block in model_part.layers.values():
@@ -385,7 +380,9 @@ def _update_expert_bias(
         tokens_per_expert_by_layer = torch.vstack(tokens_per_expert_list)
 
         if dp_cp_mesh is not None:
-            if should_manual_allreduce(tokens_per_expert_by_layer):
+            if isinstance(tokens_per_expert_by_layer, torch.distributed.tensor.DTensor):
+                tokens_per_expert_by_layer = tokens_per_expert_by_layer.full_tensor()
+            else:
                 # Perform single all-reduce to get global statistics across all processes
                 pg = dp_cp_mesh.get_group()
                 torch.distributed.all_reduce(

torchtitan/experiments/auto_parallel/job_config.py

@@ -14,25 +14,6 @@
 
 @dataclass
 class Experimental:
-    custom_import: str = ""
-    """
-    This option enables the importation of external modules.
-    Currently, it only supports dotted import modules (e.g., some_package.model_x).
-    It is the user's responsibility to ensure that the specified path can be
-    successfully imported. One method to achieve this, you can place your module
-    inside the ``torchtitan/torchtitan`` folder and execute ``pip install -e .`` to
-    make it available for import.
-    """
-
-    custom_args_module: str = ""
-    """
-    DEPRECATED (moved to Job.custom_config_module). Will be removed soon.
-
-    This option allows users to extend TorchTitan's existing JobConfig by extending
-    a user defined JobConfig dataclass. Similar to ``--experimental.custom_import``, the user
-    needs to ensure that the path can be imported.
-    """
-
     # "aten" (default), "inductor", "none"
     comms_bucket_reorder_strategy: str = "aten"