Performance improvements and optimizer fixes

config_files/training/config_lorem_ipsum_long_moe_ep_fsdp2.yaml

@@ -158,7 +158,7 @@ checkpoint_saving:
       component_key: checkpoint_saving_strategy
       variant_key: save_k_most_recent_checkpoints_strategy
       config:
-        k: -1
+        k: 0
     checkpoint_saving_execution:
       component_key: checkpoint_saving_execution
       variant_key: dcp
@@ -176,6 +176,9 @@ loss_fn:
     model:
       instance_key: model_raw
       pass_type: BY_REFERENCE
+    device_mesh:
+      instance_key: device_mesh
+      pass_type: BY_REFERENCE
 
 device_mesh:
   component_key: device_mesh
@@ -221,7 +224,7 @@ initialized_model:
       component_key: model_initialization
       variant_key: composed
       config:
-        model_type: gpt2
+        model_type: qwen_moe
         weight_init_type: scaled
         mean: 0.0
         std: 0.02
@@ -238,9 +241,6 @@ ep_model:
     device_mesh:
       instance_key: device_mesh
       pass_type: BY_REFERENCE
-    mixed_precision_settings:
-      param_dtype: BF_16
-      reduce_dtype: BF_16
     block_names: [TransformerBlock]
 
 ac_model:
@@ -279,22 +279,22 @@ model_raw:
     prediction_key: ${loss_fn.config.prediction_key}
     vocab_size: 50304
     max_seq_len: ${settings.step_profile.sequence_length}
-    d_model: 128
+    d_model: 2048
     n_heads: 8
     n_kv_heads: 4
-    num_layers: 2
-    d_ff: 128
+    num_layers: 48
+    d_ff: 6144
     attn_dropout: 0.0
     ffn_dropout: 0.0
     tie_embeddings: false
     norm_eps: 1e-6
     rope_base: 1000000.0
-    moe_num_experts: 8
-    moe_top_k: 2
-    moe_d_ff: 128
-    moe_capacity_factor: 1.25
+    moe_num_experts: 128
+    moe_top_k: 8
+    moe_d_ff: 768
     moe_min_capacity: 4
     moe_overflow_policy: residual
+    moe_expert_backend: auto
     moe_router_noise_std: 0.0
     moe_router_temperature: 1.0
     moe_router_dropout: 0.0
@@ -318,7 +318,7 @@ lr_scheduler:
 
 optimizer:
   component_key: optimizer
-  variant_key: ep_adam_w
+  variant_key: adam_w
   config:
     lr: 0.0001
     betas: [0.9, 0.95]
@@ -328,9 +328,6 @@ optimizer:
     wrapped_model:
       instance_key: initialized_model
       pass_type: BY_REFERENCE
-    device_mesh:
-      instance_key: device_mesh
-      pass_type: BY_REFERENCE
 
 gradient_clipper:
   component_key: gradient_clipper
@@ -370,8 +367,11 @@ evaluation_subscriber:
 
 mfu_calculator:
   component_key: mfu_calculator
-  variant_key: gpt2
+  variant_key: moe
   config:
+    num_experts: ${model_raw.config.moe_num_experts}
+    top_k: ${model_raw.config.moe_top_k}
+    moe_d_ff: ${model_raw.config.moe_d_ff}
     n_layer: ${model_raw.config.num_layers}
     sequence_length: ${settings.step_profile.sequence_length}
     n_embd: ${model_raw.config.d_model}

src/modalities/checkpointing/stateful/app_state.py

@@ -184,16 +184,6 @@ def load_state_dict_(app_state: AppState, state_dict: dict[str, Any]) -> None:
 
 
 class OptimizerStateRetriever(StateRetrieverIF):
-    @staticmethod
-    def _uses_standard_optimizer_state_dict(app_state: AppState) -> bool:
-        """Checks whether the optimizer state dict follows the standard torch Optimizer schema.
-
-        Standard optimizer state dicts contain top-level "state" and "param_groups" keys,
-        which are required by distributed optimizer checkpoint utilities.
-        """
-        state_dict = app_state.optimizer.state_dict()
-        return isinstance(state_dict, dict) and "state" in state_dict and "param_groups" in state_dict
-
     @staticmethod
     def get_state_dict(app_state: AppState) -> dict[str, Any]:
         """Returns the state dict of the optimizer in the AppState object.
@@ -206,10 +196,6 @@ def get_state_dict(app_state: AppState) -> dict[str, Any]:
         """
         if isinstance(app_state.optimizer, OptimizersList):
             sd = app_state.optimizer.state_dict()
-        elif not OptimizerStateRetriever._uses_standard_optimizer_state_dict(app_state):
-            # Custom optimizers (e.g. EP wrappers) may not expose the standard torch
-            # optimizer format expected by get_optimizer_state_dict.
-            sd = app_state.optimizer.state_dict()
         else:
             assert len(app_state.model_parts) == 1, "Expected a single model part for non-OptimizersList optimizer."
             sd = get_optimizer_state_dict(
@@ -231,8 +217,6 @@ def load_state_dict_(app_state: AppState, state_dict: dict[str, Any]) -> None:
         """
         if isinstance(app_state.optimizer, OptimizersList):
             app_state.optimizer.load_state_dict(state_dict)
-        elif not OptimizerStateRetriever._uses_standard_optimizer_state_dict(app_state):
-            app_state.optimizer.load_state_dict(state_dict)
         else:
             assert len(app_state.model_parts) == 1, "Expected a single model part for non-OptimizersList optimizer."
             set_optimizer_state_dict(

src/modalities/config/config.py

@@ -87,6 +87,7 @@ class MoECrossEntropyLossConfig(BaseModel):
     target_key: str
     prediction_key: str
     model: Any
+    device_mesh: Optional[PydanticDeviceMeshIFType] = None
     tag: str = "MoECrossEntropyLoss"
 
     class Config:
@@ -177,19 +178,6 @@ class AdamWOptimizerConfig(BaseModel):
     fused: bool | None = None
 
 
-class EPAdamWConfig(BaseModel):
-    wrapped_model: PydanticPytorchModuleOrListType
-    device_mesh: PydanticDeviceMeshIFType
-    lr: float
-    betas: tuple[float, float]
-    eps: float
-    weight_decay: float
-    weight_decay_groups_excluded: list[str]
-
-    class Config:
-        arbitrary_types_allowed = True
-
-
 class DummyLRSchedulerConfig(BaseModel):
     optimizer: PydanticOptimizerIFType
 
@@ -338,7 +326,6 @@ class EPWrappedModelConfig(BaseModel):
     model: PydanticPytorchModuleOrListType
     block_names: list[str]
     device_mesh: PydanticDeviceMeshIFType
-    mixed_precision_settings: FSDP2MixedPrecisionSettings
 
 
 class DebuggingEnrichedModelConfig(BaseModel):
@@ -545,6 +532,12 @@ class GPT2MFUCalculatorConfig(BaseModel):
     device_mesh: Optional[PydanticDeviceMeshIFType] = None
 
 
+class MoEMFUCalculatorConfig(GPT2MFUCalculatorConfig):
+    num_experts: Annotated[int, Field(strict=True, gt=0)]
+    top_k: Annotated[int, Field(strict=True, gt=0)]
+    moe_d_ff: Annotated[int, Field(strict=True, gt=0)]
+
+
 class ParallelDegreeConfig(BaseModel):
     device_mesh: PydanticDeviceMeshIFType
     parallelism_methods: list[ParallelismDegrees]

src/modalities/loss_functions.py

@@ -1,10 +1,16 @@
 from abc import ABC, abstractmethod
-from typing import overload
+from typing import Optional, overload
 
 import torch
+from torch.distributed.device_mesh import DeviceMesh
 from torch.nn import CrossEntropyLoss
 
 from modalities.batch import InferenceResultBatch
+from modalities.running_env.fsdp.device_mesh import (
+    ParallelismDegrees,
+    get_mesh_for_parallelism_method,
+    has_parallelism_method,
+)
 
 
 class Loss(ABC):
@@ -88,20 +94,39 @@ def _parse_arguments(
 
 
 class MoECrossEntropyLoss(Loss):
-    """Cross entropy loss with optional MoE auxiliary losses from model layers."""
+    """Cross entropy loss plus the global-batch load balancing loss (and router z-loss) of the MoE layers."""
 
     def __init__(
         self,
         target_key: str,
         prediction_key: str,
         model,
+        device_mesh: Optional[DeviceMesh] = None,
         tag: str = "MoECrossEntropyLoss",
     ):
+        """
+        Args:
+            target_key (str): Key of the targets in the batch.
+            prediction_key (str): Key of the logits in the model output.
+            model: The (unwrapped or FSDP-wrapped) MoE model whose MoE blocks provide the router statistics.
+            device_mesh (Optional[DeviceMesh]): Device mesh; the load balancing token counts are aggregated over
+                its data parallel dimensions. If None, they are aggregated over the default (world) group.
+            tag (str): Tag of the loss.
+        """
+        from modalities.models.moe.qwen_model import GlobalBatchLoadBalancingLoss
+
         super().__init__(tag)
         self.target_key = target_key
         self.prediction_key = prediction_key
-        self.model = model
         self.loss_fun = CrossEntropyLoss(reduction="mean")
+        process_groups = None
+        if device_mesh is not None:
+            process_groups = [
+                get_mesh_for_parallelism_method(device_mesh, method).get_group()
+                for method in (ParallelismDegrees.DP_REPLICATE, ParallelismDegrees.DP_SHARD)
+                if has_parallelism_method(device_mesh, method)
+            ]
+        self.load_balancing_loss = GlobalBatchLoadBalancingLoss(model=model, process_groups=process_groups)
 
     def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor:
         labels = forward_batch.get_targets(self.target_key)
@@ -113,10 +138,9 @@ def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor:
             labels.contiguous().long().view(-1),
         )
 
-        for layer in self.model.layers.values():
-            if hasattr(layer, "aux_loss") and layer.aux_loss is not None:
-                loss = loss + layer.aux_loss.to(loss.dtype)
-
+        moe_loss = self.load_balancing_loss()
+        if moe_loss is not None:
+            loss = loss + moe_loss.to(loss.dtype)
         return loss