fix moduledict with AP meta-pytorch/autoparallel#260

Author: xmfan

torchtitan/components/optimizer.py

@@ -351,18 +351,9 @@ def should_manual_allreduce(tokens_per_expert_by_layer):
             tokens_per_expert_by_layer, torch.distributed.tensor.DTensor
         )
 
-    def get_transformer_blocks(model_part):
-        if isinstance(model_part.layers, nn.ModuleDict):
-            # regular torchtitan
-            blocks = model_part.layers.values()
-        else:
-            # TODO: fix autoparallel to preserve the module dict
-            blocks = model_part.layers.children()
-        return blocks
-
     def _should_register_moe_balancing_hook(model_parts: list[nn.Module]) -> bool:
         for model_part in model_parts:
-            for transformer_block in get_transformer_blocks(model_part):
+            for transformer_block in model_part.layers.values():
                 if is_moe_block(transformer_block):
                     # Assumption: load_balance_coeff is set universally on all moe blocks.
                     return bool(transformer_block.moe.load_balance_coeff)
@@ -384,8 +375,7 @@ def _update_expert_bias(
         # default compute stream. Need to assess if this is OK performance-wise.
         tokens_per_expert_list = []
         for model_part in model_parts:
-            blocks = get_transformer_blocks(model_part)
-            for transformer_block in blocks:
+            for transformer_block in model_part.layers.values():
                 if not is_moe_block(transformer_block):
                     continue
                 if transformer_block.moe.load_balance_coeff is None:
@@ -414,8 +404,7 @@ def _update_expert_bias(
         moe_layer_idx = 0
         with torch.no_grad():
             for model_part in model_parts:
-                blocks = get_transformer_blocks(model_part)
-                for transformer_block in blocks:
+                for transformer_block in model_part.layers.values():
                     if not is_moe_block(transformer_block):
                         continue
                     moe = transformer_block.moe