From a801e99700c9583423fa76a6126fe4d2cc966521 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 23 Jun 2026 11:06:35 +0000 Subject: [PATCH 1/2] refactor: Merge moe model factory into general model factory --- src/modalities/models/model_factory.py | 94 +++++++++++++++++- src/modalities/models/moe/__init__.py | 2 - src/modalities/models/moe/model_factory.py | 96 ------------------- src/modalities/registry/components.py | 3 +- .../test_distributed_multidim_dataloader.py | 2 + 5 files changed, 96 insertions(+), 101 deletions(-) delete mode 100644 src/modalities/models/moe/model_factory.py diff --git a/src/modalities/models/model_factory.py b/src/modalities/models/model_factory.py index acef23f71..c8c9c9b9a 100644 --- a/src/modalities/models/model_factory.py +++ b/src/modalities/models/model_factory.py @@ -5,6 +5,7 @@ import itertools import json import time +import warnings from dataclasses import asdict, dataclass from functools import partial from pathlib import Path @@ -18,6 +19,7 @@ from torch.distributed.fsdp import FSDPModule as FSDP2 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP1 from torch.distributed.fsdp import ShardingStrategy +from torch.distributed.tensor import DTensor from torch.distributed.tensor.parallel import ( ColwiseParallel, PrepareModuleInput, @@ -41,9 +43,10 @@ TransformerMLP, ) from modalities.models.model import ActivationType +from modalities.models.parallelism.expert_parallelism import ExpertParallel from modalities.nn.model_initialization.initialization_if import ModelInitializationIF from modalities.running_env.env_utils import FSDP2MixedPrecisionSettings, MixedPrecisionSettings -from modalities.running_env.fsdp.device_mesh import ParallelismDegrees +from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, get_mesh_for_parallelism_method from modalities.running_env.fsdp.fsdp_auto_wrapper import FSDPTransformerAutoWrapPolicyFactory from modalities.training.activation_checkpointing.activation_checkpointing import ( ActivationCheckpointing, @@ -253,6 +256,95 @@ def get_fsdp2_wrapped_model( ) return model + @staticmethod + def get_ep_wrapped_model( + model, + block_names: list[str], + device_mesh: DeviceMesh, + mixed_precision_settings: FSDP2MixedPrecisionSettings, + ) -> nn.Module: + """Apply expert parallel wrapping for selected MoE blocks.""" + block_types = [] + missing_block_names = [] + for name in block_names: + block_type = get_module_class_from_name(model, name) + if block_type is None: + missing_block_names.append(name) + else: + block_types.append(block_type) + + if len(missing_block_names) > 0 and ( + not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 + ): + warnings.warn( + f"Could not resolve some requested MoE block names and they will be ignored: {missing_block_names}", + stacklevel=2, + ) + + block_types = tuple(block_types) + + if len(block_types) == 0: + raise ValueError(f"None of the requested MoE block names were found: {block_names}") + + ep_mesh = get_mesh_for_parallelism_method(device_mesh, ParallelismDegrees.EP) + target_dtype = mixed_precision_settings.param_dtype.value + + wrapped_blocks = 0 + for module in model.modules(): + if isinstance(module, block_types): + if hasattr(module, "experts"): + ep_target = module + elif (ffn := getattr(module, "ffn", None)) is not None and hasattr(ffn, "experts"): + ep_target = ffn + else: + raise ValueError( + f"Module {type(module).__name__} has no EP-compatible experts location. " + "Expected `experts` or `ffn.experts`." + ) + + if getattr(ep_target, "_ep_enabled", False): + continue + + experts = ep_target.experts + missing = [a for a in ("w1", "w2") if not hasattr(experts, a)] + if missing: + raise ValueError( + f"Module {type(ep_target).__name__}.experts is not grouped-experts compatible. " + f"Missing: {missing}" + ) + if experts.w1.ndim != 3 or experts.w2.ndim != 3: + raise ValueError( + f"Expected grouped expert parameters with ndim=3. Got w1.ndim={experts.w1.ndim}, " + f"w2.ndim={experts.w2.ndim}" + ) + + ep_target._ep_mesh = ep_mesh + ep_target._ep_group = ep_mesh.get_group() + ep_target._ep_size = ep_mesh.size() + ep_target._ep_rank = ep_mesh.get_local_rank() + + ep_target.experts = ExpertParallel()._apply(ep_target.experts, ep_mesh) + setattr(ep_target.experts, "_ep_enabled", True) + + for pname, p in list(ep_target.experts._parameters.items()): + if isinstance(p, DTensor) and p.dtype != target_dtype: + local = p.to_local().to(target_dtype) + ep_target.experts._parameters[pname] = nn.Parameter( + DTensor.from_local(local, p.device_mesh, p.placements, run_check=False), + requires_grad=p.requires_grad, + ) + + wrapped_blocks += 1 + + if wrapped_blocks == 0: + raise ValueError(f"No blocks matched the requested types: {[t.__name__ for t in block_types]}") + + model._ep_wrapped = True + model._ep_mesh = ep_mesh + model._ep_num_wrapped_blocks = wrapped_blocks + + return model + @staticmethod def get_weight_initialized_model(model: nn.Module, model_initializer: ModelInitializationIF) -> nn.Module: """ diff --git a/src/modalities/models/moe/__init__.py b/src/modalities/models/moe/__init__.py index 5e55327a1..198dbd8f9 100644 --- a/src/modalities/models/moe/__init__.py +++ b/src/modalities/models/moe/__init__.py @@ -1,10 +1,8 @@ from modalities.models.moe.loss_functions import MoECrossEntropyLoss -from modalities.models.moe.model_factory import get_ep_wrapped_model from modalities.models.moe.qwen_model import QwenModel, QwenModelConfig __all__ = [ "MoECrossEntropyLoss", "QwenModel", "QwenModelConfig", - "get_ep_wrapped_model", ] diff --git a/src/modalities/models/moe/model_factory.py b/src/modalities/models/moe/model_factory.py deleted file mode 100644 index 35e9b7e10..000000000 --- a/src/modalities/models/moe/model_factory.py +++ /dev/null @@ -1,96 +0,0 @@ -import warnings - -import torch.distributed as dist -import torch.nn as nn -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor - -from modalities.models.parallelism.expert_parallelism import ExpertParallel -from modalities.running_env.env_utils import FSDP2MixedPrecisionSettings -from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, get_mesh_for_parallelism_method -from modalities.util import get_module_class_from_name - - -def get_ep_wrapped_model( - model, - block_names: list[str], - device_mesh: DeviceMesh, - mixed_precision_settings: FSDP2MixedPrecisionSettings, -) -> nn.Module: - block_types = [] - missing_block_names = [] - for name in block_names: - block_type = get_module_class_from_name(model, name) - if block_type is None: - missing_block_names.append(name) - else: - block_types.append(block_type) - - if len(missing_block_names) > 0 and (not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0): - warnings.warn( - f"Could not resolve some requested MoE block names and they will be ignored: {missing_block_names}", - stacklevel=2, - ) - - block_types = tuple(block_types) - - if len(block_types) == 0: - raise ValueError(f"None of the requested MoE block names were found: {block_names}") - - ep_mesh = get_mesh_for_parallelism_method(device_mesh, ParallelismDegrees.EP) - target_dtype = mixed_precision_settings.param_dtype.value - - wrapped_blocks = 0 - for module in model.modules(): - if isinstance(module, block_types): - if hasattr(module, "experts"): - ep_target = module - elif (ffn := getattr(module, "ffn", None)) is not None and hasattr(ffn, "experts"): - ep_target = ffn - else: - raise ValueError( - f"Module {type(module).__name__} has no EP-compatible experts location. " - "Expected `experts` or `ffn.experts`." - ) - - if getattr(ep_target, "_ep_enabled", False): - continue - - experts = ep_target.experts - missing = [a for a in ("w1", "w2") if not hasattr(experts, a)] - if missing: - raise ValueError( - f"Module {type(ep_target).__name__}.experts is not grouped-experts compatible. Missing: {missing}" - ) - if experts.w1.ndim != 3 or experts.w2.ndim != 3: - raise ValueError( - f"Expected grouped expert parameters with ndim=3. Got w1.ndim={experts.w1.ndim}, " - f"w2.ndim={experts.w2.ndim}" - ) - - ep_target._ep_mesh = ep_mesh - ep_target._ep_group = ep_mesh.get_group() - ep_target._ep_size = ep_mesh.size() - ep_target._ep_rank = ep_mesh.get_local_rank() - - ep_target.experts = ExpertParallel()._apply(ep_target.experts, ep_mesh) - ep_target.experts._ep_enabled = True - - for pname, p in list(ep_target.experts._parameters.items()): - if isinstance(p, DTensor) and p.dtype != target_dtype: - local = p.to_local().to(target_dtype) - ep_target.experts._parameters[pname] = nn.Parameter( - DTensor.from_local(local, p.device_mesh, p.placements, run_check=False), - requires_grad=p.requires_grad, - ) - - wrapped_blocks += 1 - - if wrapped_blocks == 0: - raise ValueError(f"No blocks matched the requested types: {[t.__name__ for t in block_types]}") - - model._ep_wrapped = True - model._ep_mesh = ep_mesh - model._ep_num_wrapped_blocks = wrapped_blocks - - return model diff --git a/src/modalities/registry/components.py b/src/modalities/registry/components.py index 71eb2c8ad..f36229917 100644 --- a/src/modalities/registry/components.py +++ b/src/modalities/registry/components.py @@ -100,7 +100,6 @@ from modalities.models.huggingface.huggingface_model import HuggingFacePretrainedModel, HuggingFacePretrainedModelConfig from modalities.models.model_factory import GPT2ModelFactory, ModelFactory from modalities.models.moe.loss_functions import MoECrossEntropyLoss -from modalities.models.moe.model_factory import get_ep_wrapped_model from modalities.models.moe.qwen_model import QwenModel, QwenModelConfig from modalities.models.parallelism.pipeline_parallelism import ComponentSelectorFromPipeline, PipelineFactory from modalities.models.parallelism.pipeline_parallelism_configs import ( @@ -196,7 +195,7 @@ class ComponentEntity: # models ComponentEntity("model", "gpt2", GPT2ModelFactory.get_gpt2_model, GPT2LLMConfig), ComponentEntity("model", "moe", QwenModel, QwenModelConfig), - ComponentEntity("model", "ep_wrapped", get_ep_wrapped_model, EPWrappedModelConfig), + ComponentEntity("model", "ep_wrapped", ModelFactory.get_ep_wrapped_model, EPWrappedModelConfig), ComponentEntity( "model", "gpt2_tp", maybe_model_list(GPT2ModelFactory.get_gpt2_tensor_parallelized_model), GPT2ModelTPConfig ), diff --git a/tests/dataloader/distributed/test_distributed_multidim_dataloader.py b/tests/dataloader/distributed/test_distributed_multidim_dataloader.py index a2767bae4..935ad66d2 100644 --- a/tests/dataloader/distributed/test_distributed_multidim_dataloader.py +++ b/tests/dataloader/distributed/test_distributed_multidim_dataloader.py @@ -93,6 +93,7 @@ def _test_dataloader_produces_different_samples_in_different_dp_ranks(process_id tensor_parallel_degree=2, pipeline_parallel_degree=2, context_parallel_degree=1, + expert_parallel_degree=1, enable_loss_parallel=False, world_size=world_size, ) @@ -117,6 +118,7 @@ def _test_dataloader_produces_same_samples_in_connected_non_dp_ranks(process_id: tensor_parallel_degree=2, pipeline_parallel_degree=2, context_parallel_degree=1, + expert_parallel_degree=1, enable_loss_parallel=False, world_size=world_size, ) From daea062a6e656242cd63014752f06f4294c54625 Mon Sep 17 00:00:00 2001 From: rrutmann Date: Tue, 23 Jun 2026 12:02:32 +0000 Subject: [PATCH 2/2] refactor: Merge MoECrossEntropyLoss into modalities core --- src/modalities/loss_functions.py | 33 +++++++++++++++++ src/modalities/models/moe/__init__.py | 2 +- src/modalities/models/moe/loss_functions.py | 39 --------------------- src/modalities/registry/components.py | 3 +- tests/models/moe/test_loss_functions.py | 2 +- 5 files changed, 36 insertions(+), 43 deletions(-) delete mode 100644 src/modalities/models/moe/loss_functions.py diff --git a/src/modalities/loss_functions.py b/src/modalities/loss_functions.py index e3be6100d..bc0b1cdec 100644 --- a/src/modalities/loss_functions.py +++ b/src/modalities/loss_functions.py @@ -87,6 +87,39 @@ def _parse_arguments( return labels, lm_logits +class MoECrossEntropyLoss(Loss): + """Cross entropy loss with optional MoE auxiliary losses from model layers.""" + + def __init__( + self, + target_key: str, + prediction_key: str, + model, + tag: str = "MoECrossEntropyLoss", + ): + super().__init__(tag) + self.target_key = target_key + self.prediction_key = prediction_key + self.model = model + self.loss_fun = CrossEntropyLoss(reduction="mean") + + def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: + labels = forward_batch.get_targets(self.target_key) + lm_logits = forward_batch.get_predictions(self.prediction_key) + + labels = labels.to(lm_logits.device) + loss = self.loss_fun( + lm_logits.contiguous().view(-1, lm_logits.size(-1)), + 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) + + return loss + + def nce_loss( embedding1: torch.Tensor, embedding2: torch.Tensor, device: torch.device, is_asymmetric: bool, temperature: float ) -> torch.Tensor: diff --git a/src/modalities/models/moe/__init__.py b/src/modalities/models/moe/__init__.py index 198dbd8f9..35647c543 100644 --- a/src/modalities/models/moe/__init__.py +++ b/src/modalities/models/moe/__init__.py @@ -1,4 +1,4 @@ -from modalities.models.moe.loss_functions import MoECrossEntropyLoss +from modalities.loss_functions import MoECrossEntropyLoss from modalities.models.moe.qwen_model import QwenModel, QwenModelConfig __all__ = [ diff --git a/src/modalities/models/moe/loss_functions.py b/src/modalities/models/moe/loss_functions.py deleted file mode 100644 index 57f30da69..000000000 --- a/src/modalities/models/moe/loss_functions.py +++ /dev/null @@ -1,39 +0,0 @@ -import torch -from torch.nn import CrossEntropyLoss - -from modalities.batch import InferenceResultBatch -from modalities.loss_functions import Loss - - -class MoECrossEntropyLoss(Loss): - """Cross entropy loss with optional MoE auxiliary losses from model layers.""" - - def __init__( - self, - target_key: str, - prediction_key: str, - model, - tag: str = "MoECrossEntropyLoss", - ): - super().__init__(tag) - self.target_key = target_key - self.prediction_key = prediction_key - self.model = model - self.loss_fun = CrossEntropyLoss(reduction="mean") - - def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: - labels = forward_batch.get_targets(self.target_key) - lm_logits = forward_batch.get_predictions(self.prediction_key) - - labels = labels.to(lm_logits.device) - loss = self.loss_fun( - lm_logits.contiguous().view(-1, lm_logits.size(-1)), - labels.contiguous().long().view(-1), - ) - - # Aux loss - 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) - - return loss diff --git a/src/modalities/registry/components.py b/src/modalities/registry/components.py index f36229917..9f5fda81f 100644 --- a/src/modalities/registry/components.py +++ b/src/modalities/registry/components.py @@ -85,7 +85,7 @@ ProgressSubscriberFactory, ResultsSubscriberFactory, ) -from modalities.loss_functions import CLMCrossEntropyLoss +from modalities.loss_functions import CLMCrossEntropyLoss, MoECrossEntropyLoss from modalities.models.coca.coca_model import CoCa, CoCaConfig from modalities.models.coca.collator import CoCaCollateFnConfig, CoCaCollatorFn from modalities.models.components.layer_norms import ( @@ -99,7 +99,6 @@ from modalities.models.gpt2.llama3_like_initialization import Llama3Initializer, Llama3InitializerConfig from modalities.models.huggingface.huggingface_model import HuggingFacePretrainedModel, HuggingFacePretrainedModelConfig from modalities.models.model_factory import GPT2ModelFactory, ModelFactory -from modalities.models.moe.loss_functions import MoECrossEntropyLoss from modalities.models.moe.qwen_model import QwenModel, QwenModelConfig from modalities.models.parallelism.pipeline_parallelism import ComponentSelectorFromPipeline, PipelineFactory from modalities.models.parallelism.pipeline_parallelism_configs import ( diff --git a/tests/models/moe/test_loss_functions.py b/tests/models/moe/test_loss_functions.py index 346b69818..5149aef1b 100644 --- a/tests/models/moe/test_loss_functions.py +++ b/tests/models/moe/test_loss_functions.py @@ -2,7 +2,7 @@ from torch.nn import CrossEntropyLoss from modalities.batch import InferenceResultBatch -from modalities.models.moe.loss_functions import MoECrossEntropyLoss +from modalities.loss_functions import MoECrossEntropyLoss class DummyLayer: