diff --git a/config_files/training/config_lorem_ipsum_long_moe_ep_fsdp2.yaml b/config_files/training/config_lorem_ipsum_long_moe_ep_fsdp2.yaml index b17bb89e1..f9e58a536 100644 --- a/config_files/training/config_lorem_ipsum_long_moe_ep_fsdp2.yaml +++ b/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} diff --git a/src/modalities/checkpointing/stateful/app_state.py b/src/modalities/checkpointing/stateful/app_state.py index 25c1efc91..2da3ab236 100644 --- a/src/modalities/checkpointing/stateful/app_state.py +++ b/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( diff --git a/src/modalities/config/config.py b/src/modalities/config/config.py index fea034546..308f97283 100644 --- a/src/modalities/config/config.py +++ b/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] diff --git a/src/modalities/loss_functions.py b/src/modalities/loss_functions.py index bc0b1cdec..b2cde9fbd 100644 --- a/src/modalities/loss_functions.py +++ b/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 diff --git a/src/modalities/models/model_factory.py b/src/modalities/models/model_factory.py index c8c9c9b9a..b1cc267e7 100644 --- a/src/modalities/models/model_factory.py +++ b/src/modalities/models/model_factory.py @@ -19,7 +19,6 @@ 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, @@ -46,7 +45,7 @@ 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, get_mesh_for_parallelism_method +from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, get_expert_meshes from modalities.running_env.fsdp.fsdp_auto_wrapper import FSDPTransformerAutoWrapPolicyFactory from modalities.training.activation_checkpointing.activation_checkpointing import ( ActivationCheckpointing, @@ -215,11 +214,17 @@ def get_fsdp2_wrapped_model( modules = list(model.modules()) - # Collect EP parameters to exclude from FSDP2 sharding - ep_params = { - p for m in model.modules() if getattr(m, "_ep_enabled", False) for p in m.parameters(recurse=False) - } - ignored_params = ep_params if ep_params else None + # Expert-parallel experts are FSDP-sharded on their own mesh before the blocks are wrapped, + # so that the block-level FSDP units do not manage the expert parameters. + expert_modules = [m for m in modules if getattr(m, "_ep_enabled", False)] + if len(expert_modules) > 0: + ModelFactory._fully_shard_experts( + expert_modules=expert_modules, + device_mesh=device_mesh, + mp_policy=mp_policy, + reshard_after_forward=reshard_after_forward, + ) + fsdp_blocks: list[nn.Module] = [] # we first shard all the blocks grouped_modules: list[nn.Module] = [] @@ -235,8 +240,8 @@ def get_fsdp2_wrapped_model( grouped_modules, **fsdp_config, reshard_after_forward=reshard_block_after_forward, - ignored_params=ignored_params, ) + fsdp_blocks.append(grouped_modules[0]) grouped_modules = list() if len(grouped_modules) > 0: @@ -245,23 +250,82 @@ def get_fsdp2_wrapped_model( grouped_modules, **fsdp_config, reshard_after_forward=reshard_block_after_forward, - ignored_params=ignored_params, ) + fsdp_blocks.append(grouped_modules[0]) # finally, we shard the entire model - fully_shard(model, **fsdp_config, reshard_after_forward=reshard_after_forward, ignored_params=ignored_params) + fully_shard(model, **fsdp_config, reshard_after_forward=reshard_after_forward) + if len(expert_modules) > 0 and layers_per_fsdp_unit == 1: + ModelFactory._set_explicit_moe_prefetching(root=model, fsdp_blocks=fsdp_blocks) logger.info( f"Rank {dist.get_rank()} sharded number of parameters: " f"{get_local_number_of_trainable_parameters(model)}" ) return model + @staticmethod + def _fully_shard_experts( + expert_modules: list[nn.Module], + device_mesh: DeviceMesh, + mp_policy: MixedPrecisionPolicy, + reshard_after_forward: bool, + ) -> None: + """FSDP-shards expert-parallel experts over the expert FSDP mesh (dp_shard // ep ranks). + + Each expert's gradient receives contributions from the tokens of all data parallel ranks, + so the gradient divide factor is set to the full data parallel degree to match the dense gradients. + """ + expert_meshes = get_expert_meshes(device_mesh) + if expert_meshes is None: + raise ValueError( + "Found expert-parallel modules, but the device mesh has no expert meshes. " + "Set expert_parallel_degree > 1 in the device mesh config." + ) + efsdp_degree = expert_meshes.expert_fsdp.size(-1) + ep_degree = expert_meshes.ep.size() + for experts in expert_modules: + num_experts = experts.w1.shape[0] + # If there are fewer local experts than efsdp ranks, shard on the hidden dim instead. + shard_placement_fn = (lambda param: Shard(1)) if efsdp_degree * ep_degree > num_experts else None + fully_shard( + experts, + mesh=expert_meshes.expert_fsdp, + mp_policy=mp_policy, + reshard_after_forward=reshard_after_forward, + shard_placement_fn=shard_placement_fn, + ) + experts.set_gradient_divide_factor(float(expert_meshes.gradient_divide_factor)) + if expert_meshes.expert_fsdp.size() > 1: + # A custom divide factor makes FSDP reduce-scatter with NCCL's PREMUL_SUM. On Leonardo + # (torch 2.10, bf16, reduce-scatter across nodes) this returned all-zero expert gradients. + # Plain SUM followed by a local division is correct and costs only one extra elementwise op. + # NOTE: not for single-rank groups: there, torch would apply the divide factor twice. + experts.set_force_sum_reduction_for_comms(True) + + @staticmethod + def _set_explicit_moe_prefetching(root: FSDP2, fsdp_blocks: list[nn.Module]) -> None: + """Explicit FSDP prefetching for MoE models (as in TorchTitan). + + The device-to-host synchronization in the expert-parallel token dispatch blocks the CPU and + delays FSDP's implicit prefetching, so we prefetch the next block and its experts explicitly. + """ + + def with_experts(block: nn.Module) -> list[nn.Module]: + return [block] + [m for m in block.modules() if getattr(m, "_ep_enabled", False)] + + if len(fsdp_blocks) == 0: + return + root.set_modules_to_forward_prefetch(with_experts(fsdp_blocks[0])) + for block, next_block in zip(fsdp_blocks[:-1], fsdp_blocks[1:]): + block.set_modules_to_forward_prefetch(with_experts(next_block)) + for block, prev_block in zip(fsdp_blocks[1:], fsdp_blocks[:-1]): + block.set_modules_to_backward_prefetch(with_experts(prev_block)) + @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 = [] @@ -286,8 +350,10 @@ def get_ep_wrapped_model( 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 + expert_meshes = get_expert_meshes(device_mesh) + if expert_meshes is None: + raise ValueError("EP wrapping requires a device mesh with expert_parallel_degree > 1.") + ep_mesh = expert_meshes.ep wrapped_blocks = 0 for module in model.modules(): @@ -318,22 +384,22 @@ def get_ep_wrapped_model( f"w2.ndim={experts.w2.ndim}" ) + if experts.w1.shape[0] % ep_mesh.size() != 0: + raise ValueError( + f"Number of experts ({experts.w1.shape[0]}) must be divisible by the " + f"expert parallel degree ({ep_mesh.size()})." + ) + 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) + # NOTE: expert parameters stay in full precision (master weights); the low precision + # compute copies are created by FSDP's mixed precision policy. 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: diff --git a/src/modalities/models/moe/qwen_model.py b/src/modalities/models/moe/qwen_model.py index b20dd05ef..8ed806fa8 100644 --- a/src/modalities/models/moe/qwen_model.py +++ b/src/modalities/models/moe/qwen_model.py @@ -2,9 +2,11 @@ from typing import Literal, Optional, overload import torch +import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from pydantic import BaseModel +from torch.optim.optimizer import register_optimizer_step_post_hook from modalities.models.components.rotary_embedding import ( apply_rotary_pos_emb, @@ -38,9 +40,12 @@ class QwenModelConfig(BaseModel): moe_num_experts: int = 128 moe_top_k: int = 8 moe_d_ff: int = 768 - moe_capacity_factor: float = 1.25 + # None means dropless routing (default). If set, each expert processes at most + # max(ceil(moe_capacity_factor * num_tokens * moe_top_k / moe_num_experts), moe_min_capacity) tokens. + moe_capacity_factor: Optional[float] = None moe_min_capacity: int = 4 moe_overflow_policy: Literal["drop", "residual"] = "residual" + moe_expert_backend: Literal["auto", "grouped_mm", "loop"] = "auto" moe_router_noise_std: float = 0.0 moe_router_temperature: float = 1.0 moe_router_dropout: float = 0.0 @@ -58,7 +63,7 @@ def reset_parameters(self): nn.init.ones_(self.weight) def forward(self, x): - return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight + return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) class RotaryEmbedding(nn.Module): @@ -151,83 +156,140 @@ def forward(self, x, mask=None): return self.o_proj(out.transpose(1, 2).contiguous().view(B, T, -1)) +def _local_tensor(t: torch.Tensor) -> torch.Tensor: + return t.to_local() if DTensor is not None and isinstance(t, DTensor) else t + + class GroupedExperts(nn.Module): + # torch._grouped_mm needs every token group to start at a 16-byte aligned offset along the + # token dimension (including in the backward 2d x 2d grouped GEMM). 16 rows covers bf16 and fp8. + GROUPED_MM_ALIGNMENT = 16 + def __init__( self, - num_experts, - d_model, - d_ff, - ffn_dropout, + num_experts: int, + d_model: int, + d_ff: int, + ffn_dropout: float, + backend: Literal["auto", "grouped_mm", "loop"] = "auto", + init_std: float = 0.02, ): super().__init__() self.num_experts = num_experts self.d_model = d_model self.d_ff = d_ff + self.backend = backend + self.init_std = init_std self.dropout = nn.Dropout(ffn_dropout) if ffn_dropout > 0 else nn.Identity() + # w1: gate projection, w2: up projection, w3: down projection self.w1 = nn.Parameter(torch.empty(self.num_experts, self.d_ff, self.d_model)) self.w2 = nn.Parameter(torch.empty(self.num_experts, self.d_ff, self.d_model)) self.w3 = nn.Parameter(torch.empty(self.num_experts, self.d_model, self.d_ff)) def reset_parameters(self): - nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5)) - nn.init.kaiming_uniform_(self.w2, a=math.sqrt(5)) - nn.init.kaiming_uniform_(self.w3, a=math.sqrt(5)) - - def _forward_local(self, routed_input: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: - w1 = self.w1.to_local() if DTensor is not None and isinstance(self.w1, DTensor) else self.w1 - w2 = self.w2.to_local() if DTensor is not None and isinstance(self.w2, DTensor) else self.w2 - w3 = self.w3.to_local() if DTensor is not None and isinstance(self.w3, DTensor) else self.w3 - # F.linear requires matching dtypes between inputs and weights. Under mixed precision, - # routed_input can be BF16 while local expert weights remain FP32. + # NOTE: kaiming init on a 3D tensor would use fan_in = d_ff * d_model instead of the per-expert fan_in. + for w in (self.w1, self.w2, self.w3): + nn.init.trunc_normal_(w, mean=0.0, std=self.init_std, a=-3 * self.init_std, b=3 * self.init_std) + + def _use_grouped_mm(self, x: torch.Tensor) -> bool: + if self.backend == "grouped_mm": + return True + if self.backend == "loop": + return False + # auto: grouped_mm on CUDA. On sm80 (A100) torch 2.10 uses a fallback path that syncs on the group + # offsets, but it is still 2-4x faster than the per-expert loop (scripts/benchmark_moe_experts.py). + return x.is_cuda and hasattr(torch, "_grouped_mm") + + def forward( + self, + routed_input: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + num_tokens_per_expert_cpu: Optional[list[int]] = None, + ) -> torch.Tensor: + """Runs the (local) experts on expert-major sorted tokens. + + Args: + routed_input (torch.Tensor): Tokens sorted by expert, shape (num_rows, d_model). Rows beyond + sum(num_tokens_per_expert) are tokens dropped by a capacity limit; their output is zero. + num_tokens_per_expert (torch.Tensor): Number of tokens per local expert (on device). + num_tokens_per_expert_cpu (Optional[list[int]]): The same counts on the host, if already available. + Providing them avoids a host synchronization in the loop backend. + """ + w1, w2, w3 = _local_tensor(self.w1), _local_tensor(self.w2), _local_tensor(self.w3) if routed_input.dtype != w1.dtype: - w1 = w1.to(dtype=routed_input.dtype) - w2 = w2.to(dtype=routed_input.dtype) - w3 = w3.to(dtype=routed_input.dtype) - local_num_tokens = ( - num_tokens_per_expert.to_local() - if DTensor is not None and isinstance(num_tokens_per_expert, DTensor) - else num_tokens_per_expert - ) + # without FSDP mixed precision, weights may be fp32 while activations are bf16 + w1, w2, w3 = (w.to(dtype=routed_input.dtype) for w in (w1, w2, w3)) + num_tokens_per_expert = _local_tensor(num_tokens_per_expert) - outputs: list[torch.Tensor] = [] - start = 0 - total_rows = routed_input.shape[0] - - for expert_idx, num_tokens in enumerate(local_num_tokens.tolist()): - requested_tokens = int(num_tokens) - end = start + requested_tokens - local_end = min(end, total_rows) - expert_input = routed_input[start:local_end] - real_tokens = int(expert_input.shape[0]) - - out_parts: list[torch.Tensor] = [] - if real_tokens > 0: - x1 = F.linear(expert_input, w1[expert_idx]) - x2 = F.linear(expert_input, w2[expert_idx]) - out_parts.append(self.dropout(F.linear(F.silu(x1) * x2, w3[expert_idx]))) - - pad = requested_tokens - real_tokens - if pad > 0: - out_parts.append(routed_input.new_zeros((pad, self.d_model))) - - if out_parts: - outputs.append(torch.cat(out_parts, dim=0) if len(out_parts) > 1 else out_parts[0]) - - start = end - - if not outputs: - return routed_input.new_zeros((0, self.d_model)) - - out = torch.cat(outputs, dim=0) - if out.shape[0] < total_rows: - out = torch.cat([out, routed_input.new_zeros((total_rows - out.shape[0], self.d_model))], dim=0) - elif out.shape[0] > total_rows: - out = out[:total_rows] + use_grouped_mm = self._use_grouped_mm(routed_input) + if num_tokens_per_expert_cpu is None and not use_grouped_mm: + num_tokens_per_expert_cpu = num_tokens_per_expert.tolist() # host sync (non-EP loop path only) + + num_rows = routed_input.shape[0] + num_valid_rows = sum(num_tokens_per_expert_cpu) if num_tokens_per_expert_cpu is not None else num_rows + x = routed_input[:num_valid_rows] if num_valid_rows < num_rows else routed_input + + if use_grouped_mm: + out = self._run_grouped_mm(x, num_tokens_per_expert, w1, w2, w3) + else: + out = self._run_loop(x, num_tokens_per_expert_cpu, w1, w2, w3) + + if num_valid_rows < num_rows: + out = torch.cat([out, out.new_zeros(num_rows - num_valid_rows, self.d_model)]) return out - def forward(self, routed_input: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: - return self._forward_local(routed_input, num_tokens_per_expert) + def _run_loop( + self, + x: torch.Tensor, + num_tokens_per_expert_cpu: list[int], + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + ) -> torch.Tensor: + outputs: list[torch.Tensor] = [] + for expert_idx, expert_input in enumerate(torch.split(x, num_tokens_per_expert_cpu, dim=0)): + if expert_input.shape[0] == 0: + continue + h = F.silu(F.linear(expert_input, w1[expert_idx])) * F.linear(expert_input, w2[expert_idx]) + outputs.append(self.dropout(F.linear(h, w3[expert_idx]))) + if len(outputs) == 0: + # Keep the expert weights in the autograd graph so that every rank produces gradients + # (required for collectives on gradients, e.g. FSDP reduce-scatter). + h = F.silu(F.linear(x, w1[0])) * F.linear(x, w2[0]) + return F.linear(h, w3[0]) + return torch.cat(outputs, dim=0) if len(outputs) > 1 else outputs[0] + + def _run_grouped_mm( + self, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + ) -> torch.Tensor: + # Pad every expert group to a multiple of GROUPED_MM_ALIGNMENT rows. The padded length is bounded + # statically, so this needs no host synchronization. Padding rows are zeros and never read back. + align = self.GROUPED_MM_ALIGNMENT + num_rows = x.shape[0] + num_local_experts = num_tokens_per_expert.shape[0] + with torch.no_grad(): + counts = num_tokens_per_expert.to(torch.int64) + padded_counts = (counts + align - 1) // align * align + padded_start = torch.cumsum(padded_counts, dim=0) - padded_counts + start = torch.cumsum(counts, dim=0) - counts + dest = torch.repeat_interleave(padded_start - start, counts, output_size=num_rows) + torch.arange( + num_rows, device=x.device + ) + offs = torch.cumsum(padded_counts, dim=0).to(torch.int32) + padded_len = num_rows + num_local_experts * (align - 1) + x_padded = torch.index_copy(x.new_zeros(padded_len, x.shape[-1]), 0, dest, x) + + h = F.silu(torch._grouped_mm(x_padded, w1.transpose(-2, -1), offs=offs)) * torch._grouped_mm( + x_padded, w2.transpose(-2, -1), offs=offs + ) + out = self.dropout(torch._grouped_mm(h, w3.transpose(-2, -1), offs=offs)) + return out.index_select(0, dest) class MoEBlock(nn.Module): @@ -237,7 +299,7 @@ def __init__( moe_d_ff: int, moe_num_experts: int, moe_top_k: int, - moe_capacity_factor: float, + moe_capacity_factor: Optional[float], moe_min_capacity: int, moe_overflow_policy: str, moe_router_noise_std: float, @@ -246,6 +308,7 @@ def __init__( moe_aux_loss_coef: float, moe_z_loss_coef: float, ffn_dropout: float, + moe_expert_backend: Literal["auto", "grouped_mm", "loop"] = "auto", ): super().__init__() self.num_experts = moe_num_experts @@ -261,74 +324,160 @@ def __init__( self.router = nn.Linear(d_model, self.num_experts, bias=False) self.experts = GroupedExperts( - num_experts=moe_num_experts, d_model=d_model, d_ff=moe_d_ff, ffn_dropout=ffn_dropout + num_experts=moe_num_experts, + d_model=d_model, + d_ff=moe_d_ff, + ffn_dropout=ffn_dropout, + backend=moe_expert_backend, ) - self.last_aux_loss: Optional[torch.Tensor] = None + self.last_tokens_per_expert: Optional[torch.Tensor] = None + self.last_router_prob_mean: Optional[torch.Tensor] = None + self.last_z_loss: Optional[torch.Tensor] = None def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, D = x.shape E = self.num_experts K = self.top_k N = B * T - x_flat = x.view(N, D) + x_flat = x.reshape(N, D) - logits = self.router(self.router_dropout(x_flat).to(self.router.weight.dtype)).float() + # Router in fp32 for numerically stable routing decisions. + logits = F.linear(self.router_dropout(x_flat).float(), self.router.weight.float()) if self.router_noise_std > 0 and self.training: logits = logits + torch.randn_like(logits) * self.router_noise_std logits = logits / self.router_temperature - probs = torch.softmax(logits, dim=-1) - topk_val, topk_idx = torch.topk(probs, k=K, dim=-1) - topk_w = (topk_val / (topk_val.sum(dim=-1, keepdim=True) + 1e-9)).to(x_flat.dtype) - - capacity = max(math.ceil(self.capacity_factor * N / E), self.min_capacity) - - dispatch_mask = F.one_hot(topk_idx, num_classes=E).to(x_flat.dtype) - positions = torch.cumsum(dispatch_mask.sum(dim=1), dim=0) - capacity_mask = (positions <= capacity).to(x_flat.dtype) - final_mask = dispatch_mask * capacity_mask.unsqueeze(1) - - load = final_mask.sum(dim=[0, 1]) - importance = probs.sum(dim=0) - - flat_valid = capacity_mask.gather(1, topk_idx).bool().reshape(-1) - flat_token_ids = torch.arange(N, device=x.device).unsqueeze(1).expand(N, K).reshape(-1)[flat_valid] - flat_expert_ids = topk_idx.reshape(-1)[flat_valid] - flat_weights = topk_w.reshape(-1)[flat_valid] - - if flat_expert_ids.numel() > 0: - sort_idx = torch.argsort(flat_expert_ids) - token_ids_sorted = flat_token_ids[sort_idx] - expert_ids_sorted = flat_expert_ids[sort_idx] - weights_sorted = flat_weights[sort_idx] + scores = torch.softmax(logits, dim=-1) + topk_val, topk_idx = torch.topk(scores, k=K, dim=-1) + topk_w = topk_val / (topk_val.sum(dim=-1, keepdim=True) + 1e-9) - routed_output = self.experts(x_flat[token_ids_sorted], torch.bincount(expert_ids_sorted, minlength=E)) - weighted_output = routed_output * weights_sorted.unsqueeze(-1) - - out = x_flat.new_zeros((N, D)) - out.index_add_(0, token_ids_sorted, weighted_output) - assigned = x_flat.new_zeros((N,)) - assigned.index_add_(0, token_ids_sorted, weights_sorted) + # Everything below is free of host synchronization (no boolean masking, no bincount, no .item()). + with torch.no_grad(): + flat_expert_ids = topk_idx.reshape(-1) + num_tokens_per_expert = torch.zeros(E, dtype=torch.int64, device=x.device).scatter_add_( + 0, flat_expert_ids, torch.ones_like(flat_expert_ids) + ) + # stable sort keeps token order within each expert (earlier tokens win under a capacity limit) + sort_idx = torch.argsort(flat_expert_ids, stable=True) + num_tokens_for_experts = num_tokens_per_expert + token_kept = None + if self.capacity_factor is not None: + sort_idx, num_tokens_for_experts, token_kept = self._apply_capacity( + sort_idx, flat_expert_ids, num_tokens_per_expert, N + ) + token_ids_sorted = sort_idx // K + + routed_input = x_flat.index_select(0, token_ids_sorted) + num_tokens_cpu = None + if self.capacity_factor is not None and not getattr(self.experts, "_ep_enabled", False): + num_tokens_cpu = num_tokens_for_experts.tolist() # host sync; EP dispatch gets these counts for free + if num_tokens_cpu is None: + routed_output = self.experts(routed_input, num_tokens_for_experts) else: - out = x_flat.new_zeros((N, D)) - assigned = x_flat.new_zeros((N,)) - - not_assigned = assigned < 1e-6 - if not_assigned.any() and self.overflow_policy == "residual": - out[not_assigned] = x_flat[not_assigned] - - aux = None - if self.aux_loss_coef > 0: - imp = importance / (importance.sum() + 1e-9) - ld = load / (load.sum() + 1e-9) - aux = self.aux_loss_coef * E * torch.sum(imp * ld) - if self.z_loss_coef > 0: - z_loss = torch.mean(torch.logsumexp(logits, dim=-1) ** 2) - aux = (aux if aux is not None else torch.tensor(0.0, device=x.device)) + self.z_loss_coef * z_loss - - self.last_aux_loss = aux + routed_output = self.experts(routed_input, num_tokens_for_experts, num_tokens_cpu) + + # Un-sort to (token, k) order and combine the top-k expert outputs with the routing weights. + routed_output = torch.index_copy(torch.empty_like(routed_output), 0, sort_idx, routed_output) + out = torch.bmm(topk_w.to(routed_output.dtype).unsqueeze(1), routed_output.view(N, K, D)).squeeze(1) + + if token_kept is not None and self.overflow_policy == "residual": + # tokens that were dropped by all of their top-k experts pass through unchanged + all_dropped = token_kept.view(N, K).sum(dim=-1) == 0 + out = torch.where(all_dropped.unsqueeze(-1), x_flat, out) + + # Router statistics of this micro-batch for the global-batch load balancing loss, which is computed + # in the loss function (see GlobalBatchLoadBalancingLoss): local token counts per expert (f_i is + # aggregated across data parallel ranks and micro-batches there) and the mean routing probability P_i. + self.last_tokens_per_expert = num_tokens_per_expert + self.last_router_prob_mean = scores.mean(dim=0) + self.last_z_loss = ( + self.z_loss_coef * torch.mean(torch.logsumexp(logits, dim=-1) ** 2) if self.z_loss_coef > 0 else None + ) return out.view(B, T, D) + def _apply_capacity( + self, + sort_idx: torch.Tensor, + flat_expert_ids: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + num_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Drops assignments beyond the expert capacity without host synchronization. + + Returns the new sort index (kept assignments grouped by expert first, dropped ones at the end), + the number of kept tokens per expert, and a per-assignment (token-major) kept indicator. + """ + E = self.num_experts + capacity = max(math.ceil(self.capacity_factor * num_tokens * self.top_k / E), self.min_capacity) + expert_ids_sorted = flat_expert_ids[sort_idx] + expert_start = torch.cumsum(num_tokens_per_expert, dim=0) - num_tokens_per_expert + position_in_expert = torch.arange(sort_idx.shape[0], device=sort_idx.device) - expert_start[expert_ids_sorted] + kept_sorted = position_in_expert < capacity + reorder = torch.argsort(expert_ids_sorted + E * (~kept_sorted).to(expert_ids_sorted.dtype), stable=True) + token_kept = torch.empty_like(kept_sorted, dtype=torch.int32).scatter_(0, sort_idx, kept_sorted.to(torch.int32)) + return sort_idx[reorder], num_tokens_per_expert.clamp(max=capacity), token_kept + + +class GlobalBatchLoadBalancingLoss: + """Global-batch load balancing loss (Qiu et al., 2025, https://arxiv.org/abs/2501.11873), as used by Qwen3, + plus the optional router z-loss. + + For every MoE layer: L_lb = coef * E * sum_i f_i * P_i, where + - P_i is the mean routing probability of expert i over the local micro-batch (carries the gradient), + - f_i is the fraction of routed tokens assigned to expert i over the global batch: the token counts + are summed over all data parallel ranks and over the micro-batches of the current optimizer step + (gradient accumulation). Within a step, f_i is the running estimate over the micro-batches seen so far. + As data parallel gradients are averaged, the optimized objective is E * sum_i f_i^global * P_i^global, + i.e. load balance is encouraged over the global batch instead of within every micro-batch. + + The accumulated counts are reset after every optimizer step (global optimizer step post-hook). + Without gradients (evaluation), the local micro-batch statistics are used without communication. + """ + + def __init__(self, model: nn.Module, process_groups: Optional[list[dist.ProcessGroup]] = None): + """ + Args: + model (nn.Module): The model containing the MoE blocks. + process_groups (Optional[list[dist.ProcessGroup]]): Groups over which the token counts are summed, + i.e. the data parallel groups. If None, the default (world) group is used when distributed + training is initialized. + """ + self.model = model + self.process_groups = process_groups + self._accumulated_tokens_per_expert: Optional[torch.Tensor] = None + self._reset_hook = register_optimizer_step_post_hook(lambda optimizer, args, kwargs: self.reset()) + + def reset(self) -> None: + self._accumulated_tokens_per_expert = None + + def _all_reduce(self, tensor: torch.Tensor) -> torch.Tensor: + if not (dist.is_available() and dist.is_initialized()): + return tensor + for group in self.process_groups if self.process_groups is not None else [None]: + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group) + return tensor + + def __call__(self) -> Optional[torch.Tensor]: + blocks = [m for m in self.model.modules() if isinstance(m, MoEBlock) and m.last_tokens_per_expert is not None] + if len(blocks) == 0: + return None + tokens_per_expert = torch.stack([b.last_tokens_per_expert for b in blocks]) # (num_layers, num_experts) + if torch.is_grad_enabled(): + # one collective per micro-batch for all layers + tokens_per_expert = self._all_reduce(tokens_per_expert.clone()) + if self._accumulated_tokens_per_expert is not None: + tokens_per_expert = tokens_per_expert + self._accumulated_tokens_per_expert + self._accumulated_tokens_per_expert = tokens_per_expert + tokens_fraction = tokens_per_expert.float() / tokens_per_expert.sum(dim=-1, keepdim=True).clamp_min(1) + + loss = torch.zeros((), device=tokens_fraction.device) + for block, f in zip(blocks, tokens_fraction): + if block.aux_loss_coef > 0: + loss = loss + block.aux_loss_coef * block.num_experts * torch.sum(f * block.last_router_prob_mean) + if block.last_z_loss is not None: + loss = loss + block.last_z_loss + return loss + class TransformerBlock(nn.Module): def __init__( @@ -345,9 +494,10 @@ def __init__( moe_d_ff: int = 768, moe_num_experts: int = 128, moe_top_k: int = 8, - moe_capacity_factor: float = 1.25, + moe_capacity_factor: Optional[float] = None, moe_min_capacity: int = 4, moe_overflow_policy: str = "residual", + moe_expert_backend: str = "auto", moe_router_noise_std: float = 0.0, moe_router_temperature: float = 1.0, moe_router_dropout: float = 0.0, @@ -374,6 +524,7 @@ def __init__( moe_capacity_factor=moe_capacity_factor, moe_min_capacity=moe_min_capacity, moe_overflow_policy=moe_overflow_policy, + moe_expert_backend=moe_expert_backend, moe_router_noise_std=moe_router_noise_std, moe_router_temperature=moe_router_temperature, moe_router_dropout=moe_router_dropout, @@ -387,10 +538,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.ffn(self.pre_ffn_norm(x)) return x - @property - def aux_loss(self) -> Optional[torch.Tensor]: - return getattr(self.ffn, "last_aux_loss", None) - class QwenModel(NNModel): def __init__( @@ -412,9 +559,10 @@ def __init__( rope_base: float = 1000000.0, moe_num_experts: int = 128, moe_top_k: int = 8, - moe_capacity_factor: float = 1.25, + moe_capacity_factor: Optional[float] = None, moe_min_capacity: int = 4, moe_overflow_policy: str = "residual", + moe_expert_backend: str = "auto", moe_router_noise_std: float = 0.0, moe_router_temperature: float = 1.0, moe_router_dropout: float = 0.0, @@ -450,6 +598,7 @@ def __init__( moe_capacity_factor=moe_capacity_factor, moe_min_capacity=moe_min_capacity, moe_overflow_policy=moe_overflow_policy, + moe_expert_backend=moe_expert_backend, moe_router_noise_std=moe_router_noise_std, moe_router_temperature=moe_router_temperature, moe_router_dropout=moe_router_dropout, diff --git a/src/modalities/models/parallelism/expert_parallelism.py b/src/modalities/models/parallelism/expert_parallelism.py index eaa33937c..458ea27ed 100644 --- a/src/modalities/models/parallelism/expert_parallelism.py +++ b/src/modalities/models/parallelism/expert_parallelism.py @@ -9,57 +9,31 @@ from torch.distributed.tensor import DeviceMesh, Shard, distribute_module, distribute_tensor -def _permute_tokens( - x: Tensor, - num_tokens_per_expert_group: Tensor, - ep_degree: int, - num_local_experts: int, -) -> tuple[tuple, Tensor, Tensor, Tensor]: +def expert_major_permutation(counts: Tensor, num_sources: int, num_local_experts: int, num_rows: int) -> Tensor: """ - Reorder tokens from the post-all-to-all layout to per-local-expert contiguous layout. - - After the all-to-all, received tokens are ordered as: - [e0_from_rank0 tokens, e1_from_rank0 tokens, ..., e0_from_rank1 tokens, ...] - - We reorder to: - [all tokens for local_expert_0, all tokens for local_expert_1, ...] - - Returns (original_shape, x_permuted, permuted_indices, new_num_tokens_per_expert). + Computes, without any host synchronization, the gather index that reorders rows from the + source-major layout produced by the all-to-all + [src0: e0 | e1 | ... , src1: e0 | e1 | ..., ...] + into the expert-major layout expected by the experts + [e0: src0 | src1 | ..., e1: src0 | src1 | ..., ...]. + + Args: + counts (Tensor): Number of rows per (source, local expert), flattened source-major. + num_sources (int): Number of sources (EP degree). + num_local_experts (int): Number of experts on this rank. + num_rows (int): Total number of rows, i.e. counts.sum(). Must be known on the host. + + Returns: + Tensor: Index such that x[index] is in expert-major order. """ - counts = num_tokens_per_expert_group.view(ep_degree, num_local_experts) # (ep_degree, num_local_experts) - - flat_counts = counts.flatten() # length = ep_degree * num_local_experts - - offsets = flat_counts.cumsum(0) - flat_counts - - # build permuted_indices - indices_per_expert: list[Tensor] = [] - for e in range(num_local_experts): - for r in range(ep_degree): - count = int(counts[r, e].item()) - if count > 0: - start = int(offsets[r * num_local_experts + e].item()) - indices_per_expert.append(torch.arange(start, start + count, device=x.device, dtype=torch.long)) - - if indices_per_expert: - permuted_indices = torch.cat(indices_per_expert) - else: - permuted_indices = torch.zeros(0, dtype=torch.long, device=x.device) - - new_num_tokens_per_expert = counts.sum(dim=0) # (num_local_experts,) - original_shape = x.shape - x_permuted = x[permuted_indices] if permuted_indices.numel() > 0 else x.new_zeros((0, x.shape[-1])) - return original_shape, x_permuted, permuted_indices, new_num_tokens_per_expert - - -def _unpermute_tokens(out: Tensor, original_shape: tuple, permuted_indices: Tensor) -> Tensor: - """ - Inverse of _permute_tokens: scatter expert outputs back to the all-to-all layout. - """ - out_unpermuted = out.new_zeros(original_shape) - if permuted_indices.numel() > 0: - out_unpermuted[permuted_indices] = out - return out_unpermuted + counts = counts.view(num_sources, num_local_experts) + flat_counts = counts.reshape(-1) + src_start = (torch.cumsum(flat_counts, dim=0) - flat_counts).view(num_sources, num_local_experts) + seg_counts = counts.t().reshape(-1) + seg_src_start = src_start.t().reshape(-1) + seg_dst_start = torch.cumsum(seg_counts, dim=0) - seg_counts + shift = torch.repeat_interleave(seg_src_start - seg_dst_start, seg_counts, output_size=num_rows) + return shift + torch.arange(num_rows, device=counts.device, dtype=shift.dtype) class ExpertParallel: @@ -69,15 +43,20 @@ class ExpertParallel: Shards GroupedExperts parameters on the expert dimension (Shard(0)) across EP ranks, and wraps forward() with all-to-all token dispatch/combine collectives. + The only host synchronization per forward is a single device-to-host copy of the + (ep_degree x num_local_experts) token count matrix, which is required to issue the + variable-sized token all-to-all. All index computations stay on the device. + Usage: module.experts = ExpertParallel()._apply(module.experts, ep_mesh) """ def __init__(self) -> None: + # Per-call state, written in _token_dispatch and consumed in _token_combine of the same forward call. self.input_splits: list[int] | None = None self.output_splits: list[int] | None = None - self.original_shape: tuple | None = None self.permuted_indices: Tensor | None = None + self.num_dropped_rows: int = 0 def _partition_fn(self, name: str, mod: nn.Module, device_mesh: DeviceMesh) -> None: for param_name, param in mod.named_parameters(recurse=False): @@ -86,47 +65,60 @@ def _partition_fn(self, name: str, mod: nn.Module, device_mesh: DeviceMesh) -> N nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)])), ) - def _token_dispatch(self, mod: nn.Module, inputs: tuple, device_mesh: DeviceMesh) -> tuple[Tensor, Tensor]: - routed_input, num_tokens_per_expert = inputs + def _token_dispatch(self, mod: nn.Module, inputs: tuple, device_mesh: DeviceMesh) -> tuple: + routed_input, num_tokens_per_expert = inputs[0], inputs[1] ep_degree = device_mesh.shape[0] num_local_experts = num_tokens_per_expert.shape[0] // ep_degree + group = device_mesh.get_group() with torch.no_grad(): - num_tokens_per_expert_group = all_to_all_single( - num_tokens_per_expert, None, None, group=device_mesh.get_group() - ) - + num_tokens_per_expert_group = all_to_all_single(num_tokens_per_expert, None, None, group=group) num_tokens_per_expert_group = torch.ops._c10d_functional.wait_tensor(num_tokens_per_expert_group) - input_splits = ( - num_tokens_per_expert.view(ep_degree, -1).sum(dim=1).to(torch.device("cpu"), non_blocking=True) - ) - - output_splits = ( - num_tokens_per_expert_group.view(ep_degree, -1).sum(dim=1).to(torch.device("cpu"), non_blocking=False) + # The single host synchronization of this layer: both count matrices in one copy. + counts_cpu = ( + torch.stack([num_tokens_per_expert, num_tokens_per_expert_group]) + .to("cpu") + .view(2, ep_degree, num_local_experts) ) - self.input_splits = input_splits.tolist() - self.output_splits = output_splits.tolist() - - routed_input = all_to_all_single_autograd( - routed_input, - self.output_splits, - self.input_splits, - device_mesh.get_group(), - ) - - self.original_shape, routed_input, self.permuted_indices, num_tokens_per_expert_group = _permute_tokens( - routed_input, num_tokens_per_expert_group, ep_degree, num_local_experts - ) - return routed_input, num_tokens_per_expert_group + self.input_splits = counts_cpu[0].sum(dim=1).tolist() + self.output_splits = counts_cpu[1].sum(dim=1).tolist() + num_tokens_per_local_expert_cpu = counts_cpu[1].sum(dim=0).tolist() + + # Rows beyond sum(input_splits) are tokens dropped by a capacity limit; they are not dispatched. + num_valid_rows = sum(self.input_splits) + self.num_dropped_rows = routed_input.shape[0] - num_valid_rows + if self.num_dropped_rows > 0: + routed_input = routed_input[:num_valid_rows] + + routed_input = all_to_all_single_autograd(routed_input, self.output_splits, self.input_splits, group) + + num_received_rows = sum(self.output_splits) + if num_local_experts > 1: + with torch.no_grad(): + self.permuted_indices = expert_major_permutation( + num_tokens_per_expert_group, ep_degree, num_local_experts, num_received_rows + ) + routed_input = routed_input.index_select(0, self.permuted_indices) + else: + self.permuted_indices = None + + num_tokens_per_local_expert = num_tokens_per_expert_group.view(ep_degree, num_local_experts).sum(dim=0) + return routed_input, num_tokens_per_local_expert, num_tokens_per_local_expert_cpu def _token_combine(self, mod: nn.Module, routed_output: Tensor, device_mesh: DeviceMesh) -> Tensor: - routed_output = _unpermute_tokens(routed_output, self.original_shape, self.permuted_indices) + if self.permuted_indices is not None: + routed_output = torch.index_copy(torch.empty_like(routed_output), 0, self.permuted_indices, routed_output) routed_output = all_to_all_single_autograd( routed_output, self.input_splits, self.output_splits, device_mesh.get_group(), ) + if self.num_dropped_rows > 0: + routed_output = torch.cat( + [routed_output, routed_output.new_zeros(self.num_dropped_rows, routed_output.shape[-1])] + ) + self.permuted_indices = None return routed_output def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: diff --git a/src/modalities/nn/model_initialization/parameter_name_filters.py b/src/modalities/nn/model_initialization/parameter_name_filters.py index ebc8a23ec..8f5ca6ddf 100644 --- a/src/modalities/nn/model_initialization/parameter_name_filters.py +++ b/src/modalities/nn/model_initialization/parameter_name_filters.py @@ -13,6 +13,7 @@ class WeightInitTypes(Enum): class SupportWeightInitModels(Enum): GPT2 = "gpt2" COCA = "coca" + QWEN_MOE = "qwen_moe" class RegexFilter(BaseModel): @@ -74,4 +75,30 @@ class RegexFilter(BaseModel): WeightInitTypes.SCALED: RegexFilter(weights=[], biases=[]), WeightInitTypes.SCALED_EMBED: RegexFilter(weights=[], biases=[]), }, + # MoE model (modalities.models.moe.qwen_model.QwenModel). The optional + # "_checkpoint_wrapped_module." infix appears when activation checkpointing is applied before initialization. + SupportWeightInitModels.QWEN_MOE: { + WeightInitTypes.PLAIN: RegexFilter( + weights=[ + r"layers\.\d+\.(_checkpoint_wrapped_module\.)?attn\.(q_proj|k_proj|v_proj|o_proj)\.weight", + r"layers\.\d+\.(_checkpoint_wrapped_module\.)?ffn\.router\.weight", + r"layers\.\d+\.(_checkpoint_wrapped_module\.)?ffn\.experts\.(w1|w2|w3)", + r"token_emb\.weight", + r"lm_head\.weight", + ], + ), + # scaled: output projections of the attention and of the experts + WeightInitTypes.SCALED: RegexFilter( + weights=[ + r"layers\.\d+\.(_checkpoint_wrapped_module\.)?attn\.o_proj\.weight", + r"layers\.\d+\.(_checkpoint_wrapped_module\.)?ffn\.experts\.w3", + ] + ), + WeightInitTypes.SCALED_EMBED: RegexFilter( + weights=[ + r"token_emb\.weight", + r"lm_head\.weight", + ] + ), + }, } diff --git a/src/modalities/optimizers/ep_adamw.py b/src/modalities/optimizers/ep_adamw.py deleted file mode 100644 index 006f9faf9..000000000 --- a/src/modalities/optimizers/ep_adamw.py +++ /dev/null @@ -1,154 +0,0 @@ -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.nn import Module -from torch.optim import AdamW, Optimizer - -from modalities.optimizers.optimizer_factory import _build_optimizer_groups_via_weight_decay_split - - -def _get_ep_param_ids(model: Module) -> set: - return {id(p) for m in model.modules() if getattr(m, "_ep_enabled", False) for p in m.parameters(recurse=False)} - - -def _get_dense_optimizer_groups(model, ep_param_ids, weight_decay, weight_decay_groups_excluded): - weight_decay_groups = model.weight_decay_groups - params = { - name: p - for name, p in model.named_parameters() - if p.requires_grad and id(p) not in ep_param_ids and (not isinstance(p, DTensor) or p.to_local().numel() > 0) - } - return _build_optimizer_groups_via_weight_decay_split( - weight_decay, weight_decay_groups_excluded, weight_decay_groups, params - ) - - -class EPAdamW(Optimizer): - """ - ZeRO stage-1 for EP (DTensor) params + standard AdamW for dense params. - - Each dp_shard rank stores optimizer states for 1/dp_shard of the EP params. - After each step, updated EP param values are broadcast from owner to all ranks. - Dense params are handled by a separate AdamW (FSDP2 shards them independently). - """ - - def __init__( - self, - model: Module, - device_mesh, - lr: float, - betas: tuple[float, float], - eps: float, - weight_decay: float, - weight_decay_groups_excluded: list[str], - ): - self._dp_mesh = device_mesh["dp_shard"] - self._dp_group = self._dp_mesh.get_group() - self._dp_rank = dist.get_rank(self._dp_group) - self._dp_size = dist.get_world_size(self._dp_group) - - ep_param_ids = _get_ep_param_ids(model) - self._all_ep_params = [p for p in model.parameters() if id(p) in ep_param_ids] - - # rank r owns params[r::dp_size] - self._owned_ep_params = self._all_ep_params[self._dp_rank :: self._dp_size] - - dense_groups = _get_dense_optimizer_groups(model, ep_param_ids, weight_decay, weight_decay_groups_excluded) - - if self._owned_ep_params: - self._ep_adamw = AdamW(self._owned_ep_params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) - else: - self._ep_adamw = None - self._dense_adamw = AdamW(dense_groups, lr=lr, betas=betas, eps=eps) - - # unified param groups for lr_scheduler compatibility: - # group 0 = all EP params, groups 1+ = dense weight-decay split - ep_group = {"params": self._all_ep_params, "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay} - all_groups = [ep_group] + [{**g, "lr": lr, "betas": betas, "eps": eps} for g in dense_groups] - super().__init__(all_groups, {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay}) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # all-reduce - for p in self._all_ep_params: - if p.grad is None: - continue - if isinstance(p.grad, DTensor): - local_g = p.grad.to_local() - dist.all_reduce(local_g, op=dist.ReduceOp.SUM, group=self._dp_group) - local_g.div_(self._dp_size) - else: - dist.all_reduce(p.grad, op=dist.ReduceOp.SUM, group=self._dp_group) - p.grad.div_(self._dp_size) - - # Sync lr - if self._ep_adamw is not None: - self._ep_adamw.param_groups[0]["lr"] = self.param_groups[0]["lr"] - for i, group in enumerate(self._dense_adamw.param_groups): - group["lr"] = self.param_groups[i + 1]["lr"] - - # Update ep params - if self._ep_adamw is not None: - self._ep_adamw.step() - - # Update dense params - self._dense_adamw.step() - - # broadcast updated EP param local tensors - for i, p in enumerate(self._all_ep_params): - owner_local_rank = i % self._dp_size - owner_global_rank = dist.get_global_rank(self._dp_group, owner_local_rank) - if isinstance(p, DTensor): - local_tensor = p.to_local() - elif isinstance(p.data, DTensor): - local_tensor = p.data.to_local() - else: - local_tensor = p.data - dist.broadcast(local_tensor, src=owner_global_rank, group=self._dp_group) - - return loss - - def zero_grad(self, set_to_none: bool = True): - for p in self._all_ep_params: - if set_to_none: - p.grad = None - elif p.grad is not None: - p.grad.detach_() - p.grad.zero_() - self._dense_adamw.zero_grad(set_to_none=set_to_none) - - def state_dict(self) -> dict: - return { - "ep_adamw": self._ep_adamw.state_dict() if self._ep_adamw is not None else {}, - "dense_adamw": self._dense_adamw.state_dict(), - } - - def load_state_dict(self, state_dict: dict) -> None: - if self._ep_adamw is not None and state_dict["ep_adamw"]: - self._ep_adamw.load_state_dict(state_dict["ep_adamw"]) - self._dense_adamw.load_state_dict(state_dict["dense_adamw"]) - - -def get_ep_adam_w( - wrapped_model, - device_mesh, - lr: float, - betas: tuple[float, float], - eps: float, - weight_decay: float, - weight_decay_groups_excluded: list[str], -) -> EPAdamW: - return EPAdamW( - model=wrapped_model, - device_mesh=device_mesh, - lr=lr, - betas=betas, - eps=eps, - weight_decay=weight_decay, - weight_decay_groups_excluded=weight_decay_groups_excluded, - ) diff --git a/src/modalities/optimizers/optimizer_factory.py b/src/modalities/optimizers/optimizer_factory.py index 2eebe89e1..71d7c345f 100644 --- a/src/modalities/optimizers/optimizer_factory.py +++ b/src/modalities/optimizers/optimizer_factory.py @@ -92,11 +92,40 @@ def get_optimizer_groups(model: FSDP, weight_decay: float, weight_decay_groups_e model, weight_decay, weight_decay_groups_excluded ) + optimizer_groups, optimizer_groups_names = _split_optimizer_groups_by_device_mesh( + optimizer_groups, optimizer_groups_names + ) _assert_completeness_of_optimizer_groups(model, optimizer_groups) _print_optimizer_groups_overview(optimizer_groups, optimizer_groups_names) return optimizer_groups +def _split_optimizer_groups_by_device_mesh( + optimizer_groups: OptimizerGroups, optimizer_groups_names: list[str] +) -> tuple[OptimizerGroups, list[str]]: + """ + Splits optimizer groups such that each group only contains DTensors on the same device mesh. + The foreach/fused optimizer implementations cannot operate on DTensors from different meshes + within one group (e.g., expert-parallel experts and FSDP-sharded dense parameters in MoE models). + For models whose parameters all live on one mesh, the groups are returned unchanged. + """ + split_groups: OptimizerGroups = [] + split_names: list[str] = [] + for group, name in zip(optimizer_groups, optimizer_groups_names): + params_by_mesh: dict[object, list[nn.Parameter]] = {} + for param in group["params"]: + mesh = param.device_mesh if isinstance(param, DTensor) else None + params_by_mesh.setdefault(mesh, []).append(param) + if len(params_by_mesh) <= 1: + split_groups.append(group) + split_names.append(name) + continue + for i, params in enumerate(params_by_mesh.values()): + split_groups.append({**group, "params": params}) + split_names.append(f"{name}_mesh{i}") + return split_groups, split_names + + def _check_existence_of_weight_decay_groups_excluded(model: nn.Module, weight_decay_groups_excluded: list[str]) -> None: """ checks the existence of all groups diff --git a/src/modalities/registry/components.py b/src/modalities/registry/components.py index 9f5fda81f..8dc5622df 100644 --- a/src/modalities/registry/components.py +++ b/src/modalities/registry/components.py @@ -36,7 +36,6 @@ DummyLRSchedulerConfig, DummyProgressSubscriberConfig, DummyResultSubscriberConfig, - EPAdamWConfig, EPWrappedModelConfig, EvaluationResultToDiscSubscriberConfig, FSDP1ActivationCheckpointedModelConfig, @@ -54,6 +53,7 @@ LLMDataLoaderConfig, MemMapDatasetConfig, MoECrossEntropyLossConfig, + MoEMFUCalculatorConfig, OneCycleLRSchedulerConfig, PackedMemMapDatasetContinuousConfig, PackedMemMapDatasetMegatronConfig, @@ -113,7 +113,6 @@ ComposedInitializationRoutines, ComposedModelInitializationConfig, ) -from modalities.optimizers.ep_adamw import get_ep_adam_w from modalities.optimizers.lr_schedulers import DummyLRScheduler, LRSchedulerFactory from modalities.optimizers.optimizer_factory import OptimizerFactory from modalities.optimizers.optimizer_list import OptimizersList @@ -136,7 +135,7 @@ from modalities.utils.debug_components import Debugging, HookRegistration from modalities.utils.debugging_configs import DebuggingConfig, NaNHookConfig, PrintForwardHookConfig from modalities.utils.maybe_list_parameter import MaybeListDecorator, maybe_list_parameter -from modalities.utils.mfu import GPT2MFUCalculator +from modalities.utils.mfu import GPT2MFUCalculator, MoEMFUCalculator from modalities.utils.number_conversion import ( LocalNumBatchesFromNumSamplesConfig, LocalNumBatchesFromNumTokensConfig, @@ -266,7 +265,6 @@ class ComponentEntity: ComponentEntity( "optimizer", "adam_w", maybe_model_list_for_optimizer(OptimizerFactory.get_adam_w), AdamWOptimizerConfig ), - ComponentEntity("optimizer", "ep_adam_w", maybe_model_list_for_optimizer(get_ep_adam_w), EPAdamWConfig), ComponentEntity( "optimizer", "fsdp1_checkpointed", @@ -418,6 +416,7 @@ class ComponentEntity: ), # MFU calculators ComponentEntity("mfu_calculator", "gpt2", GPT2MFUCalculator, GPT2MFUCalculatorConfig), + ComponentEntity("mfu_calculator", "moe", MoEMFUCalculator, MoEMFUCalculatorConfig), # Number conversion ComponentEntity( "number_conversion", diff --git a/src/modalities/running_env/fsdp/device_mesh.py b/src/modalities/running_env/fsdp/device_mesh.py index f4f3b7e26..cbaabec4c 100644 --- a/src/modalities/running_env/fsdp/device_mesh.py +++ b/src/modalities/running_env/fsdp/device_mesh.py @@ -52,7 +52,6 @@ def _validate(self): self.data_parallel_shard_degree = self.world_size // ( self.data_parallel_replicate_degree * self.context_parallel_degree - * self.expert_parallel_degree * self.tensor_parallel_degree * self.pipeline_parallel_degree ) @@ -61,14 +60,12 @@ def _validate(self): self.data_parallel_replicate_degree = self.world_size // ( self.data_parallel_shard_degree * self.context_parallel_degree - * self.expert_parallel_degree * self.tensor_parallel_degree * self.pipeline_parallel_degree ) if ( self.data_parallel_shard_degree * self.data_parallel_replicate_degree - * self.expert_parallel_degree * self.tensor_parallel_degree * self.pipeline_parallel_degree * self.context_parallel_degree @@ -77,13 +74,29 @@ def _validate(self): raise ConfigError( f"Invalid parallel dims: data_parallel_shard_degree({self.data_parallel_shard_degree}) * " f"data_parallel_replicate_degree({self.data_parallel_replicate_degree}) * " - f"expert_parallel_degree({self.expert_parallel_degree}) * " f"tensor_parallel_degree({self.tensor_parallel_degree}) *" f"* pipeline_parallel_degree({self.pipeline_parallel_degree}) *" f"context_parallel_degree({self.context_parallel_degree})!= WORLD_SIZE({self.world_size})" ) if self.enable_loss_parallel and self.tensor_parallel_degree <= 1: raise ConfigError(f"{self.enable_loss_parallel=} requires tensor_parallel_degree > 1") + if self.expert_parallel_degree > 1: + # Expert parallelism is folded into data parallel sharding: + # the dp_shard ranks act as data parallel ranks for dense layers and are split into + # (dp_shard // ep) x ep for the experts. + if self.data_parallel_shard_degree % self.expert_parallel_degree != 0: + raise ConfigError( + f"expert_parallel_degree({self.expert_parallel_degree}) must divide " + f"data_parallel_shard_degree({self.data_parallel_shard_degree})." + ) + unsupported = { + "tensor_parallel_degree": self.tensor_parallel_degree, + "pipeline_parallel_degree": self.pipeline_parallel_degree, + "context_parallel_degree": self.context_parallel_degree, + } + unsupported = {k: v for k, v in unsupported.items() if v > 1} + if unsupported: + raise ConfigError(f"Expert parallelism is not yet supported in combination with {unsupported}.") return self @@ -92,6 +105,7 @@ class ParallelismDegrees(Enum): DP_SHARD = "dp_shard" CP = "cp" EP = "ep" + EFSDP = "efsdp" TP = "tp" PP = "pp" @@ -132,7 +146,6 @@ def get_device_mesh( data_parallel_replicate_degree, data_parallel_shard_degree, context_parallel_degree, - expert_parallel_degree, tensor_parallel_degree, ], [ @@ -140,7 +153,6 @@ def get_device_mesh( ParallelismDegrees.DP_REPLICATE.value, ParallelismDegrees.DP_SHARD.value, ParallelismDegrees.CP.value, - ParallelismDegrees.EP.value, ParallelismDegrees.TP.value, ], strict=True, @@ -151,11 +163,66 @@ def get_device_mesh( names = tuple(names) device_mesh = init_device_mesh(device_type, dims, mesh_dim_names=names) logger.info(f"{device_mesh=} | {world_size=} | {enable_loss_parallel=}") + if expert_parallel_degree > 1: + _init_expert_meshes(device_mesh, expert_parallel_degree) # TODO: Torch Titan had some more checks here. We need to check if we also need those: # https://github.com/pytorch/torchtitan/blob/b291ad662493b63d25b038a30a915082d3617baf/torchtitan/distributed/parallel_dims.py#L86-L104 return device_mesh +class ExpertMeshes: + """Device meshes used for the experts of MoE layers when expert parallelism is enabled. + + Expert parallelism is folded into data parallel sharding: dp_shard is split into + (efsdp = dp_shard // ep, ep), with ep as the innermost dimension so that EP groups + consist of consecutive ranks (i.e., stay within a node for ep <= GPUs per node). + + Attributes: + ep (DeviceMesh): 1D mesh over which experts are distributed (Shard(0) on the expert dim). + expert_fsdp (DeviceMesh): Mesh over which each expert shard is additionally FSDP-sharded + (1D "efsdp", or 2D ("dp_replicate", "efsdp") for HSDP). + gradient_divide_factor (int): Number of data parallel ranks contributing tokens to each expert, + i.e. the factor that expert gradients have to be divided by to match the dense gradients. + """ + + def __init__(self, ep: DeviceMesh, expert_fsdp: DeviceMesh, gradient_divide_factor: int): + self.ep = ep + self.expert_fsdp = expert_fsdp + self.gradient_divide_factor = gradient_divide_factor + + +_EXPERT_MESHES_ATTR = "_modalities_expert_meshes" + + +def _init_expert_meshes(device_mesh: DeviceMesh, expert_parallel_degree: int) -> None: + dp_shard_degree = device_mesh[ParallelismDegrees.DP_SHARD.value].size() + efsdp_degree = dp_shard_degree // expert_parallel_degree + # NOTE: pass the dim as an index; torch 2.10's _unflatten rejects valid dim names (inverted membership check). + sparse_mesh = device_mesh._unflatten( + device_mesh.mesh_dim_names.index(ParallelismDegrees.DP_SHARD.value), + (efsdp_degree, expert_parallel_degree), + (ParallelismDegrees.EFSDP.value, ParallelismDegrees.EP.value), + ) + if has_parallelism_method(device_mesh, ParallelismDegrees.DP_REPLICATE): + expert_fsdp = sparse_mesh[(ParallelismDegrees.DP_REPLICATE.value, ParallelismDegrees.EFSDP.value)] + dp_replicate_degree = device_mesh[ParallelismDegrees.DP_REPLICATE.value].size() + else: + expert_fsdp = sparse_mesh[ParallelismDegrees.EFSDP.value] + dp_replicate_degree = 1 + expert_meshes = ExpertMeshes( + ep=sparse_mesh[ParallelismDegrees.EP.value], + expert_fsdp=expert_fsdp, + gradient_divide_factor=dp_replicate_degree * dp_shard_degree, + ) + setattr(device_mesh, _EXPERT_MESHES_ATTR, expert_meshes) + logger.info(f"{sparse_mesh=} | {expert_parallel_degree=} | {efsdp_degree=}") + + +def get_expert_meshes(device_mesh: DeviceMesh) -> ExpertMeshes | None: + """Returns the expert meshes if expert parallelism is enabled for the given device mesh, else None.""" + return getattr(device_mesh, _EXPERT_MESHES_ATTR, None) + + def get_parallel_degree(device_mesh: DeviceMesh, parallelism_methods: list[ParallelismDegrees]) -> int: """Gets the number of parallel ranks (i.e., the parallelism degree) from the device mesh for a specific parallelism method. diff --git a/src/modalities/training/activation_checkpointing/activation_checkpointing.py b/src/modalities/training/activation_checkpointing/activation_checkpointing.py index 0c194c350..a52227dd7 100644 --- a/src/modalities/training/activation_checkpointing/activation_checkpointing.py +++ b/src/modalities/training/activation_checkpointing/activation_checkpointing.py @@ -76,6 +76,8 @@ class ActivationCheckpointing: "ops.aten._scaled_dot_product_efficient_attention.default": ops.aten._scaled_dot_product_efficient_attention.default, # noqa "ops.aten._scaled_dot_product_flash_attention.default": ops.aten._scaled_dot_product_flash_attention.default, "ops._c10d_functional.reduce_scatter_tensor.default": ops._c10d_functional.reduce_scatter_tensor.default, + "ops._c10d_functional.all_to_all_single.default": ops._c10d_functional.all_to_all_single.default, + "ops.aten._grouped_mm.default": ops.aten._grouped_mm.default, # for low precision training, it's useful to always save # the result of max, since the absolute maximum is # used to compute the scaling factor for quantization. diff --git a/src/modalities/training/gradient_clipping/ep_gradient_clipper.py b/src/modalities/training/gradient_clipping/ep_gradient_clipper.py index 2efc5ed58..175d241db 100644 --- a/src/modalities/training/gradient_clipping/ep_gradient_clipper.py +++ b/src/modalities/training/gradient_clipping/ep_gradient_clipper.py @@ -16,7 +16,13 @@ class EPGradientClipper(FSDP2GradientClipper): - """FSDP2 clipper wrapper that handles EP DTensor gradients safely.""" + """FSDP2 gradient clipper for models whose parameters live on different device meshes. + + With expert parallelism, the dense parameters are DTensors on the data parallel mesh, while the expert + parameters are DTensors on the (expert FSDP x EP) mesh. Foreach ops cannot mix DTensors of different + meshes, so the gradients are grouped by mesh (as in TorchTitan's EP-aware clip_grad_norm_). + The norm of each group is reduced with a single collective, and the group norms are combined locally. + """ def __init__( self, @@ -38,52 +44,49 @@ def __init__( @torch.no_grad() def clip_gradients(self) -> torch.Tensor: - grads = [p.grad for model in self.models for p in model.parameters() if p.grad is not None] + params_by_mesh: dict[Optional[DeviceMesh], list[torch.nn.Parameter]] = {} + for model in self.models: + for p in model.parameters(): + if p.grad is None: + continue + mesh = p.grad.device_mesh if isinstance(p.grad, DTensor) else None + params_by_mesh.setdefault(mesh, []).append(p) - if len(grads) == 0: - device = ( - torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") + norm_type = self.norm_type.value + group_norms: list[torch.Tensor] = [] + for params in params_by_mesh.values(): + group_norm = torch.nn.utils.get_total_norm( + tensors=[p.grad for p in params], norm_type=norm_type, foreach=self.foreach ) + if isinstance(group_norm, DTensor): + # one collective per mesh (reduces the partial norm over all shards) + group_norm = group_norm.full_tensor() + group_norms.append(group_norm) + + if len(group_norms) == 0: + device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else "cpu" total_norm = torch.tensor(0.0, device=device) + elif math.isinf(norm_type): + total_norm = torch.max(torch.stack(group_norms)) else: - norm_type_val = self.norm_type.value - first_grad = grads[0] - first_device = first_grad.to_local().device if isinstance(first_grad, DTensor) else first_grad.device - norm_scalars: list[torch.Tensor] = [] - - for grad in grads: - grad_norm = torch.linalg.vector_norm(grad, ord=norm_type_val) - if isinstance(grad_norm, DTensor): - grad_norm = grad_norm.full_tensor() - norm_scalars.append(grad_norm.to(first_device)) - - if math.isinf(norm_type_val): - total_norm = torch.max(torch.stack(norm_scalars)) - else: - total_norm = torch.linalg.vector_norm(torch.stack(norm_scalars), ord=norm_type_val) - - if self.error_if_nonfinite and (torch.isnan(total_norm) or torch.isinf(total_norm)): - raise RuntimeError( - f"The total norm of order {norm_type_val} for gradients is non-finite: {total_norm.item()}" - ) + total_norm = torch.linalg.vector_norm(torch.stack(group_norms), ord=norm_type) if has_parallelism_method(self.device_mesh, ParallelismDegrees.PP): pp_mesh = get_mesh_for_parallelism_method( device_mesh=self.device_mesh, parallelism_method=ParallelismDegrees.PP ) - if math.isinf(self.norm_type.value): + if math.isinf(norm_type): dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=pp_mesh.get_group()) else: - total_norm **= self.norm_type.value + total_norm **= norm_type dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=pp_mesh.get_group()) - total_norm **= 1.0 / self.norm_type.value + total_norm **= 1.0 / norm_type - # do not use torch.nn.utils.clip_grads_with_norm_ here: it batches grads with - # torch._foreach_mul_, which fails when the list mixes DTensors from different meshes. - clip_coef = self.max_norm / (total_norm + 1e-6) - clip_coef_clamped = torch.clamp(clip_coef, max=1.0) + if self.error_if_nonfinite and not torch.isfinite(total_norm): + raise RuntimeError(f"The total norm of order {norm_type} for gradients is non-finite: {total_norm.item()}") - for grad in grads: - grad_device = grad.to_local().device if isinstance(grad, DTensor) else grad.device - grad.mul_(clip_coef_clamped.to(grad_device)) + for params in params_by_mesh.values(): + torch.nn.utils.clip_grads_with_norm_( + parameters=params, max_norm=self.max_norm, total_norm=total_norm, foreach=self.foreach + ) return total_norm diff --git a/src/modalities/utils/mfu.py b/src/modalities/utils/mfu.py index 434dd714e..d893fa183 100644 --- a/src/modalities/utils/mfu.py +++ b/src/modalities/utils/mfu.py @@ -195,3 +195,39 @@ def compute(self, num_samples_per_second: torch.Tensor) -> torch.Tensor: theoretical_flops_per_token=self._theoretical_flops_per_token, theoretical_gpu_peak_performance=self._theoretical_gpu_peak_performance, ) + + +class MoEMFUCalculator(GPT2MFUCalculator): + """ + MFU calculator for MoE transformers. Only the parameters that are active per token are counted, + i.e. all dense parameters plus top_k of the num_experts SwiGLU experts in every layer. + """ + + def __init__( + self, + n_layer: int, + sequence_length: int, + n_embd: int, + world_size: int, + model_parts: FSDPX | list[FSDP2], + num_experts: int, + top_k: int, + moe_d_ff: int, + device_mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None, + ): + super().__init__( + n_layer=n_layer, + sequence_length=sequence_length, + n_embd=n_embd, + world_size=world_size, + model_parts=model_parts, + device_mesh=device_mesh, + ) + inactive_expert_params = n_layer * 3 * n_embd * moe_d_ff * (num_experts - top_k) + self._num_params = self._num_params - inactive_expert_params + self._theoretical_flops_per_token = GPT2MFUCalculator._get_theoretical_flops_per_token( + num_params=self._num_params, + n_layer=self._n_layer, + sequence_length=self._sequence_length, + n_embd=self._n_embd, + ) diff --git a/tests/end2end_tests/test_moe_ep_fsdp2_e2e.py b/tests/end2end_tests/test_moe_ep_fsdp2_e2e.py index 70e7203e9..52ba25af0 100644 --- a/tests/end2end_tests/test_moe_ep_fsdp2_e2e.py +++ b/tests/end2end_tests/test_moe_ep_fsdp2_e2e.py @@ -7,7 +7,9 @@ import pytest import torch +import torch.distributed as dist import torch.multiprocessing as mp +from torch.distributed.tensor import DTensor from modalities.__main__ import Main, load_app_config_dict from modalities.batch import EvaluationResultBatch @@ -46,6 +48,20 @@ def _patch_for_short_test_run(config_dict: dict[str, Any], checkpoint_root_path: config_dict["train_dataloader"]["config"]["pin_memory"] = False config_dict["test_dataloader"]["config"]["pin_memory"] = False + # tiny model, independent of the model size in the training config + config_dict["model_raw"]["config"].update( + d_model=128, + n_heads=8, + n_kv_heads=4, + num_layers=2, + d_ff=128, + moe_num_experts=8, + moe_top_k=2, + moe_d_ff=128, + ) + + # the training config disables checkpointing (k: 0); this test checks saving with the tiny model + config_dict["checkpoint_saving"]["config"]["checkpoint_saving_strategy"]["config"]["k"] = -1 config_dict["settings"]["paths"]["checkpoint_saving_path"] = checkpoint_root_path config_dict["checkpoint_saving"]["config"]["checkpoint_saving_execution"]["config"][ "checkpoint_path" @@ -118,6 +134,16 @@ def _worker_impl(process_id: int, config_file_path: Path, tmp_path: Path) -> Non loss_value = message.payload.losses["train loss avg"].value assert torch.isfinite(loss_value), f"Found non-finite train loss: {loss_value}" + # Every rank must hold identical (unsharded) dense and expert weights after training. + # This fails if gradients of some parameters are not synchronized across all data parallel ranks. + world_size = dist.get_world_size() + for name, param in components.model_raw.named_parameters(): + full = param.detach().full_tensor() if isinstance(param, DTensor) else param.detach() + gathered = [torch.empty_like(full) for _ in range(world_size)] + dist.all_gather(gathered, full.contiguous()) + for other in gathered[1:]: + torch.testing.assert_close(other, gathered[0], msg=f"{name} differs across ranks") + if process_id == 0: checkpoint_info_file_path = checkpoint_root_path / "last_checkpoint_info.json" assert checkpoint_info_file_path.exists(), "Expected checkpoint info file from DCP save." diff --git a/tests/models/moe/test_loss_functions.py b/tests/models/moe/test_loss_functions.py index 5149aef1b..2dcbddfa8 100644 --- a/tests/models/moe/test_loss_functions.py +++ b/tests/models/moe/test_loss_functions.py @@ -3,57 +3,106 @@ from modalities.batch import InferenceResultBatch from modalities.loss_functions import MoECrossEntropyLoss - - -class DummyLayer: - def __init__(self, aux_loss): - self.aux_loss = aux_loss - - -class DummyModel: - def __init__(self, aux_losses: list[torch.Tensor | None]): - self.layers = {str(i): DummyLayer(aux) for i, aux in enumerate(aux_losses)} - - -def test_moe_cross_entropy_loss_adds_aux_losses(): - logits = torch.tensor( - [ - [[1.2, 0.3, -0.5], [0.1, 1.8, -0.3]], - [[0.5, -0.4, 1.1], [0.7, 0.2, -0.1]], - ], - dtype=torch.float32, +from modalities.models.moe.qwen_model import MoEBlock, QwenModel + +VOCAB, NUM_EXPERTS, TOP_K, AUX_COEF, Z_COEF = 32, 4, 2, 0.01, 0.001 + + +def _build_model() -> QwenModel: + torch.manual_seed(0) + model = QwenModel( + vocab_size=VOCAB, + max_seq_len=16, + d_model=16, + n_heads=4, + n_kv_heads=2, + d_ff=32, + num_layers=2, + moe_d_ff=24, + moe_num_experts=NUM_EXPERTS, + moe_top_k=TOP_K, + moe_aux_loss_coef=AUX_COEF, + moe_z_loss_coef=Z_COEF, + moe_expert_backend="loop", ) - targets = torch.tensor([[0, 1], [2, 0]], dtype=torch.long) + # expert weights are allocated uninitialized (initialized later by the model initializer in training) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + return model + + +def _forward(model: QwenModel, seed: int) -> tuple[InferenceResultBatch, torch.Tensor, list[MoEBlock]]: + input_ids = torch.randint(0, VOCAB, (2, 6), generator=torch.Generator().manual_seed(seed)) + logits = model(input_ids) + targets = torch.roll(input_ids, -1, dims=1) + batch = InferenceResultBatch(targets={"targets": targets}, predictions={"logits": logits}) + ce = CrossEntropyLoss()(logits.view(-1, VOCAB), targets.view(-1)) + blocks = [m for m in model.modules() if isinstance(m, MoEBlock)] + return batch, ce, blocks + + +def _expected_moe_loss(blocks: list[MoEBlock], tokens_per_expert: torch.Tensor) -> torch.Tensor: + f = tokens_per_expert.float() / tokens_per_expert.sum(dim=-1, keepdim=True) + expected = torch.zeros(()) + for block, f_layer in zip(blocks, f): + expected = expected + AUX_COEF * NUM_EXPERTS * torch.sum(f_layer * block.last_router_prob_mean) + expected = expected + block.last_z_loss + return expected + + +def test_moe_loss_without_router_statistics_matches_plain_ce(): + model = _build_model() + loss_fn = MoECrossEntropyLoss(target_key="targets", prediction_key="logits", model=model) + logits = torch.randn(2, 3, VOCAB) + targets = torch.randint(0, VOCAB, (2, 3)) + batch = InferenceResultBatch(targets={"targets": targets}, predictions={"logits": logits}) + + expected = CrossEntropyLoss()(logits.view(-1, VOCAB), targets.view(-1)) + assert torch.allclose(loss_fn(batch), expected) - batch = InferenceResultBatch( - targets={"targets": targets}, - predictions={"logits": logits}, - ) - aux_1 = torch.tensor(0.2) - aux_2 = torch.tensor(0.3) - model = DummyModel(aux_losses=[aux_1, None, aux_2]) +def test_moe_loss_uses_token_counts_accumulated_over_the_global_batch(): + model = _build_model() loss_fn = MoECrossEntropyLoss(target_key="targets", prediction_key="logits", model=model) + optimizer = torch.optim.SGD(model.parameters(), lr=0.0) + + # micro-batch 1: only its own token counts + batch, ce, blocks = _forward(model, seed=1) + counts_1 = torch.stack([b.last_tokens_per_expert for b in blocks]) + torch.testing.assert_close(loss_fn(batch), ce + _expected_moe_loss(blocks, counts_1)) + # micro-batch 2 of the same optimizer step: counts of micro-batches 1 and 2 + batch, ce, blocks = _forward(model, seed=2) + counts_2 = torch.stack([b.last_tokens_per_expert for b in blocks]) loss = loss_fn(batch) - base_ce = CrossEntropyLoss(reduction="mean")(logits.view(-1, logits.size(-1)), targets.view(-1)) + torch.testing.assert_close(loss, ce + _expected_moe_loss(blocks, counts_1 + counts_2)) - assert torch.allclose(loss, base_ce + aux_1 + aux_2) + # the load balancing loss is differentiable w.r.t. the router + loss.backward() + assert all(b.router.weight.grad is not None and b.router.weight.grad.abs().sum() > 0 for b in blocks) + # an optimizer step starts a new global batch + optimizer.step() + batch, ce, blocks = _forward(model, seed=3) + counts_3 = torch.stack([b.last_tokens_per_expert for b in blocks]) + torch.testing.assert_close(loss_fn(batch), ce + _expected_moe_loss(blocks, counts_3)) -def test_moe_cross_entropy_loss_without_aux_matches_plain_ce(): - logits = torch.randn(2, 3, 5) - targets = torch.randint(0, 5, (2, 3), dtype=torch.long) - batch = InferenceResultBatch( - targets={"labels": targets}, - predictions={"pred": logits}, - ) +def test_moe_loss_in_evaluation_uses_local_counts_and_keeps_training_state(): + model = _build_model() + loss_fn = MoECrossEntropyLoss(target_key="targets", prediction_key="logits", model=model) - model = DummyModel(aux_losses=[None, None]) - loss_fn = MoECrossEntropyLoss(target_key="labels", prediction_key="pred", model=model) + batch, _, blocks = _forward(model, seed=1) + counts_train = torch.stack([b.last_tokens_per_expert for b in blocks]) + loss_fn(batch) - loss = loss_fn(batch) - expected = CrossEntropyLoss(reduction="mean")(logits.view(-1, logits.size(-1)), targets.view(-1)) + with torch.no_grad(): + batch, ce, blocks = _forward(model, seed=2) + counts_eval = torch.stack([b.last_tokens_per_expert for b in blocks]) + torch.testing.assert_close(loss_fn(batch), ce + _expected_moe_loss(blocks, counts_eval)) - assert torch.allclose(loss, expected) + # evaluation did not modify the accumulated training statistics + batch, ce, blocks = _forward(model, seed=3) + counts_3 = torch.stack([b.last_tokens_per_expert for b in blocks]) + torch.testing.assert_close(loss_fn(batch), ce + _expected_moe_loss(blocks, counts_train + counts_3)) diff --git a/tests/models/moe/test_qwen_model.py b/tests/models/moe/test_qwen_model.py index d4d90b592..68fd789b6 100644 --- a/tests/models/moe/test_qwen_model.py +++ b/tests/models/moe/test_qwen_model.py @@ -1,10 +1,11 @@ +import pytest import torch from modalities.models.moe.qwen_model import GroupedExperts, QwenModel def _build_tiny_qwen_model() -> QwenModel: - return QwenModel( + model = QwenModel( vocab_size=32, max_seq_len=16, d_model=16, @@ -15,12 +16,16 @@ def _build_tiny_qwen_model() -> QwenModel: moe_d_ff=24, moe_num_experts=4, moe_top_k=2, - moe_capacity_factor=1.25, moe_min_capacity=1, moe_overflow_policy="residual", moe_aux_loss_coef=0.01, moe_z_loss_coef=0.0, ) + # expert weights are allocated uninitialized (initialized later by the model initializer in training) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + return model def test_qwen_model_forward_dict_output_shape(): @@ -35,26 +40,64 @@ def test_qwen_model_forward_dict_output_shape(): assert output["logits"].shape == (batch_size, seq_len, 32) -def test_grouped_experts_forward_local_preserves_input_dtype(): - experts = GroupedExperts(num_experts=2, d_model=8, d_ff=12, ffn_dropout=0.0) +@pytest.mark.parametrize("backend", ["loop", "grouped_mm"]) +def test_grouped_experts_forward_preserves_input_dtype(backend: str): + # grouped_mm requires 16-byte aligned strides (d_ff * 2 bytes for bf16) + experts = GroupedExperts(num_experts=2, d_model=8, d_ff=16, ffn_dropout=0.0, backend=backend) experts.reset_parameters() # Input in bf16 while expert weights are initialized in fp32. routed_input = torch.randn(4, 8, dtype=torch.bfloat16) - num_tokens_per_expert = torch.tensor([2, 2], dtype=torch.long) + num_tokens_per_expert = torch.tensor([3, 1], dtype=torch.long) - out = experts._forward_local(routed_input=routed_input, num_tokens_per_expert=num_tokens_per_expert) + out = experts(routed_input, num_tokens_per_expert) assert out.shape == routed_input.shape assert out.dtype == routed_input.dtype -def test_transformer_block_exposes_aux_loss_after_forward(): +@pytest.mark.parametrize("backend", ["loop", "grouped_mm"]) +def test_grouped_experts_matches_per_expert_reference(backend: str): + torch.manual_seed(2) + experts = GroupedExperts(num_experts=3, d_model=8, d_ff=12, ffn_dropout=0.0, backend=backend) + experts.reset_parameters() + counts = [4, 0, 3] + routed_input = torch.randn(sum(counts), 8) + + out = experts(routed_input, torch.tensor(counts)) + + expected = [] + for e, x in enumerate(torch.split(routed_input, counts)): + h = torch.nn.functional.silu(x @ experts.w1[e].T) * (x @ experts.w2[e].T) + expected.append(h @ experts.w3[e].T) + torch.testing.assert_close(out, torch.cat(expected)) + + +def test_grouped_experts_zeroes_dropped_rows(): + experts = GroupedExperts(num_experts=2, d_model=8, d_ff=12, ffn_dropout=0.0, backend="loop") + experts.reset_parameters() + # 5 rows, but only 3 are assigned to experts; the trailing 2 rows were dropped by a capacity limit + out = experts(torch.randn(5, 8), torch.tensor([2, 1])) + assert out.shape == (5, 8) + assert torch.all(out[3:] == 0) + + +def test_grouped_experts_init_std_is_per_expert(): + experts = GroupedExperts(num_experts=4, d_model=256, d_ff=128, ffn_dropout=0.0) + experts.reset_parameters() + for w in (experts.w1, experts.w2, experts.w3): + assert 0.015 < w.std().item() < 0.025 + + +def test_moe_block_records_router_statistics_after_forward(): torch.manual_seed(1) model = _build_tiny_qwen_model() input_ids = torch.randint(0, 32, (2, 4), dtype=torch.long) _ = model({"input_ids": input_ids}) - first_layer = next(iter(model.layers.values())) - assert first_layer.aux_loss is not None + moe_block = next(iter(model.layers.values())).ffn + # local token counts per expert (top_k assignments per token) and the mean routing probabilities + assert moe_block.last_tokens_per_expert.sum().item() == 2 * 4 * moe_block.top_k + assert moe_block.last_router_prob_mean.requires_grad + torch.testing.assert_close(moe_block.last_router_prob_mean.sum(), torch.tensor(1.0)) diff --git a/tests/optimizers/test_ep_adamw.py b/tests/optimizers/test_ep_adamw.py deleted file mode 100644 index bb366627b..000000000 --- a/tests/optimizers/test_ep_adamw.py +++ /dev/null @@ -1,92 +0,0 @@ -import torch -import torch.nn as nn - -from modalities.models.model import NNModel -from modalities.optimizers.ep_adamw import EPAdamW - - -class DummyDPShardMesh: - def __init__(self): - self._group = object() - - def get_group(self): - return self._group - - -class EPSubmodule(nn.Module): - def __init__(self): - super().__init__() - self.ep_weight = nn.Parameter(torch.tensor([1.0, -1.0])) - self._ep_enabled = True - - -class TinyModel(NNModel): - def __init__(self): - super().__init__(weight_decay_groups={"linear": ["linear"], "embedding": [], "layernorm": ["norm"]}) - self.linear = nn.Linear(2, 2, bias=False) - self.norm = nn.LayerNorm(2) - self.experts = EPSubmodule() - - def forward(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - x = inputs["x"] - return {"y": self.linear(x)} - - -def _patch_distributed_ops(monkeypatch): - from modalities.optimizers import ep_adamw as ep_adamw_module - - monkeypatch.setattr(ep_adamw_module.dist, "get_rank", lambda group=None: 0) - monkeypatch.setattr(ep_adamw_module.dist, "get_world_size", lambda group=None: 1) - monkeypatch.setattr(ep_adamw_module.dist, "all_reduce", lambda tensor, op=None, group=None: tensor) - monkeypatch.setattr(ep_adamw_module.dist, "broadcast", lambda tensor, src=0, group=None: tensor) - monkeypatch.setattr(ep_adamw_module.dist, "get_global_rank", lambda group, group_rank: group_rank) - - -def test_ep_adamw_state_dict_and_load_state_dict(monkeypatch): - _patch_distributed_ops(monkeypatch) - - model = TinyModel() - optimizer = EPAdamW( - model=model, - device_mesh={"dp_shard": DummyDPShardMesh()}, - lr=1e-2, - betas=(0.9, 0.95), - eps=1e-8, - weight_decay=0.1, - weight_decay_groups_excluded=["layernorm"], - ) - - state = optimizer.state_dict() - assert "ep_adamw" in state - assert "dense_adamw" in state - - optimizer.load_state_dict(state) - - -def test_ep_adamw_step_updates_parameters_and_zero_grad(monkeypatch): - _patch_distributed_ops(monkeypatch) - - model = TinyModel() - optimizer = EPAdamW( - model=model, - device_mesh={"dp_shard": DummyDPShardMesh()}, - lr=1e-2, - betas=(0.9, 0.95), - eps=1e-8, - weight_decay=0.1, - weight_decay_groups_excluded=["layernorm"], - ) - - before = [p.detach().clone() for p in model.parameters()] - for p in model.parameters(): - p.grad = torch.ones_like(p) - - optimizer.step() - after = list(model.parameters()) - - for p_before, p_after in zip(before, after): - assert not torch.allclose(p_before, p_after) - - optimizer.zero_grad(set_to_none=True) - for p in model.parameters(): - assert p.grad is None