diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py index ca6bdd354ce..9c31b280875 100644 --- a/megatron/core/distributed/finalize_model_grads.py +++ b/megatron/core/distributed/finalize_model_grads.py @@ -1,7 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from functools import partial -from typing import Callable, List, Optional, Union +from typing import Callable, Dict, List, Optional, Union import torch from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors @@ -275,6 +275,44 @@ def _allreduce_position_embedding_grads( ) +def _allreduce_router_grads(model: List[torch.nn.Module], config: TransformerConfig): + """ + All-reduce router grads. + + Reduce grads across all the pp stages to ensure that parameters of the router stay in sync. + """ + + if parallel_state.get_pipeline_model_parallel_world_size() > 1: + grads_dict: Dict[str, List[torch.Tensor]] = {} + for model_chunk in model: + for name, param in get_attr_wrapped_model(model_chunk, 'named_parameters')(): + if param.requires_grad and getattr(param, 'flextron_router_pp_sync', False): + grad = param.main_grad + if name in grads_dict: + # Add all the virtual PP rank's gradients to + # the first local virtual PP rank. + grads_dict[name][0].add_(grad) + # Append to the end for later update after cross-rank reduce. + grads_dict[name].append(grad) + else: + grads_dict[name] = [grad] + + if grads_dict: + # All-reduce the gradient on the first VPP rank. + grads = [param_grad[0] for _, param_grad in grads_dict.items()] + coalesced = _flatten_dense_tensors(grads) + torch.distributed.all_reduce( + coalesced, group=parallel_state.get_pipeline_model_parallel_group() + ) + for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)): + buf.copy_(synced) + + # Update the gradients on other VPP ranks. + for grads in grads_dict.values(): + for grad in grads[1:]: + grad.copy_(grads[0]) + + def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.nn.Module]): """ Reset the temporary tensors of the model. @@ -457,6 +495,9 @@ def finalize_model_grads( if config.timers is not None: config.timers('conditional-embedder-grads-all-reduce').stop() + if getattr(config, 'flextron', False): + _allreduce_router_grads(model, config) + # All-reduce layer-norm grads (for sequence parallelism) and non-tensor parallel modules. if config.timers is not None: config.timers('non-tensor-parallel-grads-all-reduce', log_level=1).start( diff --git a/megatron/elastification/__init__.py b/megatron/elastification/__init__.py new file mode 100644 index 00000000000..26496bfed70 --- /dev/null +++ b/megatron/elastification/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. diff --git a/megatron/elastification/arguments.py b/megatron/elastification/arguments.py new file mode 100644 index 00000000000..a12df474e77 --- /dev/null +++ b/megatron/elastification/arguments.py @@ -0,0 +1,380 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import math + + +def convert_per_lists_to_int_lists(config): + """Convert all *_per_list attributes to *_int_list using model dimensions. + + Called once after model dimensions are known so downstream code can always + use the int-list path without branching on which list type is active. + After this call every *_per_list is None and every *_int_list is set. + """ + conversions = [ + ('emb_per_list', 'emb_int_list', config.hidden_size), + ('mlp_per_list', 'mlp_int_list', config.ffn_hidden_size), + ('mamba_per_list', 'mamba_int_list', config.mamba_num_heads), + ('moe_expert_per_list', 'moe_expert_int_list', config.num_moe_experts), + ] + for per_attr, int_attr, ref_dim in conversions: + per_val = getattr(config, per_attr, None) + if per_val is not None: + setattr(config, int_attr, [math.floor(x * ref_dim) for x in per_val]) + setattr(config, per_attr, None) + + +def sort_budget_list_descending(args): + """Sort ``budget_list`` descending and permute ``budget_probs`` to match. + + The Flextron router's interpolation branch and the elasticity hooks + implicitly assume budget_list is descending (largest first). Sort once + here so downstream code can rely on that invariant regardless of the + order the user passed on the CLI. Idempotent: a list already in + descending order is unchanged. + """ + bl = getattr(args, 'budget_list', None) + if bl is None or len(bl) <= 1: + return + + bp = getattr(args, 'budget_probs', None) + if bp is not None: + assert len(bp) == len(bl), ( + f'budget_probs length {len(bp)} does not match budget_list length {len(bl)}' + ) + + order = sorted(range(len(bl)), key=lambda i: bl[i], reverse=True) + args.budget_list = [bl[i] for i in order] + if bp is not None: + args.budget_probs = [bp[i] for i in order] + + +def validate_flextron_per_int_lists(args): + """ + Enforce mutual exclusion between ratio per-lists and integer choice lists. + + For each module, at most one of (*_per_list, *_int_list) may be set. If neither + is set, *_per_list defaults to [1.0]. Skips when flextron-related args were not + registered on the parser. + """ + pairs = ( + ('mamba', 'mamba_per_list', 'mamba_int_list'), + ('mlp', 'mlp_per_list', 'mlp_int_list'), + ('emb', 'emb_per_list', 'emb_int_list'), + ('moe-expert', 'moe_expert_per_list', 'moe_expert_int_list'), + ) + for cli_name, per_attr, int_attr in pairs: + # Default to None when the attribute is missing - happens when + # flextron args weren't registered on the parser (the docstring + # promises we skip in that case). + per_val = getattr(args, per_attr, None) + int_val = getattr(args, int_attr, None) + per_set = per_val is not None + int_set = int_val is not None + if per_set: + for x in per_val: + assert 0.0 <= x <= 1.0, f'--{cli_name}-per-list values must be in [0, 1], got {x}.' + assert not ( + per_set and int_set + ), f'Use either --{cli_name}-per-list or --{cli_name}-int-list for {cli_name}, not both.' + if not per_set and not int_set: + setattr(args, per_attr, [1.0]) + + +def add_flextron_args(parser): + group = parser.add_argument_group(title='flextron') + # Distillation flags + group.add_argument('--distillation', action='store_true', help='Enable self-distillation.') + group.add_argument('--distill-coeff', type=float, default=0.0, help='Distillation coefficient.') + group.add_argument('--distill-only', action='store_true', help='Distillation only.') + # Basic Flextron flags + group.add_argument('--flextron', action='store_true', help='Enable Flextron.') + group.add_argument('--binary-mask', action='store_true', help='Use binary mask in Flextron.') + group.add_argument('--slice', action='store_true', help='Use slice in Flextron.') + group.add_argument('--enable-router', action='store_true', help='Enable router in Flextron.') + group.add_argument( + '--add-skipping', action='store_true', help='Add layer skipping in Flextron.' + ) + group.add_argument('--no-attn-skip', action='store_true', help='No attn skip in Flextron.') + group.add_argument( + '--lr-mult-router', + type=float, + default=1.0, + help='Learning rate multiplier for router in Flextron.', + ) + group.add_argument('--flex-strict', action='store_true', help='Strict loading of Flextron.') + group.add_argument('--is-flex-eval', action='store_true', help='Is Flextron evaluation.') + group.add_argument('--freeze-router', action='store_true', help='Freeze router in Flextron.') + group.add_argument('--freeze-model', action='store_true', help='Freeze model in Flextron.') + group.add_argument( + '--flex-hetero-ffn', action='store_true', help='Use flex hetero FFN in Flextron.' + ) + group.add_argument( + '--flex-hetero-mamba', action='store_true', help='Use flex hetero Mamba in Flextron.' + ) + group.add_argument( + '--flex-hetero-moe-expert', + action='store_true', + help='Use flex hetero MoE expert in Flextron.', + ) + group.add_argument( + '--router-std', type=float, default=0.1, help='Router init std for Flextron.' + ) + group.add_argument( + '--normalize-router-logits', + action='store_true', + help='Normalize router logits in Flextron.', + ) + group.add_argument('--soft-mask', action='store_true', help='Soft mask in Flextron.') + + # Flextron hyperparameters + group.add_argument( + '--budget-probs', + nargs='+', + type=float, + default=None, + help='List of budget probabilities for Flextron.', + ) + group.add_argument( + '--prefill-chunk-size', type=int, default=16384, help='Prefill chunk size for Flextron.' + ) + group.add_argument( + '--mem-infer-seq-len', + type=int, + default=131072, + help='Memory infer sequence length for Flextron.', + ) + group.add_argument( + '--mem-batch-size', type=int, default=1, help='Memory batch size for Flextron.' + ) + group.add_argument( + '--original-model-sample-prob', + type=float, + default=0.33, + help='Probability of sampling the original model in Flextron.', + ) + group.add_argument( + '--force-router-skip', + nargs='+', + type=int, + default=None, + help='Force router skip for Flextron router.', + ) + group.add_argument( + '--force-mlp', nargs='+', type=float, default=None, help='Force MLP for Flextron router.' + ) + group.add_argument( + '--force-mamba', + nargs='+', + type=float, + default=None, + help='Force Mamba for Flextron router.', + ) + group.add_argument( + '--force-emb', + nargs='+', + type=float, + default=None, + help='Force Embedding for Flextron router.', + ) + group.add_argument( + '--skip-num-attn-layer-constraint', + type=int, + default=None, + help='Skip number of attention layer constraint for Flextron router.', + ) + group.add_argument( + '--skip-total-layer-constraint', + type=int, + default=None, + help='Skip total layer constraint for Flextron router.', + ) + group.add_argument( + '--disable-budget', action='store_true', help='Disable budget for Flextron router.' + ) + group.add_argument( + '--curr-iteration', type=int, default=None, help='Current iteration for Flextron router.' + ) + group.add_argument( + '--hard-sample-th', + type=float, + default=0.996, + help='Hard sample threshold for Flextron router.', + ) + group.add_argument( + '--router-beta', type=float, default=1.0, help='Beta value for Flextron router.' + ) + group.add_argument( + '--loss-alpha', type=float, default=1.0, help='Alpha coefficient for Flextron loss.' + ) + group.add_argument('--tau-init', type=float, default=1.0, help='Tau init for Flextron router.') + group.add_argument( + '--tau-decay', type=float, default=0.9999, help='Tau decay for Flextron router.' + ) + group.add_argument( + '--router-inter-dim', + type=int, + default=128, + help='Intermediate dimension for Flextron router.', + ) + group.add_argument( + '--linear-scaler-start', + type=float, + default=1.0, + help='Linear scaler start for Flextron router.', + ) + group.add_argument( + '--linear-scaler-end', + type=float, + default=10.0, + help='Linear scaler end for Flextron router.', + ) + group.add_argument( + '--override-selected-budget', + nargs='+', + type=float, + default=None, + help='Override selected budget for Flextron router.', + ) + group.add_argument('--router-gbs', type=int, default=32, help='Router gbs for Flextron router.') + # Model configuration lists + group.add_argument( + '--budget-list', + nargs='+', + type=float, + default=[1.0], + help='List of budget values for Flextron.', + ) + group.add_argument( + '--mamba-per-list', + nargs='+', + type=float, + default=None, + help='List of Mamba percentage values for Flextron (mutually exclusive with --mamba-int-list).', + ) + group.add_argument( + '--mlp-per-list', + nargs='+', + type=float, + default=None, + help='List of MLP percentage values for Flextron (mutually exclusive with --mlp-int-list).', + ) + group.add_argument( + '--emb-per-list', + nargs='+', + type=float, + default=None, + help='List of embedding percentage values for Flextron (mutually exclusive with --emb-int-list).', + ) + group.add_argument( + '--moe-expert-per-list', + nargs='+', + type=float, + default=None, + help='List of MoE expert percentage values for Flextron (mutually exclusive with --moe-expert-int-list).', + ) + group.add_argument( + '--mamba-int-list', + nargs='+', + type=int, + default=None, + help='List of Mamba integer router choices for Flextron (mutually exclusive with --mamba-per-list).', + ) + group.add_argument( + '--mlp-int-list', + nargs='+', + type=int, + default=None, + help='List of MLP integer router choices for Flextron (mutually exclusive with --mlp-per-list).', + ) + group.add_argument( + '--emb-int-list', + nargs='+', + type=int, + default=None, + help='List of embedding integer router choices for Flextron (mutually exclusive with --emb-per-list).', + ) + group.add_argument( + '--moe-expert-int-list', + nargs='+', + type=int, + default=None, + help='List of MoE expert integer router choices for Flextron (mutually exclusive with --moe-expert-per-list).', + ) + group.add_argument( + '--budget-type', type=str, default='param', choices=['param', 'mem'], help='Type of budget.' + ) + # Memory quantization profile + group.add_argument( + '--memory-profile', + type=str, + default='bf16', + help='Named memory quantization preset from memory_profiles.yaml ' + '(e.g. bf16, fp8_kv, fp8_all, int8). ' + 'Individual --bpe-* overrides take priority.', + ) + group.add_argument( + '--memory-profile-path', + type=str, + default=None, + help='Path to a custom memory_profiles.yaml. ' + 'Defaults to the bundled megatron/elastification/memory_profiles.yaml.', + ) + group.add_argument( + '--bpe-params', + type=float, + default=None, + help='Override bytes-per-element for model parameters ' '(2=BF16, 1=FP8/INT8, 0.5625=FP4).', + ) + group.add_argument( + '--bpe-kv-cache', type=float, default=None, help='Override bytes-per-element for KV cache.' + ) + group.add_argument( + '--bpe-ssm-cache', + type=float, + default=None, + help='Override bytes-per-element for Mamba SSM state cache.', + ) + group.add_argument( + '--bpe-max-buffer', + type=float, + default=None, + help='Override bytes-per-element for MoE dispatch buffer.', + ) + group.add_argument( + '--param-budget-target', + type=str, + default=None, + choices=['active', 'total'], + help='Whether param budget loss supervises on active params ' + '(top-k experts only) or total params. ' + 'Overrides the preset value from --memory-profile.', + ) + group.add_argument( + '--layer-ranking-list', nargs='+', type=int, default=None, help='List of layer ranking.' + ) + group.add_argument( + '--log-budgets', + nargs='+', + type=str, + default=["all"], + help='Budget values to log distillation loss for (space-separated list or "all").', + ) + # Additional parameters + + group.add_argument( + '--basemodel-type', + type=str, + default='nemotronh_8b', + choices=['nemotronh_8b'], + help='Base model type for parameter loss calculation.', + ) + + # Budget configuration + group.add_argument( + '--flextron-config-file', + type=str, + default=None, + help='Configuration file for Flextron budget settings.', + ) + + return parser diff --git a/megatron/elastification/flextron_config.py b/megatron/elastification/flextron_config.py new file mode 100644 index 00000000000..620e7ef2557 --- /dev/null +++ b/megatron/elastification/flextron_config.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +""" +FlextronConfig — all Flextron/elastification config fields in one place. + +Previously these lived as fields on TransformerConfig (megatron/core). +They are now injected onto the model config at runtime via inject_flextron_config +so that megatron/core stays clean. +""" + +import dataclasses +from dataclasses import dataclass +from typing import List, Optional + + +@dataclass +class FlextronConfig: + # ── Core flags ──────────────────────────────────────────────────────────── + flextron: bool = False + binary_mask: bool = False + add_skipping: bool = False + no_attn_skip: bool = False + slice: bool = False + soft_mask: bool = False + + # ── Router ──────────────────────────────────────────────────────────────── + enable_router: bool = False + router_inter_dim: int = 128 + hard_sample_th: float = 0.996 + router_beta: float = 1.0 + loss_alpha: float = 1.0 + tau_init: float = 1.0 + tau_decay: float = 0.9999 + router_std: float = 0.1 + router_gbs: int = 32 + normalize_router_logits: bool = False + linear_scaler_start: Optional[float] = None + linear_scaler_end: Optional[float] = None + + # ── Budget ──────────────────────────────────────────────────────────────── + budget_probs: Optional[List[float]] = None + budget_list: Optional[List[float]] = None + budget_type: str = 'param' + disable_budget: bool = False + + # ── Training / eval control ─────────────────────────────────────────────── + basemodel_type: str = 'nemotronh_8b' + is_flex_eval: bool = False + freeze_router: bool = False + freeze_model: bool = False + curr_iteration: Optional[int] = None + original_model_sample_prob: float = 0.33 + override_selected_budget: Optional[List[float]] = None + + # ── Layer-skip constraints ──────────────────────────────────────────────── + skip_num_attn_layer_constraint: Optional[int] = None + skip_total_layer_constraint: Optional[int] = None + layer_ranking_list: Optional[List[int]] = None + + # ── Force overrides (eval / frozen-router mode) ─────────────────────────── + force_router_skip: Optional[List[int]] = None + force_mlp: Optional[List[float]] = None + force_mamba: Optional[List[float]] = None + force_emb: Optional[List[float]] = None + + # ── Choice lists (converted to int at model-setup time) ─────────────────── + mamba_per_list: Optional[List[float]] = None + mlp_per_list: Optional[List[float]] = None + emb_per_list: Optional[List[float]] = None + moe_expert_per_list: Optional[List[float]] = None + mamba_int_list: Optional[List[int]] = None + mlp_int_list: Optional[List[int]] = None + emb_int_list: Optional[List[int]] = None + moe_expert_int_list: Optional[List[int]] = None + + # ── Heterogeneous per-layer routing ─────────────────────────────────────── + flex_hetero_ffn: bool = False + flex_hetero_mamba: bool = False + flex_hetero_moe_expert: bool = False + + # ── Memory / inference sizing ───────────────────────────────────────────── + prefill_chunk_size: int = 16384 + mem_infer_seq_len: int = 131072 + mem_batch_size: int = 1 + + # ── Distillation ────────────────────────────────────────────────────────── + distillation: bool = False + distill_coeff: float = 0.0 + distill_only: bool = False + + +def inject_flextron_config(args, config) -> None: + """Copy all FlextronConfig fields from parsed args onto an existing config object. + + Safe to call even when flextron args were not registered on the parser — + falls back to FlextronConfig defaults via getattr. After this call every + FlextronConfig field is accessible directly as config.. + """ + # Validate per-list/int-list mutual exclusion, apply default fallbacks, + # and sort the budget list descending before copying onto config so + # downstream code sees the resolved + ordered state. + from megatron.elastification.arguments import ( + sort_budget_list_descending, + validate_flextron_per_int_lists, + ) + + validate_flextron_per_int_lists(args) + sort_budget_list_descending(args) + + defaults = FlextronConfig() + for f in dataclasses.fields(defaults): + value = getattr(args, f.name, getattr(defaults, f.name)) + setattr(config, f.name, value) diff --git a/megatron/elastification/flextron_elasticity_hooks.py b/megatron/elastification/flextron_elasticity_hooks.py new file mode 100644 index 00000000000..da815043865 --- /dev/null +++ b/megatron/elastification/flextron_elasticity_hooks.py @@ -0,0 +1,1832 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +""" +Flextron Elasticity Hooks + +Applies elasticity masking through PyTorch hooks without modifying original +modules. One manager class per module type (MambaMixer, SelfAttention, +TransformerLayer, MoELayer, TopKRouter, TEGroupedMLP, HybridStack). +""" + +import math +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn as nn + +from megatron.core import parallel_state +from megatron.core.tensor_parallel.utils import split_tensor_along_last_dim +from megatron.core.transformer.moe.moe_utils import get_capacity, group_limited_topk + + +class FlextronMambaElasticityManager: + """ + Manages elasticity for MambaMixer using pure PyTorch hooks. + Based on the exact implementation from original flextron_os MambaMixer. + """ + + def __init__(self, config, layer_idx=0): + self.config = config + self.layer_idx = layer_idx + self.enabled = getattr(config, 'flextron', False) + + if not self.enabled: + return + + # Current elasticity parameters - store the full router outputs + self.current_router_emb = None + self.current_router_mamba = None + + # Hook handles for cleanup + self.hook_handles = [] + + def _init_embedding_masks(self): + """Initialize embedding dimension masks.""" + mask_list = [] + for emb_int in self.config.emb_int_list: + assert ( + 0 <= emb_int <= self.config.hidden_size + ), f'emb_int_list entries must be in [0, hidden_size={self.config.hidden_size}], got {emb_int}.' + mask = torch.zeros(self.config.hidden_size, dtype=torch.bool) + mask[:emb_int] = True + mask_list.append(mask) + self.emb_masks_lookup = { + emb_int: idx for idx, emb_int in enumerate(self.config.emb_int_list) + } + self.emb_masks = torch.stack(mask_list, dim=0).to(device='cuda').to(dtype=torch.bfloat16) + + def _init_mamba_masks(self): + """Initialize Mamba-specific masks for different layers.""" + in_proj_mask_list = [] + conv1d_mask_list = [] + + world_size = parallel_state.get_tensor_model_parallel_world_size() + + in_proj_z_shard = [i for i in range(self.mamba_mixer.d_inner_local_tp)] + in_proj_x_shard = [ + i + for i in range(self.mamba_mixer.d_inner_local_tp, 2 * self.mamba_mixer.d_inner_local_tp) + ] + in_proj_B_shard = [ + i + for i in range( + 2 * self.mamba_mixer.d_inner_local_tp, + 2 * self.mamba_mixer.d_inner_local_tp + + self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state, + ) + ] + in_proj_C_shard = [ + i + for i in range( + 2 * self.mamba_mixer.d_inner_local_tp + + self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state, + 2 * self.mamba_mixer.d_inner_local_tp + + 2 * self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state, + ) + ] + in_proj_dt_shard = [ + i + for i in range( + 2 * self.mamba_mixer.d_inner_local_tp + + 2 * self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state, + 2 * self.mamba_mixer.d_inner_local_tp + + 2 * self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state + + self.mamba_mixer.nheads_local_tp, + ) + ] + + conv1d_x_shard = [i for i in range(self.mamba_mixer.d_inner_local_tp)] + conv1d_B_shard = [ + i + for i in range( + self.mamba_mixer.d_inner_local_tp, + self.mamba_mixer.d_inner_local_tp + + self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state, + ) + ] + conv1d_C_shard = [ + i + for i in range( + self.mamba_mixer.d_inner_local_tp + + self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state, + self.mamba_mixer.d_inner_local_tp + + 2 * self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state, + ) + ] + + out_proj_x_shard = [i for i in range(self.mamba_mixer.d_inner_local_tp)] + + tp_size = parallel_state.get_tensor_model_parallel_world_size() + for mamba_int in self.config.mamba_int_list: + assert ( + 0 <= mamba_int <= self.mamba_mixer.nheads + ), f"mamba_int_list entries must be in [0, nheads={self.mamba_mixer.nheads}], got {mamba_int}." + assert ( + mamba_int % tp_size == 0 + ), f"mamba_int_list entries must be evenly divisible by tp_size={tp_size}, got {mamba_int}." + mamba_nhead_idx = mamba_int // tp_size + + in_proj_mask = torch.zeros( + self.mamba_mixer.d_inner_local_tp * 2 + + self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state * 2 + + self.mamba_mixer.nheads_local_tp, + dtype=torch.bool, + ) + in_proj_mask[in_proj_z_shard[: mamba_nhead_idx * self.mamba_mixer.headdim]] = True + in_proj_mask[in_proj_x_shard[: mamba_nhead_idx * self.mamba_mixer.headdim]] = True + in_proj_mask[in_proj_B_shard] = True + in_proj_mask[in_proj_C_shard] = True + in_proj_mask[in_proj_dt_shard[:mamba_nhead_idx]] = True + in_proj_mask_list.append(in_proj_mask) + + conv1d_mask = torch.zeros( + self.mamba_mixer.d_inner_local_tp + + self.mamba_mixer.ngroups_local_tp * self.mamba_mixer.d_state * 2, + dtype=torch.bool, + ) + conv1d_mask[conv1d_x_shard[: mamba_nhead_idx * self.mamba_mixer.headdim]] = True + conv1d_mask[conv1d_B_shard] = True + conv1d_mask[conv1d_C_shard] = True + conv1d_mask_list.append(conv1d_mask) + + self.mamba_masks_lookup = { + mamba_int: idx for idx, mamba_int in enumerate(self.config.mamba_int_list) + } + + in_proj_mask_list = [item.to(in_proj_mask_list[0].device) for item in in_proj_mask_list] + self.in_proj_mask_list = ( + torch.stack(in_proj_mask_list, dim=0).to(device='cuda').to(dtype=torch.bfloat16) + ) + + conv1d_mask_list = [item.to(conv1d_mask_list[0].device) for item in conv1d_mask_list] + self.conv1d_mask_list = ( + torch.stack(conv1d_mask_list, dim=0).to(device='cuda').to(dtype=torch.bfloat16) + ) + + def attach_hooks(self, mamba_mixer): + """Attach hooks to MambaMixer following the original flextron_os pattern.""" + if not self.enabled: + return + + self.mamba_mixer = mamba_mixer + + emb_effective_per_list = [x / self.config.hidden_size for x in self.config.emb_int_list] + mamba_effective_per_list = [x / mamba_mixer.nheads for x in self.config.mamba_int_list] + + # Setup hook - runs first to initialize masks for this forward pass + def setup_masks_hook(module, input): + if self.config.flextron: + self._init_embedding_masks() + self._init_mamba_masks() + return input + + # Cleanup hook - runs last to remove masks after forward pass + def cleanup_masks_hook(module, input, output): + if self.config.flextron: + self.emb_masks = None + self.in_proj_mask_list = None + self.conv1d_mask_list = None + self.emb_masks_lookup = {} + self.mamba_masks_lookup = {} + return output + + # Hook 1: Input masking and router_emb processing + def input_mask_hook(module, input): + if self.config.flextron and self.current_router_emb is not None: + hidden_states = input[0] + + # Apply embedding mask + if self.config.soft_mask: + soft_mask = torch.zeros( + self.emb_masks[0].shape, + dtype=torch.bfloat16, + device=self.emb_masks[0].device, + ) + for mask, per_logit in zip(self.emb_masks, self.current_router_emb[0]): + soft_mask.add_(mask * per_logit) + mask = soft_mask + masked_input = hidden_states * mask[None, None, :] + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + mask = self.emb_masks[self.emb_masks_lookup[emb_choice]] + masked_input = hidden_states * mask[None, None, :] + masked_input = masked_input * router_emb_logits + + return tuple([masked_input] + list(input[1:])) + + return input + + # Hook 2: in_proj pre-hook for eps modification + def in_proj_pre_hook(module, input): + if self.config.flextron and self.current_router_emb is not None: + + # Set eps to the pruned value + if self.config.soft_mask: + soft_eps = 0 + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_eps += self.config.layernorm_epsilon * emb_per * per_logit + module.eps = soft_eps.float().detach().item() + else: + emb_choice = self.current_router_emb[1] + emb_effective_per = emb_choice / self.config.hidden_size + module.eps = self.config.layernorm_epsilon * emb_effective_per + return input + + # Hook 3: in_proj post-hook for router scaling + def in_proj_post_hook(module, input, output): + if self.config.flextron and self.current_router_mamba is not None: + # Apply router_emb scaling to in_proj output + xz, bias = output + + if self.config.soft_mask: + # Soft scaling with embedding router + soft_xz = torch.zeros_like(xz) + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_xz.add_(xz * per_logit * (emb_per**0.5)) + xz = soft_xz + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + emb_effective_per = emb_choice / self.config.hidden_size + xz = xz * router_emb_logits * (emb_effective_per**0.5) + + # Apply mamba router logic (hard mask only) + if not self.config.soft_mask: + if self.config.flex_hetero_mamba: + mamba_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('M') - 1 + ) + router_mamba_logits = torch.max(self.current_router_mamba[0][mamba_idx]) + mamba_per = self.current_router_mamba[1][mamba_idx] + else: + router_mamba_logits, mamba_per = ( + torch.max(self.current_router_mamba[0]), + self.current_router_mamba[1], + ) + + # Apply mamba masking + if self.config.soft_mask: + soft_in_proj_mask = torch.zeros_like(self.in_proj_mask_list[0]) + if self.config.flex_hetero_mamba: + mamba_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('M') - 1 + ) + for mask, per_logit in zip( + self.in_proj_mask_list, self.current_router_mamba[0][mamba_idx] + ): + soft_in_proj_mask.add_(mask * per_logit) + else: + for mask, per_logit in zip( + self.in_proj_mask_list, self.current_router_mamba[0] + ): + soft_in_proj_mask.add_(mask * per_logit) + in_proj_mask = soft_in_proj_mask + else: + in_proj_mask = self.in_proj_mask_list[self.mamba_masks_lookup[mamba_per]] + + xz = xz * in_proj_mask.to(device=xz.device)[None, None, :] + + if not self.config.soft_mask: + xz = xz * router_mamba_logits + + # Reset eps to original + module.eps = self.config.layernorm_epsilon + + return (xz, bias) + return output + + # Hook 4: conv1d output masking + def conv1d_mask_hook(module, input, output): + if self.config.flextron and self.current_router_mamba is not None: + if not self.config.soft_mask: + if self.config.flex_hetero_mamba: + mamba_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('M') - 1 + ) + router_mamba_logits = torch.max(self.current_router_mamba[0][mamba_idx]) + mamba_per = self.current_router_mamba[1][mamba_idx] + else: + router_mamba_logits, mamba_per = ( + torch.max(self.current_router_mamba[0]), + self.current_router_mamba[1], + ) + + # Apply conv1d masking + if self.config.soft_mask: + soft_conv1d_mask = torch.zeros_like(self.conv1d_mask_list[0]) + if self.config.flex_hetero_mamba: + mamba_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('M') - 1 + ) + for mask, per_logit in zip( + self.conv1d_mask_list, self.current_router_mamba[0][mamba_idx] + ): + soft_conv1d_mask.add_(mask * per_logit) + else: + for mask, per_logit in zip( + self.conv1d_mask_list, self.current_router_mamba[0] + ): + soft_conv1d_mask.add_(mask * per_logit) + conv1d_mask = soft_conv1d_mask + else: + conv1d_mask = self.conv1d_mask_list[self.mamba_masks_lookup[mamba_per]] + masked_output = output * conv1d_mask.to(device=output.device)[None, :, None] + + if not self.config.soft_mask: + masked_output = masked_output * router_mamba_logits + + return masked_output + return output + + # Hook 5a: RMSNorm pre-hook for eps modification + def norm_pre_hook(module, input): + if self.config.flextron and self.current_router_mamba is not None: + if self.config.soft_mask: + soft_eps = 0 + if self.config.flex_hetero_mamba: + mamba_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('M') - 1 + ) + for mamba_per, per_logit in zip( + mamba_effective_per_list, self.current_router_mamba[0][mamba_idx] + ): + soft_eps += self.config.layernorm_epsilon * mamba_per * per_logit + else: + for mamba_per, per_logit in zip( + mamba_effective_per_list, self.current_router_mamba[0] + ): + soft_eps += self.config.layernorm_epsilon * mamba_per * per_logit + module.eps = soft_eps.float().detach().item() + else: + if self.config.flex_hetero_mamba: + mamba_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('M') - 1 + ) + mamba_per = self.current_router_mamba[1][mamba_idx] + else: + mamba_per = self.current_router_mamba[1] + mamba_effective_per = mamba_per / self.mamba_mixer.nheads + module.eps = self.config.layernorm_epsilon * mamba_effective_per + + return input + + # Hook 5b: RMSNorm post-hook for scaling and eps restoration + def norm_post_hook(module, input, output): + if self.config.flextron and self.current_router_mamba is not None: + # Restore original eps + module.eps = self.config.layernorm_epsilon + + if self.config.soft_mask: + soft_scaled_output = torch.zeros_like(output) + if self.config.flex_hetero_mamba: + mamba_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('M') - 1 + ) + for mamba_per, per_logit in zip( + mamba_effective_per_list, self.current_router_mamba[0][mamba_idx] + ): + soft_scaled_output.add_(output * (mamba_per**0.5) * per_logit) + else: + for mamba_per, per_logit in zip( + mamba_effective_per_list, self.current_router_mamba[0] + ): + soft_scaled_output.add_(output * (mamba_per**0.5) * per_logit) + return soft_scaled_output + else: + if self.config.flex_hetero_mamba: + mamba_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('M') - 1 + ) + router_mamba_logits = torch.max(self.current_router_mamba[0][mamba_idx]) + mamba_per = self.current_router_mamba[1][mamba_idx] + else: + router_mamba_logits, mamba_per = ( + torch.max(self.current_router_mamba[0]), + self.current_router_mamba[1], + ) + mamba_effective_per = mamba_per / self.mamba_mixer.nheads + return output * (mamba_effective_per**0.5) * router_mamba_logits + + return output + + # Hook 6: Final output masking + def output_mask_hook(module, input, output): + if self.config.flextron and self.current_router_emb is not None: + out, out_bias = output + + # Apply embedding mask + if self.config.soft_mask: + soft_mask = torch.zeros( + self.emb_masks[0].shape, + dtype=torch.bfloat16, + device=self.emb_masks[0].device, + ) + for mask, per_logit in zip(self.emb_masks, self.current_router_emb[0]): + soft_mask.add_(mask * per_logit) + mask = soft_mask + masked_out = out * mask[None, None, :] + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + mask = self.emb_masks[self.emb_masks_lookup[emb_choice]] + masked_out = out * mask[None, None, :] + masked_out = masked_out * router_emb_logits + + return (masked_out, out_bias) + return output + + # IMPORTANT: Register setup hook FIRST + setup_handle = mamba_mixer.register_forward_pre_hook(setup_masks_hook) + self.hook_handles.append(setup_handle) + + # Attach main input hook + main_handle = mamba_mixer.register_forward_pre_hook(input_mask_hook) + self.hook_handles.append(main_handle) + + # Attach in_proj hooks + in_proj_pre_handle = mamba_mixer.in_proj.register_forward_pre_hook(in_proj_pre_hook) + in_proj_post_handle = mamba_mixer.in_proj.register_forward_hook(in_proj_post_hook) + self.hook_handles.append(in_proj_pre_handle) + self.hook_handles.append(in_proj_post_handle) + + # Attach conv1d hook (this will handle the standard conv1d path) + conv_handle = mamba_mixer.conv1d.register_forward_hook(conv1d_mask_hook) + self.hook_handles.append(conv_handle) + + # Attach RMSNorm hooks if rmsnorm is enabled + norm_pre_handle = mamba_mixer.norm.register_forward_pre_hook(norm_pre_hook) + norm_post_handle = mamba_mixer.norm.register_forward_hook(norm_post_hook) + self.hook_handles.append(norm_pre_handle) + self.hook_handles.append(norm_post_handle) + + # Final output hook + output_handle = mamba_mixer.register_forward_hook(output_mask_hook) + self.hook_handles.append(output_handle) + + # Cleanup hook - runs last to remove masks after forward pass + cleanup_handle = mamba_mixer.register_forward_hook(cleanup_masks_hook) + self.hook_handles.append(cleanup_handle) + + def set_elasticity_params(self, router_emb=None, router_mamba=None, **kwargs): + """Set current elasticity parameters that will be used by hooks.""" + if router_emb is not None: + self.current_router_emb = router_emb + + if router_mamba is not None: + self.current_router_mamba = router_mamba + + def detach_hooks(self): + """Remove all hooks.""" + if not hasattr(self, 'hook_handles'): + return + for handle in self.hook_handles: + handle.remove() + self.hook_handles.clear() + + def __del__(self): + """Cleanup hooks when manager is destroyed.""" + self.detach_hooks() + + +class FlextronTransformerLayerElasticityManager: + """ + Manages elasticity for TransformerLayer using pure PyTorch hooks. + Handles input/pre-MLP layernorm eps modification and MLP routing. + """ + + def __init__(self, config, layer_idx=0): + self.config = config + self.layer_idx = layer_idx + self.enabled = getattr(config, 'flextron', False) + + if not self.enabled: + return + + # Current elasticity parameters - store the full router outputs + self.current_router_emb = None + + # Hook handles for cleanup + self.hook_handles = [] + + def _init_embedding_masks(self): + """Initialize embedding dimension masks.""" + mask_list = [] + for emb_int in self.config.emb_int_list: + mask = torch.zeros(self.config.hidden_size, dtype=torch.bool) + mask[:emb_int] = True + mask_list.append(mask) + self.emb_masks_lookup = { + emb_int: idx for idx, emb_int in enumerate(self.config.emb_int_list) + } + self.emb_masks = torch.stack(mask_list, dim=0).to(device='cuda').to(dtype=torch.bfloat16) + + def initialize_masks(self, transformer_layer): + """Initialize masks based on the MoE module configuration.""" + if not self.enabled: + return + + self.transformer_layer = transformer_layer + self._init_embedding_masks() + + def attach_hooks(self, transformer_layer): + """Attach hooks to MLP/MoE layer for layer skipping only.""" + if not self.enabled: + return + + self.initialize_masks(transformer_layer) + + emb_effective_per_list = [x / self.config.hidden_size for x in self.config.emb_int_list] + + # Hook 2: Pre-MLP layernorm pre-hook for eps modification + def pre_mlp_layernorm_pre_hook(module, input): + + if self.config.flextron and self.current_router_emb is not None: + hidden_states = input[0] + # Apply embedding mask + if self.config.soft_mask: + soft_mask = torch.zeros( + self.emb_masks[0].shape, + dtype=torch.bfloat16, + device=self.emb_masks[0].device, + ) + for mask, per_logit in zip(self.emb_masks, self.current_router_emb[0]): + soft_mask.add_(mask * per_logit) + mask = soft_mask + masked_input = hidden_states * mask[None, None, :] + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + mask = self.emb_masks[self.emb_masks_lookup[emb_choice]] + masked_input = hidden_states * mask[None, None, :] + masked_input = masked_input * router_emb_logits + + # Modify eps for this forward pass + if self.config.soft_mask: + soft_eps = 0 + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_eps += self.config.layernorm_epsilon * emb_per * per_logit + module.eps = soft_eps.float().detach().item() + else: + emb_choice = self.current_router_emb[1] + emb_effective_per = emb_choice / self.config.hidden_size + module.eps = self.config.layernorm_epsilon * emb_effective_per + + return tuple([masked_input] + list(input[1:])) + + return input + + # Hook 3: Pre-MLP layernorm post-hook for scaling and eps restoration + def pre_mlp_layernorm_post_hook(module, input, output): + if self.config.flextron and self.current_router_emb is not None: + + # Restore original eps + module.eps = self.config.layernorm_epsilon + + # Apply scaling + if self.config.soft_mask: + soft_scaled_output = torch.zeros_like(output) + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_scaled_output.add_(output * (emb_per**0.5) * per_logit) + scaled_output = soft_scaled_output + else: + emb_choice = self.current_emb_choice + emb_effective_per = emb_choice / self.config.hidden_size + router_emb_logits = torch.max(self.current_router_emb[0]) + scaled_output = output * (emb_effective_per**0.5) * router_emb_logits + return scaled_output + + return output + + # Attach the pre-MLP layernorm hooks + pre_mlp_ln_pre_handle = transformer_layer.pre_mlp_layernorm.register_forward_pre_hook( + pre_mlp_layernorm_pre_hook + ) + pre_mlp_ln_post_handle = transformer_layer.pre_mlp_layernorm.register_forward_hook( + pre_mlp_layernorm_post_hook + ) + self.hook_handles.append(pre_mlp_ln_pre_handle) + self.hook_handles.append(pre_mlp_ln_post_handle) + + def set_elasticity_params(self, router_emb=None, **kwargs): + """Set current elasticity parameters that will be used by hooks.""" + if router_emb is not None: + self.current_router_emb = router_emb + self.current_emb_choice = router_emb[1] + + def detach_hooks(self): + """Remove all hooks.""" + if not hasattr(self, 'hook_handles'): + return + for handle in self.hook_handles: + handle.remove() + self.hook_handles.clear() + + def __del__(self): + """Cleanup hooks when manager is destroyed.""" + self.detach_hooks() + + +def topk_softmax_with_capacity( + logits: torch.Tensor, + topk: int, + capacity_factor: Optional[float] = None, + pad_to_capacity: bool = False, + drop_policy: str = "probs", + use_pre_softmax: bool = False, + num_groups: Optional[int] = None, + group_topk: Optional[int] = None, + scaling_factor: Optional[float] = None, + deterministic_mode: bool = False, + score_function: str = "softmax", + expert_bias: Optional[torch.Tensor] = None, + current_router_moe_expert_0: Optional[torch.Tensor] = None, + current_router_moe_expert_per: Optional[List[float]] = None, + num_experts: int = 0, +): + """Apply capacity and padding to the top-k selection. + Args: + logits (torch.Tensor): Logits tensor. + topk (int): The number of experts to select for each token. + capacity_factor (float): The capacity factor of each expert. Will drop tokens if the number + of tokens exceeds the capacity. + pad_to_capacity (bool): Whether to need padding in token drop mode. The probs for padded + tokens will be 0. + drop_policy (str): The policy to drop tokens. Can be either "prob" or "position". + If "prob", the tokens with the lowest probabilities will be dropped. + If "position", tokens at the end of each batch will be dropped. + use_pre_softmax (bool): Whether to apply softmax or sigmoid before top-k selection. + num_groups (int): Number of groups for routed experts. + group_topk (int): Number of selected groups for each token. + scaling_factor (float): Scaling factor of routing score in top-k selection. + deterministic_mode (bool): Deprecated. + score_function (str): The score function to use. Can be either "softmax" or "sigmoid". + expert_bias (torch.Tensor): The bias added to logits for expert routing. + Returns: + Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + - routing_probs (torch.Tensor): A tensor of shape [num_tokens, num_experts] containing + the routing probabilities for each token to each expert. + - routing_map (torch.Tensor): A mask tensor of shape [num_tokens, num_experts] + indicating which experts were selected for each token. True values represent + the selected experts. + - tokens_per_expert (torch.Tensor): A tensor of shape [num_experts] containing + the number of local tokens assigned to each expert before dropping and padding. + """ + assert score_function == "sigmoid", "Only sigmoid score function is supported for now." + assert expert_bias is not None, "Expert bias is required for sigmoid score function." + assert logits.dim() == 2, f"Expected 2D logits [num_tokens, num_experts], got {logits.dim()}." + num_tokens, num_experts = logits.shape + + def compute_topk(scores, topk, num_groups=None, group_topk=None): + if group_topk: + return group_limited_topk( + scores=scores, + topk=topk, + num_tokens=num_tokens, + num_experts=num_experts, + num_groups=num_groups, + group_topk=group_topk, + ) + else: + return torch.topk(scores, k=topk, dim=1) + + if score_function == "sigmoid": + scores_for_routing = 0 + scores_for_topk = 0 + for router_moe_expert_logits, router_moe_expert_per in zip( + current_router_moe_expert_0, current_router_moe_expert_per + ): + expert_threshold = math.floor(router_moe_expert_per * num_experts) + + logits_current = logits.clone() + logits_current[:, expert_threshold:] = float('-inf') + expert_bias_current = expert_bias.clone() + expert_bias_current[expert_threshold:] = 0 + expert_bias_current = expert_bias_current * router_moe_expert_logits + + scores = ( + torch.sigmoid(logits_current.float()).type_as(logits_current) + * router_moe_expert_logits + ) + + scores_for_topk += scores + scores_for_routing += scores + expert_bias_current + + _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk) + scores = torch.gather(scores_for_topk, dim=1, index=top_indices).type_as(logits) + + probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores + + else: + raise ValueError(f"Invalid score_function: {score_function}") + + if scaling_factor: + probs = probs * scaling_factor + + # TODO Try using element-wise operations instead of scatter? + topk_masked_gates = torch.zeros_like(logits).scatter(1, top_indices, probs) + topk_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool() + tokens_per_expert = topk_map.sum(dim=0) + + if capacity_factor is None: + # TopK without capacity + return topk_masked_gates, topk_map, tokens_per_expert + else: + # TopK with capacity + expert_capacity = get_capacity( + num_tokens=num_tokens * topk, num_experts=num_experts, capacity_factor=capacity_factor + ) + + # Maskout exceeded tokens + if drop_policy == "probs": + _, capacity_indices = torch.topk( + topk_masked_gates, k=expert_capacity, dim=0, sorted=False + ) + capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1).bool() + elif drop_policy == "position": + _, capacity_indices = torch.topk(topk_map.int(), k=expert_capacity, dim=0, sorted=False) + capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1).bool() + else: + raise ValueError(f"Invalid drop_policy: {drop_policy}") + + if pad_to_capacity: + final_map = capacity_mask + final_probs = topk_masked_gates * final_map + else: + # Get exceed mask and maskout exceeded probs and indices + final_map = torch.logical_and(topk_map, capacity_mask) + final_probs = topk_masked_gates * final_map + return final_probs, final_map, tokens_per_expert + + +class FlextronTopKRouterElasticityManager: + """ + Manages elasticity for MoE Router using pure PyTorch hooks. + Handles expert masking in the routing logits before topk selection. + """ + + def __init__(self, config, layer_idx=0): + self.config = config + self.layer_idx = layer_idx + self.enabled = getattr(config, 'flextron', False) + + if not self.enabled: + return + + # Current elasticity parameters + self.current_router_moe_expert = None + + # Hook handles for cleanup + self.hook_handles = [] + + def attach_hooks(self, router): + """Attach hooks to TopKRouter for expert masking.""" + if not self.enabled: + return + + # Store original method for restoration + original_routing = router.routing + + def wrapped_routing(logits, **kwargs): + + # Apply expert masking before calling original routing + if self.config.flextron and self.current_router_moe_expert is not None: + + if self.config.soft_mask: + if self.config.flex_hetero_moe_expert: + moe_expert_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('E') - 1 + ) + current_router_moe_expert_0 = self.current_router_moe_expert[0][ + moe_expert_idx + ] + else: + current_router_moe_expert_0 = self.current_router_moe_expert[0] + + seq_length, bsz = logits.shape[:2] + logits = logits.view(-1, self.config.num_moe_experts) + + # Apply Z-Loss + logits = router.apply_z_loss(logits) + assert self.config.moe_router_load_balancing_type == "none" + + scores, routing_map, _ = topk_softmax_with_capacity( + logits, + self.config.moe_router_topk, + capacity_factor=self.config.moe_expert_capacity_factor, + pad_to_capacity=self.config.moe_pad_expert_input_to_capacity, + drop_policy=self.config.moe_token_drop_policy, + use_pre_softmax=self.config.moe_router_pre_softmax, + num_groups=self.config.moe_router_num_groups, + group_topk=self.config.moe_router_group_topk, + scaling_factor=self.config.moe_router_topk_scaling_factor, + deterministic_mode=self.config.deterministic_mode, + score_function=self.config.moe_router_score_function, + expert_bias=router.expert_bias, + current_router_moe_expert_0=current_router_moe_expert_0, + current_router_moe_expert_per=[ + x / self.config.num_moe_experts for x in self.config.moe_expert_int_list + ], + num_experts=self.config.num_moe_experts, + ) + + if self.config.moe_router_enable_expert_bias and torch.is_grad_enabled(): + with torch.no_grad(): + router.local_tokens_per_expert += routing_map.sum(dim=0) + + return scores, routing_map + else: + if self.config.flex_hetero_moe_expert: + moe_expert_idx = ( + self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('E') - 1 + ) + router_moe_expert_logits = torch.max( + self.current_router_moe_expert[0][moe_expert_idx] + ) + router_moe_expert_per = self.current_router_moe_expert[1][moe_expert_idx] + else: + router_moe_expert_logits, router_moe_expert_per = ( + torch.max(self.current_router_moe_expert[0]), + self.current_router_moe_expert[1], + ) + + expert_threshold = ( + router_moe_expert_per # always an integer count after conversion + ) + + # Apply the same logic as the commented lines + logits = logits.clone() + logits[:, expert_threshold:] = float('-inf') + logits = logits * router_moe_expert_logits + + # Mask the expert_bias on a temporary tensor and restore + # the original after the call so subsequent forwards + # (including full-model passes that bypass this branch) + # don't see a truncated bias. + if hasattr(router, 'expert_bias') and router.expert_bias is not None: + original_expert_bias = router.expert_bias + masked_bias = original_expert_bias.clone() + masked_bias[expert_threshold:] = 0 + router.expert_bias = masked_bias + try: + return original_routing(logits, **kwargs) + finally: + router.expert_bias = original_expert_bias + return original_routing(logits, **kwargs) + + else: + return original_routing(logits, **kwargs) + + router.routing = wrapped_routing + # Store reference to restore later + router._original_routing = original_routing + self.hook_handles.append(('method_replacement', router, 'routing')) + + def set_elasticity_params(self, router_moe_expert=None, **kwargs): + """Set current elasticity parameters that will be used by hooks.""" + if router_moe_expert is not None: + self.current_router_moe_expert = router_moe_expert + + def detach_hooks(self): + """Remove all hooks and restore original methods.""" + if not hasattr(self, 'hook_handles'): + return + for handle in self.hook_handles: + if isinstance(handle, tuple) and handle[0] == 'method_replacement': + # Restore original method + _, router, method_name = handle + if hasattr(router, '_original_routing'): + router.routing = router._original_routing + delattr(router, '_original_routing') + else: + # Regular hook handle + handle.remove() + self.hook_handles.clear() + + def __del__(self): + """Cleanup hooks when manager is destroyed.""" + self.detach_hooks() + + +class FlextronMoEElasticityManager: + """ + Manages elasticity for MLP/MoE layers using pure PyTorch hooks. + Now supports both traditional MLP ('-') and MoE ('E') layers with layer skipping. + """ + + def __init__(self, config, layer_idx=0): + self.config = config + self.layer_idx = layer_idx + self.enabled = getattr(config, 'flextron', False) + + if not self.enabled: + return + + # Current elasticity parameters - store the full router outputs + self.current_router_emb = None + # Hook handles for cleanup + self.hook_handles = [] + + def _init_embedding_masks(self): + """Initialize embedding dimension masks.""" + mask_list = [] + for emb_int in self.config.emb_int_list: + mask = torch.zeros(self.config.hidden_size, dtype=torch.bool) + mask[:emb_int] = True + mask_list.append(mask) + self.emb_masks_lookup = { + emb_int: idx for idx, emb_int in enumerate(self.config.emb_int_list) + } + self.emb_masks = torch.stack(mask_list, dim=0).to(device='cuda').to(dtype=torch.bfloat16) + + def initialize_masks(self, moe_module): + """Initialize masks based on the MoE module configuration.""" + if not self.enabled: + return + + self.moe_module = moe_module + self._init_embedding_masks() + + def attach_hooks(self, moe_module): + """Attach hooks to MLP/MoE layer for layer skipping only.""" + if not self.enabled: + return + + self.initialize_masks(moe_module) + + def output_mask_hook(module, input, output): + + if self.config.flextron and self.current_router_emb is not None: + out, out_bias = output + + if self.config.soft_mask: + soft_mask = torch.zeros( + self.emb_masks[0].shape, + dtype=torch.bfloat16, + device=self.emb_masks[0].device, + ) + for mask, per_logit in zip(self.emb_masks, self.current_router_emb[0]): + soft_mask.add_(mask * per_logit) + mask = soft_mask + masked_out = out * mask[None, None, :] + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + mask = self.emb_masks[self.emb_masks_lookup[emb_choice]] + masked_out = out * mask[None, None, :] + masked_out = masked_out * router_emb_logits + + return (masked_out, out_bias) + return output + + # Attach the output hook + + output_handle = moe_module.register_forward_hook(output_mask_hook) + self.hook_handles.append(output_handle) + + def set_elasticity_params(self, router_emb=None, **kwargs): + """Set current elasticity parameters that will be used by hooks.""" + if router_emb is not None: + self.current_router_emb = router_emb + + def detach_hooks(self): + """Remove all hooks.""" + if not hasattr(self, 'hook_handles'): + return + for handle in self.hook_handles: + handle.remove() + self.hook_handles.clear() + + def __del__(self): + """Cleanup hooks when manager is destroyed.""" + self.detach_hooks() + + +class FlextronGroupedMLPElasticityManager: + + def __init__(self, config, layer_idx=0): + self.config = config + self.layer_idx = layer_idx + self.enabled = getattr(config, 'flextron', False) + self.mlp_idx = self.config.hybrid_layer_pattern[: self.layer_idx + 1].count('E') - 1 + + if not self.enabled: + return + + self.current_router_mlp = None + self.current_router_emb = None + + self.hook_handles = [] + + def _init_embedding_masks(self): + """Initialize embedding dimension masks.""" + mask_list = [] + for emb_int in self.config.emb_int_list: + mask = torch.zeros(self.config.hidden_size, dtype=torch.bool) + mask[:emb_int] = True + mask_list.append(mask) + self.emb_masks_lookup = { + emb_int: idx for idx, emb_int in enumerate(self.config.emb_int_list) + } + self.emb_masks = torch.stack(mask_list, dim=0).to(device='cuda').to(dtype=torch.bfloat16) + + def _init_mlp_masks(self): + """Initialize MLP-specific masks.""" + mask_list = [] + list_mlp_mask = list(set(self.config.mlp_int_list)) + list_mlp_mask.sort(reverse=True) + for mlp_int in list_mlp_mask: + mask_temp = torch.zeros(self.config.ffn_hidden_size, dtype=torch.bool) + mask_temp[:mlp_int] = True + mask_list.append(mask_temp) + mask_list = [item.to(mask_list[0].device) for item in mask_list] + self.mlp_intermediate_masks = ( + torch.stack(mask_list, dim=0).to(device='cuda').to(dtype=torch.bfloat16) + ) + self.mlp_intermediate_masks_lookup = { + mlp_int: idx for idx, mlp_int in enumerate(list_mlp_mask) + } + + def initialize_masks(self, mlp_module): + """Initialize masks based on the MLP configuration.""" + if not self.enabled: + return + + self.mlp_module = mlp_module + self._init_embedding_masks() + self._init_mlp_masks() + + def attach_hooks(self, mlp_module): + """Attach hooks to MLP following the original flextron_os pattern.""" + if not self.enabled: + return + + self.mlp_module = mlp_module + + emb_effective_per_list = [x / self.config.hidden_size for x in self.config.emb_int_list] + + # Setup hook - runs first to initialize masks for this forward pass + def setup_masks_hook(module, input): + if self.config.flextron: + self._init_embedding_masks() + self._init_mlp_masks() + return input + + # Cleanup hook - runs last to remove masks after forward pass + def cleanup_masks_hook(module, input, output): + if self.config.flextron: + self.emb_masks = None + self.mlp_intermediate_masks = None + self.emb_masks_lookup = {} + self.mlp_intermediate_masks_lookup = {} + return output + + # IMPORTANT: Register setup hook FIRST + setup_handle = mlp_module.register_forward_pre_hook(setup_masks_hook) + self.hook_handles.append(setup_handle) + + # Hook 1: Input masking and router_emb processing + def input_mask_hook(module, input): + if self.config.flextron and self.current_router_emb is not None: + hidden_states = input[0] + + # Apply embedding mask + if self.config.soft_mask: + soft_mask = torch.zeros( + self.emb_masks[0].shape, + dtype=torch.bfloat16, + device=self.emb_masks[0].device, + ) + for mask, per_logit in zip(self.emb_masks, self.current_router_emb[0]): + soft_mask.add_(mask * per_logit) + mask = soft_mask + masked_input = hidden_states * mask[None, :] + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + mask = self.emb_masks[self.emb_masks_lookup[emb_choice]] + masked_input = hidden_states * mask[None, :] + masked_input = masked_input * router_emb_logits + + # Process router_mlp logic here. Both attributes are read + # by fc1_post_hook only when current_router_mlp is not None, + # so we only assign them in that branch (avoids + # UnboundLocalError on mlp_per when the router produced no + # MLP output this step). + if self.current_router_mlp is not None: + if self.config.flex_hetero_ffn: + router_weights = torch.max(self.current_router_mlp[0][self.mlp_idx]) + mlp_per = self.current_router_mlp[1][self.mlp_idx] + else: + router_weights, mlp_per = ( + torch.max(self.current_router_mlp[0]), + self.current_router_mlp[1], + ) + module._flextron_router_weights = router_weights + module._flextron_mlp_per = mlp_per + + return tuple([masked_input] + list(input[1:])) + return input + + # Hook 2: Linear FC1 post-hook for router scaling and masking + def fc1_post_hook(module, input, output): + + # Apply router_emb scaling and MLP masking. Both router_emb and + # router_mlp must be set: the body reads current_router_emb[0] + # for the emb scaling and current_router_mlp[0] for the mask. + # Today they're always set together by update_hook_elasticity_params, + # but guarding both makes the precondition explicit. + if ( + self.config.flextron + and self.current_router_mlp is not None + and self.current_router_emb is not None + ): + intermediate_parallel, bias_parallel = output + if self.config.soft_mask: + soft_intermediate_parallel = torch.zeros_like(intermediate_parallel) + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_intermediate_parallel.add_(intermediate_parallel * per_logit) + intermediate_parallel = soft_intermediate_parallel + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + intermediate_parallel = intermediate_parallel * router_emb_logits + + # # Apply MLP masking and router weights + + mlp_per = mlp_module._flextron_mlp_per + router_weights = getattr(mlp_module, '_flextron_router_weights', None) + + # Apply masking + if self.config.soft_mask: + soft_mask = torch.zeros( + self.mlp_intermediate_masks[0].shape, + dtype=torch.bfloat16, + device=self.mlp_intermediate_masks[0].device, + ) + if self.config.flex_hetero_ffn: + for mask, per_logit in zip( + self.mlp_intermediate_masks, self.current_router_mlp[0][self.mlp_idx] + ): + soft_mask.add_(mask * per_logit) + else: + for mask, per_logit in zip( + self.mlp_intermediate_masks, self.current_router_mlp[0] + ): + soft_mask.add_(mask * per_logit) + mask = soft_mask + else: + mask = self.mlp_intermediate_masks[self.mlp_intermediate_masks_lookup[mlp_per]] + + world_size = parallel_state.get_expert_tensor_parallel_world_size() + mask_list = split_tensor_along_last_dim(mask, world_size) + rank = parallel_state.get_expert_tensor_parallel_rank() + + mask = mask_list[rank].contiguous() + + intermediate_parallel = ( + intermediate_parallel * mask.to(device=intermediate_parallel.device)[None, :] + ) + if router_weights is not None and not self.config.soft_mask: + intermediate_parallel = intermediate_parallel * router_weights + + module.eps = self.config.layernorm_epsilon + + return (intermediate_parallel, bias_parallel) + return output + + # Hook 3: Final output masking + def output_mask_hook(module, input, output): + if self.config.flextron and self.current_router_emb is not None: + out, out_bias = output + + # Apply embedding mask + if self.config.soft_mask: + soft_mask = torch.zeros( + self.emb_masks[0].shape, + dtype=torch.bfloat16, + device=self.emb_masks[0].device, + ) + for mask, per_logit in zip(self.emb_masks, self.current_router_emb[0]): + soft_mask.add_(mask * per_logit) + mask = soft_mask + masked_out = out * mask[None, :] + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + mask = self.emb_masks[self.emb_masks_lookup[emb_choice]] + masked_out = out * mask[None, :] + masked_out = masked_out * router_emb_logits + + return (masked_out, out_bias) + return output + + # Hook 1: Input masking and router_emb processing + main_handle = mlp_module.register_forward_pre_hook(input_mask_hook) + self.hook_handles.append(main_handle) + + # Hook 2: Linear FC1 pre-hook for eps modification + fc1_post_handle = mlp_module.linear_fc1.register_forward_hook(fc1_post_hook) + self.hook_handles.append(fc1_post_handle) + + # Hook 3: Final output masking + output_handle = mlp_module.register_forward_hook(output_mask_hook) + self.hook_handles.append(output_handle) + + # Cleanup hook - runs last to remove masks after forward pass + cleanup_handle = mlp_module.register_forward_hook(cleanup_masks_hook) + self.hook_handles.append(cleanup_handle) + + def set_elasticity_params(self, router_emb=None, router_mlp=None, **kwargs): + """Set current elasticity parameters that will be used by hooks.""" + if router_emb is not None: + self.current_router_emb = router_emb + + if router_mlp is not None: + self.current_router_mlp = router_mlp + + def detach_hooks(self): + """Remove all hooks.""" + if not hasattr(self, 'hook_handles'): + return + for handle in self.hook_handles: + handle.remove() + self.hook_handles.clear() + + def __del__(self): + """Cleanup hooks when manager is destroyed.""" + self.detach_hooks() + + +class FlextronAttentionElasticityManager: + """ + Manages elasticity for Attention using pure PyTorch hooks. + Based on the exact implementation from original flextron_os Attention. + """ + + def __init__(self, config, layer_idx=0): + self.config = config + self.layer_idx = layer_idx + self.enabled = getattr(config, 'flextron', False) + + if not self.enabled: + return + + # Current elasticity parameters - store the full router outputs + self.current_router_emb = None + + # Hook handles for cleanup + self.hook_handles = [] + + def _init_embedding_masks(self): + """Initialize embedding dimension masks.""" + mask_list = [] + for emb_int in self.config.emb_int_list: + mask = torch.zeros(self.config.hidden_size, dtype=torch.bool) + mask[:emb_int] = True + mask_list.append(mask) + self.emb_masks_lookup = { + emb_int: idx for idx, emb_int in enumerate(self.config.emb_int_list) + } + self.emb_masks = torch.stack(mask_list, dim=0).to(device='cuda').to(dtype=torch.bfloat16) + + def attach_hooks(self, attention_module): + """Attach hooks to Attention following the original flextron_os pattern.""" + if not self.enabled: + return + + self.attention_module = attention_module + + emb_effective_per_list = [x / self.config.hidden_size for x in self.config.emb_int_list] + + # Setup hook - runs first to initialize masks for this forward pass + def setup_masks_hook(module, input): + if self.config.flextron: + self._init_embedding_masks() + return input + + # Cleanup hook - runs last to remove masks after forward pass + def cleanup_masks_hook(module, input, output): + if self.config.flextron: + self.emb_masks = None + self.emb_masks_lookup = {} + return output + + # IMPORTANT: Register setup hook FIRST + setup_handle = attention_module.register_forward_pre_hook(setup_masks_hook) + self.hook_handles.append(setup_handle) + + # Hook 1: Input masking and router_emb processing + def input_mask_hook(module, input): + if self.config.flextron and self.current_router_emb is not None: + hidden_states = input[0] + + # Apply embedding mask + if self.config.soft_mask: + soft_mask = torch.zeros( + self.emb_masks[0].shape, + dtype=torch.bfloat16, + device=self.emb_masks[0].device, + ) + for mask, per_logit in zip(self.emb_masks, self.current_router_emb[0]): + soft_mask.add_(mask * per_logit) + mask = soft_mask + masked_input = hidden_states * mask[None, None, :] + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + mask = self.emb_masks[self.emb_masks_lookup[emb_choice]] + masked_input = ( + hidden_states * mask.to(device=hidden_states.device)[None, None, :] + ) + masked_input = masked_input * router_emb_logits + + return tuple([masked_input] + list(input[1:])) + return input + + # Hook 2: Linear QKV pre-hook for eps modification + def linear_qkv_pre_hook(module, input): + if self.config.flextron and self.current_router_emb is not None: + # Set eps on linear_qkv (fused layernorm) + if self.config.soft_mask: + soft_eps = 0 + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_eps += self.config.layernorm_epsilon * emb_per * per_logit + module.eps = soft_eps.float().detach().item() + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + emb_effective_per = emb_choice / self.config.hidden_size + module.eps = self.config.layernorm_epsilon * emb_effective_per + + return input + + # Hook 3: Linear QKV post-hook for scaling + def linear_qkv_post_hook(module, input, output): + if self.config.flextron and self.current_router_emb is not None: + query_key_value, bias = output + if self.config.soft_mask: + soft_query_key_value = torch.zeros_like(query_key_value) + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_query_key_value.add_(query_key_value * (emb_per**0.5) * per_logit) + scaled_output = soft_query_key_value + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + emb_effective_per = emb_choice / self.config.hidden_size + scaled_output = query_key_value * router_emb_logits * (emb_effective_per**0.5) + module.eps = self.config.layernorm_epsilon + return (scaled_output, bias) + return output + + # Hook 5: Final output masking + def output_mask_hook(module, input, output): + if self.config.flextron and self.current_router_emb is not None: + out, out_bias = output + + # Apply embedding mask + if self.config.soft_mask: + soft_mask = torch.zeros( + self.emb_masks[0].shape, + dtype=torch.bfloat16, + device=self.emb_masks[0].device, + ) + for mask, per_logit in zip(self.emb_masks, self.current_router_emb[0]): + soft_mask.add_(mask * per_logit) + mask = soft_mask + masked_out = out * mask[None, None, :] + else: + router_emb_logits, emb_choice = ( + torch.max(self.current_router_emb[0]), + self.current_router_emb[1], + ) + mask = self.emb_masks[self.emb_masks_lookup[emb_choice]] + masked_out = out * mask[None, None, :] + masked_out = masked_out * router_emb_logits + + return (masked_out, out_bias) + return output + + # Hook 1: Input masking and router_emb processing + main_handle = attention_module.register_forward_pre_hook(input_mask_hook) + self.hook_handles.append(main_handle) + + # Hook 2&3: Linear QKV pre-hook for eps modification + qkv_pre_handle = attention_module.linear_qkv.register_forward_pre_hook(linear_qkv_pre_hook) + qkv_post_handle = attention_module.linear_qkv.register_forward_hook(linear_qkv_post_hook) + self.hook_handles.append(qkv_pre_handle) + self.hook_handles.append(qkv_post_handle) + + # Final output masking + output_handle = attention_module.register_forward_hook(output_mask_hook) + self.hook_handles.append(output_handle) + + # Cleanup hook - runs last to remove masks after forward pass + cleanup_handle = attention_module.register_forward_hook(cleanup_masks_hook) + self.hook_handles.append(cleanup_handle) + + def set_elasticity_params(self, router_emb=None, **kwargs): + """Set current elasticity parameters that will be used by hooks.""" + if router_emb is not None: + self.current_router_emb = router_emb + + def detach_hooks(self): + """Remove all hooks.""" + if not hasattr(self, 'hook_handles'): + return + for handle in self.hook_handles: + handle.remove() + self.hook_handles.clear() + + def __del__(self): + """Cleanup hooks when manager is destroyed.""" + self.detach_hooks() + + +class FlextronStackElasticityManager: + """ + Manages elasticity for HybridStack using pure PyTorch hooks. + Handles input masking and final norm scaling. + """ + + def __init__(self, config): + self.config = config + self.enabled = getattr(config, 'flextron', False) + + if not self.enabled: + return + + # Current elasticity parameters + self.current_emb_choice = self.config.hidden_size + self.current_router_emb = None + + # Hook handles for cleanup + self.hook_handles = [] + + # Pre-computed masks + self.emb_masks = None + self.emb_masks_lookup = {} + + def initialize_masks(self, stack): + """Initialize masks based on the stack configuration.""" + if not self.enabled: + return + + self.stack = stack + + def attach_hooks(self, stack): + """Attach hooks to HybridStack.""" + if not self.enabled: + return + + self.initialize_masks(stack) + + emb_effective_per_list = [x / self.config.hidden_size for x in self.config.emb_int_list] + + # Hook 1: Final norm pre-hook for eps modification + def final_norm_pre_hook(module, input): + if self.config.flextron and self.current_router_emb is not None: + # Modify eps for this forward pass + if self.config.soft_mask: + soft_eps = 0 + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_eps += self.config.layernorm_epsilon * emb_per * per_logit + module.eps = soft_eps.float().detach().item() + else: + emb_choice = self.current_emb_choice + emb_effective_per = emb_choice / self.config.hidden_size + module.eps = self.config.layernorm_epsilon * emb_effective_per + + return input + + # Hook 2: Final norm post-hook for scaling and eps restoration + def final_norm_post_hook(module, input, output): + if self.config.flextron and self.current_router_emb is not None: + # Restore original eps + module.eps = self.config.layernorm_epsilon + + # Apply scaling + if self.config.soft_mask: + soft_scaled_output = torch.zeros_like(output) + for emb_per, per_logit in zip( + emb_effective_per_list, self.current_router_emb[0] + ): + soft_scaled_output.add_(output * (emb_per**0.5) * per_logit) + scaled_output = soft_scaled_output + else: + emb_choice = self.current_emb_choice + emb_effective_per = emb_choice / self.config.hidden_size + router_emb_logits = torch.max(self.current_router_emb[0]) + scaled_output = output * (emb_effective_per**0.5) * router_emb_logits + return scaled_output + + return output + + # Hooks for final norm if it exists + final_norm_pre_handle = stack.final_norm.register_forward_pre_hook(final_norm_pre_hook) + final_norm_post_handle = stack.final_norm.register_forward_hook(final_norm_post_hook) + self.hook_handles.append(final_norm_pre_handle) + self.hook_handles.append(final_norm_post_handle) + + def set_elasticity_params(self, router_emb=None, **kwargs): + """Set current elasticity parameters that will be used by hooks.""" + if router_emb is not None: + self.current_router_emb = router_emb + self.current_emb_choice = router_emb[1] + + def detach_hooks(self): + """Remove all hooks and restore original forward method.""" + if not hasattr(self, 'hook_handles'): + return + for handle in self.hook_handles: + handle.remove() + self.hook_handles.clear() + + def __del__(self): + """Cleanup hooks when manager is destroyed.""" + self.detach_hooks() + + +def add_flextron_mamba_elasticity(mamba_mixer, config, layer_idx=0): + """ + Add elasticity to a MambaMixer using hooks. + + Args: + mamba_mixer: The MambaMixer instance to add elasticity to + config: Configuration object with flextron settings + layer_idx: Index of this layer in the hybrid pattern + + Returns: + FlextronMambaElasticityManager: Manager object to control elasticity + """ + if hasattr(mamba_mixer, '_flextron_manager'): + return mamba_mixer._flextron_manager + manager = FlextronMambaElasticityManager(config, layer_idx) + manager.attach_hooks(mamba_mixer) + + # Store manager reference on the mixer for easy access + mamba_mixer._flextron_manager = manager + + return manager + + +def add_flextron_transformer_layer_elasticity(transformer_layer, config, layer_idx=0): + """ + Add elasticity to a TransformerLayer using hooks. + + Args: + transformer_layer: The TransformerLayer instance to add elasticity to + config: Configuration object with flextron settings + layer_idx: Index of this layer in the hybrid pattern + + Returns: + FlextronTransformerLayerElasticityManager: Manager object to control elasticity + """ + if hasattr(transformer_layer, '_flextron_layer_manager'): + return transformer_layer._flextron_layer_manager + manager = FlextronTransformerLayerElasticityManager(config, layer_idx) + manager.attach_hooks(transformer_layer) + + # Store manager reference on the layer for easy access + transformer_layer._flextron_layer_manager = manager + + return manager + + +def add_flextron_topk_router_elasticity(router, config, layer_idx=0): + """ + Add elasticity to a TopKRouter using hooks. + + Args: + router: The TopKRouter instance to add elasticity to + config: Configuration object with flextron settings + layer_idx: Index of this layer in the hybrid pattern + + Returns: + FlextronTopKRouterElasticityManager: Manager object to control elasticity + """ + if hasattr(router, '_flextron_router_manager'): + return router._flextron_router_manager + manager = FlextronTopKRouterElasticityManager(config, layer_idx) + manager.attach_hooks(router) + + # Store manager reference on the router for easy access + router._flextron_router_manager = manager + + return manager + + +def add_flextron_moe_elasticity(moe_module, config, layer_idx=0): + """ + Add elasticity to a MoE using hooks. + + Args: + moe_module: The MoE instance to add elasticity to + config: Configuration object with flextron settings + layer_idx: Index of this layer in the hybrid pattern + + Returns: + FlextronMoEElasticityManager: Manager object to control elasticity + """ + if hasattr(moe_module, '_flextron_manager'): + return moe_module._flextron_manager + manager = FlextronMoEElasticityManager(config, layer_idx) + manager.attach_hooks(moe_module) + + # Store manager reference on the module for easy access + moe_module._flextron_manager = manager + + return manager + + +def add_flextron_grouped_mlp_elasticity(grouped_mlp_module, config, layer_idx=0): + """ + Add elasticity to a GroupedMLP using hooks. + """ + if hasattr(grouped_mlp_module, '_flextron_manager'): + return grouped_mlp_module._flextron_manager + manager = FlextronGroupedMLPElasticityManager(config, layer_idx) + manager.attach_hooks(grouped_mlp_module) + + # Store manager reference on the module for easy access + grouped_mlp_module._flextron_manager = manager + + return manager + + +def add_flextron_attention_elasticity(attention_module, config, layer_idx=0): + """ + Add elasticity to an Attention module using hooks. + + Args: + attention_module: The Attention instance to add elasticity to + config: Configuration object with flextron settings + layer_idx: Index of this layer in the hybrid pattern + + Returns: + FlextronAttentionElasticityManager: Manager object to control elasticity + """ + if hasattr(attention_module, '_flextron_manager'): + return attention_module._flextron_manager + manager = FlextronAttentionElasticityManager(config, layer_idx) + manager.attach_hooks(attention_module) + + # Store manager reference on the module for easy access + attention_module._flextron_manager = manager + + return manager + + +def add_flextron_stack_elasticity(stack, config): + """ + Add elasticity to a HybridStack using hooks. + + Args: + stack: The HybridStack instance to add elasticity to + config: Configuration object with flextron settings + + Returns: + FlextronStackElasticityManager: Manager object to control elasticity + """ + if hasattr(stack, '_flextron_manager'): + return stack._flextron_manager + manager = FlextronStackElasticityManager(config) + manager.attach_hooks(stack) + + # Store manager reference on the stack for easy access + stack._flextron_manager = manager + + return manager + + +# Convenience function to apply elasticity to all modules in a model +def apply_flextron_elasticity_to_model(model, config): + """Apply elasticity to all MambaMixer, MLP/MoE, and Attention instances in a model based on hybrid pattern.""" + managers = [] + + if not hasattr(config, 'hybrid_layer_pattern') or not config.hybrid_layer_pattern: + # No hybrid pattern, skip elasticity setup + return managers + + hybrid_pattern = config.hybrid_layer_pattern + + # Find decoder layers + decoder = getattr(model, 'decoder', None) + layers = getattr(decoder, 'layers', None) + + if decoder is None or layers is None: + return managers + + # Apply elasticity per layer based on hybrid pattern + for layer_idx, layer_char in enumerate(hybrid_pattern): + if layer_idx >= len(layers): + break + + layer = layers[layer_idx] + + if layer_char == 'E': # MoE layer (treated as MLP replacement) + if ( + 'MoETransformerLayer' == layer.__class__.__name__ + or 'TransformerLayer' == layer.__class__.__name__ + ): + layer_manager = add_flextron_transformer_layer_elasticity(layer, config, layer_idx) + managers.append(layer_manager) + + # Find MoELayer module in this layer + moe_module = None + for name, module in layer.named_modules(): + if 'MoELayer' == module.__class__.__name__: + moe_module = module + break + if moe_module is not None: + manager = add_flextron_moe_elasticity(moe_module, config, layer_idx) + managers.append(manager) + + # Also add router elasticity to the MoE router + router_module = None + for name, module in moe_module.named_modules(): + if 'TopKRouter' == module.__class__.__name__: + router_module = module + break + if router_module is not None: + router_manager = add_flextron_topk_router_elasticity( + router_module, config, layer_idx + ) + managers.append(router_manager) + + # Find TEGroupedMLP module in this layer + moe_module = None + for name, module in layer.named_modules(): + if 'TEGroupedMLP' == module.__class__.__name__: + moe_module = module + break + if moe_module is not None: + manager = add_flextron_grouped_mlp_elasticity(moe_module, config, layer_idx) + managers.append(manager) + + elif layer_char == 'M': # Mamba layer + mamba_module = None + for name, module in layer.named_modules(): + if 'MambaMixer' == module.__class__.__name__: + mamba_module = module + break + if mamba_module is not None: + manager = add_flextron_mamba_elasticity(mamba_module, config, layer_idx) + managers.append(manager) + + elif layer_char == '*': # Attention layer (TransformerLayer) + attention_module = None + for name, module in layer.named_modules(): + if 'SelfAttention' == module.__class__.__name__: + attention_module = module + break + if attention_module is not None: + manager = add_flextron_attention_elasticity(attention_module, config, layer_idx) + managers.append(manager) + + # Also add hooks to HybridStack if present + if hasattr(model, 'decoder') and hasattr(model.decoder, 'final_norm'): + stack_manager = add_flextron_stack_elasticity(model.decoder, config) + managers.append(stack_manager) + + # Store all managers on the model + model._flextron_managers = managers + return managers diff --git a/megatron/elastification/flextron_utils.py b/megatron/elastification/flextron_utils.py new file mode 100644 index 00000000000..c4fce4e9c5f --- /dev/null +++ b/megatron/elastification/flextron_utils.py @@ -0,0 +1,476 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +""" +Flextron Utilities + +Provides setup and configuration functions for Flextron elasticity. +Extracted from HybridModel to keep the core model clean. +""" + +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + +from megatron.core import mpu, parallel_state +from megatron.elastification.arguments import convert_per_lists_to_int_lists +from megatron.elastification.flextron_config import inject_flextron_config +from megatron.elastification.flextron_elasticity_hooks import apply_flextron_elasticity_to_model +from megatron.elastification.memory_config import MemoryConfig, load_memory_config +from megatron.elastification.router.flex_budget_utils import ( + get_memory_footprint, + get_num_parameters, +) +from megatron.elastification.router.hybrid_flex_router import FlextronRouter +from megatron.training import get_args + + +class FlextronModelManager: + """ + Manages Flextron functionality for a model. + Handles router, budget calculations, and loss functions. + """ + + def __init__(self, model, config): + self.model = model + self.config = config + inject_flextron_config(get_args(), config) + convert_per_lists_to_int_lists(config) + config.hybrid_layer_pattern = getattr(model, 'hybrid_layer_pattern', '') + self.router = None + self.budget_type = getattr(config, 'budget_type', 'param') + + # Load memory quantization profile from args + args = get_args() + self.memory_config = load_memory_config(args) + + # Budget calculation attributes + self.all_param = None + self.total_memory = None + + # Hook managers + self.hook_managers = [] + + def setup_router(self): + """Initialize the Flextron router if enabled.""" + if getattr(self.config, 'enable_router', False): # and self.model.pre_process: + self.router = FlextronRouter(config=self.config) + + # Make router name pipeline-stage-aware to avoid naming conflicts in PP>1 + pp_rank = mpu.get_pipeline_model_parallel_rank() + router_name = f"router_pp{pp_rank}" + + # Set the router with pipeline-specific name + setattr(self.model, router_name, self.router) + self.model.router = self.router + else: + self.model.router = None + + def setup_budget_functions(self): + """Setup budget calculation functions based on budget type.""" + self._setup_param_loss_func() + + if self.budget_type == 'mem': + self._setup_memory_loss_func() + + def setup_hooks(self): + """Setup elasticity hooks on the model.""" + if getattr(self.config, 'flextron', False): + self.hook_managers = apply_flextron_elasticity_to_model(self.model, self.config) + + def _setup_param_loss_func(self): + """Setup parameter counting for budget calculations.""" + + self.all_param, self.active_param = torch.tensor( + get_num_parameters( + hybrid_pattern=self.model.hybrid_layer_pattern, + mamba_num_heads=self.config.mamba_num_heads, + mamba_d_head=self.config.mamba_head_dim, + mamba_d_state=self.config.mamba_state_dim, + num_attention_heads=self.config.num_attention_heads, + num_query_groups=self.config.num_query_groups, + ffn_hidden_size=self.config.ffn_hidden_size, + hidden_size=self.config.hidden_size, + kv_channels=self.config.kv_channels, + vocab_size=self.model.vocab_size, + tied_vocab=self.model.share_embeddings_and_output_weights, + num_experts=self.config.num_moe_experts, + shared_expert_intermediate_size=self.config.moe_shared_expert_intermediate_size, + moe_router_topk=self.config.moe_router_topk, + ), + dtype=torch.float32, + device=torch.cuda.current_device(), + ) + + def _setup_memory_loss_func(self): + """Setup memory loss function by calculating the baseline memory footprint.""" + self.total_memory = ( + get_memory_footprint( + hybrid_pattern=self.model.hybrid_layer_pattern, + mamba_num_heads=self.config.mamba_num_heads, + mamba_d_head=self.config.mamba_head_dim, + mamba_d_state=self.config.mamba_state_dim, + num_attention_heads=self.config.num_attention_heads, + num_query_groups=self.config.num_query_groups, + ffn_hidden_size=self.config.ffn_hidden_size, + hidden_size=self.config.hidden_size, + kv_channels=self.config.kv_channels, + vocab_size=self.model.vocab_size, + tied_vocab=self.model.share_embeddings_and_output_weights, + mem_infer_seq_len=self.config.mem_infer_seq_len, + mem_batch_size=self.config.mem_batch_size, + prefill_chunk_size=self.config.prefill_chunk_size, + moe_num_experts=self.config.num_moe_experts, + shared_expert_intermediate_size=self.config.moe_shared_expert_intermediate_size, + moe_router_topk=self.config.moe_router_topk, + memory_config=self.memory_config, + ) + .float() + .to(torch.cuda.current_device()) + ) + + print( + f"Total baseline memory footprint: {self.total_memory.item():.4f} GB " + f"(profile={getattr(get_args(), 'memory_profile', 'bf16')}, " + f"param_target={self.memory_config.param_budget_target})" + ) + + def budget_loss_func(self, flextron_kwargs, budget_item=0): + """Calculate budget-based loss exactly as in the original implementation.""" + dtype, device = ( + flextron_kwargs['router_mlp'][0].dtype, + flextron_kwargs['router_mlp'][0].device, + ) + + flex_mamba_num_head = flextron_kwargs['router_mamba'][0] @ torch.tensor( + self.config.mamba_int_list, dtype=dtype, device=device + ) + flex_hidden_size = flextron_kwargs['router_emb'][0] @ torch.tensor( + self.config.emb_int_list, dtype=dtype, device=device + ) + flex_ffn_hidden_size = flextron_kwargs['router_mlp'][0] @ torch.tensor( + self.config.mlp_int_list, dtype=dtype, device=device + ) + flex_moe_expert = flextron_kwargs['router_moe_expert'][0] @ torch.tensor( + self.config.moe_expert_int_list, dtype=dtype, device=device + ) + # Attention heads are not router-controlled; pass the parent value through. + num_attention_heads = self.config.num_attention_heads + + if self.config.add_skipping: + logit_skip_selected = torch.cumsum(flextron_kwargs['router_skip'][0], 0)[:-1] + logit_skip_all = torch.ones(self.config.num_layers).to(dtype=dtype, device=device) + logit_skip_all[self.config.layer_ranking_list] = logit_skip_selected + + mamba_idxs = [ + i for i, char in enumerate(self.model.hybrid_layer_pattern) if char == 'M' + ] + mamba_idxs = torch.tensor(mamba_idxs, dtype=torch.long) + flex_mamba_num_head = flex_mamba_num_head * logit_skip_all[mamba_idxs] + flex_mamba_num_head = flex_mamba_num_head.unsqueeze(-1) + + head_idxs = [i for i, char in enumerate(self.model.hybrid_layer_pattern) if char == '*'] + head_idxs = torch.tensor(head_idxs, dtype=torch.long) + num_attention_heads = num_attention_heads * logit_skip_all[head_idxs] + num_attention_heads = num_attention_heads.unsqueeze(-1) + + moe_idxs = [i for i, char in enumerate(self.model.hybrid_layer_pattern) if char == 'E'] + moe_idxs = torch.tensor(moe_idxs, dtype=torch.long) + flex_ffn_hidden_size = flex_ffn_hidden_size * logit_skip_all[moe_idxs] + flex_ffn_hidden_size = flex_ffn_hidden_size.unsqueeze(-1) + + flex_moe_expert = flex_moe_expert * logit_skip_all[moe_idxs] + flex_moe_expert = flex_moe_expert.unsqueeze(-1) + + if not self.config.flex_hetero_ffn and not self.config.add_skipping: + flex_ffn_hidden_size = flex_ffn_hidden_size.unsqueeze(-1) + if not self.config.flex_hetero_mamba and not self.config.add_skipping: + flex_mamba_num_head = flex_mamba_num_head.unsqueeze(-1) + if not self.config.flex_hetero_moe_expert and not self.config.add_skipping: + flex_moe_expert = flex_moe_expert.unsqueeze(-1) + + current_param_all, current_param_active = get_num_parameters( + hybrid_pattern=self.model.hybrid_layer_pattern, + mamba_num_heads=flex_mamba_num_head.float(), + mamba_d_head=self.config.mamba_head_dim, + mamba_d_state=self.config.mamba_state_dim, + num_attention_heads=( + num_attention_heads.float() + if isinstance(num_attention_heads, torch.Tensor) + else num_attention_heads + ), + num_query_groups=self.config.num_query_groups, + ffn_hidden_size=flex_ffn_hidden_size.float(), + hidden_size=flex_hidden_size.unsqueeze(-1).float(), + kv_channels=self.config.kv_channels, + vocab_size=self.model.vocab_size, + tied_vocab=self.model.share_embeddings_and_output_weights, + num_experts=flex_moe_expert.float(), + shared_expert_intermediate_size=self.config.moe_shared_expert_intermediate_size, + moe_router_topk=self.config.moe_router_topk, + ) + + + if self.config.budget_type == 'param': + if self.memory_config.param_budget_target == 'active': + diff = abs(current_param_active / (budget_item * self.active_param) - 1) + else: + diff = abs(current_param_all / (budget_item * self.all_param) - 1) + elif self.config.budget_type == 'mem': + current_mem = get_memory_footprint( + hybrid_pattern=self.model.hybrid_layer_pattern, + mamba_num_heads=flex_mamba_num_head.float(), + mamba_d_head=self.config.mamba_head_dim, + mamba_d_state=self.config.mamba_state_dim, + num_attention_heads=( + num_attention_heads.float() + if isinstance(num_attention_heads, torch.Tensor) + else num_attention_heads + ), + num_query_groups=self.config.num_query_groups, + ffn_hidden_size=flex_ffn_hidden_size.float(), + hidden_size=flex_hidden_size.unsqueeze(-1).float(), + kv_channels=self.config.kv_channels, + vocab_size=self.model.vocab_size, + tied_vocab=self.model.share_embeddings_and_output_weights, + mem_infer_seq_len=self.config.mem_infer_seq_len, + mem_batch_size=self.config.mem_batch_size, + prefill_chunk_size=self.config.prefill_chunk_size, + moe_num_experts=flex_moe_expert.float(), + shared_expert_intermediate_size=self.config.moe_shared_expert_intermediate_size, + moe_router_topk=self.config.moe_router_topk, + memory_config=self.memory_config, + ).float() + diff = abs(current_mem / budget_item - 1) + else: + raise ValueError(f"Invalid budget type: {self.config.budget_type}") + + # return current_param, {} + if budget_item != 1.0 and diff < 0.05: + diff = diff * 0.0 + + # if getattr(self.config, 'disable_budget', False): + # diff = diff * 0.0 + + if budget_item == 1.0: + if self.config.flex_hetero_moe_expert: + label_moe_expert = torch.zeros_like(flextron_kwargs['router_moe_expert'][0]) + label_moe_expert[:, 0] = 1.0 + mse_loss_moe_expert = F.mse_loss( + flextron_kwargs['router_moe_expert'][0], label_moe_expert + ) + else: + label_moe_expert = torch.zeros_like(flextron_kwargs['router_moe_expert'][0]) + label_moe_expert[0] = 1.0 + mse_loss_moe_expert = F.mse_loss( + flextron_kwargs['router_moe_expert'][0], label_moe_expert + ) + + if self.config.flex_hetero_mamba: + label_mamba = torch.zeros_like(flextron_kwargs['router_mamba'][0]) + label_mamba[:, 0] = 1.0 + mse_loss_mamba = F.mse_loss(flextron_kwargs['router_mamba'][0], label_mamba) + else: + label_mamba = torch.zeros_like(flextron_kwargs['router_mamba'][0]) + label_mamba[0] = 1.0 + mse_loss_mamba = F.mse_loss(flextron_kwargs['router_mamba'][0], label_mamba) + + if self.config.flex_hetero_ffn: + label_mlp = torch.zeros_like(flextron_kwargs['router_mlp'][0]) + label_mlp[:, 0] = 1.0 + mse_loss_mlp = F.mse_loss(flextron_kwargs['router_mlp'][0], label_mlp) + else: + label_mlp = torch.zeros_like(flextron_kwargs['router_mlp'][0]) + label_mlp[0] = 1.0 + mse_loss_mlp = F.mse_loss(flextron_kwargs['router_mlp'][0], label_mlp) + + if self.config.add_skipping: + label_skip = torch.zeros_like(flextron_kwargs['router_skip'][0]) + label_skip[0] = 1.0 + mse_loss_skip = F.mse_loss(flextron_kwargs['router_skip'][0], label_skip) + else: + mse_loss_skip = 0.0 + + label_emb = torch.zeros_like(flextron_kwargs['router_emb'][0]) + label_emb[0] = 1.0 + mse_loss_emb = F.mse_loss(flextron_kwargs['router_emb'][0], label_emb) + + diff += 10 * ( + mse_loss_mamba + + mse_loss_mlp + + mse_loss_moe_expert + + mse_loss_skip + + mse_loss_emb + ) + + return diff.bfloat16(), {} + + def get_loss_func(self): + """Get the budget loss function.""" + return self.budget_loss_func + + def process_router_output(self, budget_item): + """Process router output and return flextron_kwargs.""" + if self.router is None: + return {}, None + + (router_mlp, router_skip, router_emb, router_mamba, router_moe_expert) = ( + self.router(budget_item) + ) + + flextron_kwargs = { + 'router_mlp': router_mlp, + 'router_skip': router_skip, + 'router_emb': router_emb, + 'router_mamba': router_mamba, + 'router_moe_expert': router_moe_expert, + } + + return flextron_kwargs, self.get_loss_func() + + def update_hook_elasticity_params(self, flextron_kwargs): + """Update elasticity parameters in all hook managers.""" + if not self.hook_managers: + return + + # Extract elasticity parameters from router outputs + router_emb = flextron_kwargs.get('router_emb') + router_mamba = flextron_kwargs.get('router_mamba') + router_mlp = flextron_kwargs.get('router_mlp') + router_moe_expert = flextron_kwargs.get('router_moe_expert') + router_skip = flextron_kwargs.get('router_skip') # General layer skipping + + # Update all hook managers with router outputs directly + for manager in self.hook_managers: + if hasattr(manager, 'set_elasticity_params'): + manager.set_elasticity_params( + router_emb=router_emb, + router_mamba=router_mamba, + router_mlp=router_mlp, + router_moe_expert=router_moe_expert, + router_skip=router_skip, + ) + + +def setup_flextron_model(model): + """ + Setup Flextron functionality for a model after creation. + + Args: + model: The HybridModel instance + + Returns: + FlextronModelManager: Manager object to handle Flextron operations + """ + manager = FlextronModelManager(model, model.config) + + # Setup all Flextron components + manager.setup_router() + manager.setup_budget_functions() + manager.setup_hooks() + + # Store manager on model for easy access + model._flextron_manager = manager + + return manager + + +def inject_flextron_forward_logic(model): + """ + Inject Flextron-specific forward pass logic into the model. + This replaces the router logic that was previously in HybridModel.forward(). + """ + original_forward = model.forward + + def flextron_forward( + self, + input_ids, + position_ids, + attention_mask, + decoder_input=None, + labels=None, + inference_context=None, + runtime_gather_output=None, + *, + inference_params=None, + **flextron_kwargs, + ): + + # Handle override budget settings + if getattr(self.config, 'override_selected_budget', None) is not None: + assert ( + self.config.is_flex_eval + ), "Override selected budget should only be set in flex eval mode" + # Both branches must populate the 'budget' key — downstream code + # at line 422 reads it unconditionally. Setting budget=1.0 routes + # the override-1.0 case through the regular router-forward path, + # which after training produces near-identity router outputs that + # mask down to the full model. + flextron_kwargs = {'budget': self.config.override_selected_budget[0]} + + # Initialize budget_loss + budget_loss = None + + # Handle router logic if enabled and model has Flextron manager + if ( + hasattr(self, '_flextron_manager') + and self._flextron_manager is not None + and self._flextron_manager.router is not None + ): + # Every step is router-driven, including budget=1.0 (which now + # propagates the identity-MSE regularization that the old + # ``original_model`` kill-switch silently zeroed out). Use + # ``freeze_router`` if you need to train without router gradients. + budget_item = flextron_kwargs['budget'] + + # Get router output and loss function + flextron_kwargs, loss_func = self._flextron_manager.process_router_output(budget_item) + + # Calculate loss + if loss_func: + budget_loss = loss_func(flextron_kwargs, budget_item) + + # Push router outputs into the elasticity hook managers so masks + # fire on this forward. + self._flextron_manager.update_hook_elasticity_params(flextron_kwargs) + else: + # If no Flextron manager, clear flextron_kwargs to avoid passing unknown args + flextron_kwargs = {} + + # Call original forward with processed flextron_kwargs + result = original_forward( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=decoder_input, + labels=labels, + inference_context=inference_context, + runtime_gather_output=runtime_gather_output, + inference_params=inference_params, + ) + + # Handle return values based on training mode and flextron settings + if labels is not None: + loss = result if not isinstance(result, tuple) else result[0] + + if ( + hasattr(self, '_flextron_manager') + and self._flextron_manager is not None + and self._flextron_manager.router is not None + and getattr(self.config, 'flextron', False) + and not getattr(self.config, 'is_flex_eval', False) + ): + if mpu.is_pipeline_last_stage(): + return loss, budget_loss + else: + return loss + else: + # Evaluation mode or non-flextron, return loss only + return loss + else: + # No labels, return logits + return result + + # Replace the forward method + model.forward = flextron_forward.__get__(model, model.__class__) diff --git a/megatron/elastification/loss_func.py b/megatron/elastification/loss_func.py new file mode 100644 index 00000000000..cd0be6df67c --- /dev/null +++ b/megatron/elastification/loss_func.py @@ -0,0 +1,210 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +"""Flextron loss function(s). + +Combines lm loss with the router's budget loss, optional KD distillation +loss, and per-budget reporting (full-model vs sub-budget breakdown). +""" + +import torch + +from megatron.core import parallel_state +from megatron.core.models.gpt import GPTModel +from megatron.training import get_args +from megatron.training.utils import unwrap_model + + +def _mask_loss(output_tensor, loss_mask): + """Apply mask to the unreduced loss tensor.""" + args = get_args() + if isinstance(output_tensor, tuple) and len(output_tensor) == 2: + (output_tensor, (param_loss, extra_reporting_dict)) = output_tensor + tp_reduce, is_sequence_parallel = False, False + elif isinstance(output_tensor, tuple): + # Special distillation flags indicating whether to perform additional tensor-parallel adjustments. + output_tensor, tp_reduce, is_sequence_parallel = output_tensor + param_loss = None + else: + tp_reduce, is_sequence_parallel = False, False + param_loss = None + + num_tokens = loss_mask.sum().float() + + if param_loss is not None: + if param_loss > 0: + pass + else: + param_loss = -args.router_beta * param_loss + + if is_sequence_parallel: + # Sequence-parallel tensor derived from intermediate activation - need to split loss mask. + idx = parallel_state.get_tensor_model_parallel_rank() + loss_mask = torch.tensor_split(loss_mask, args.tensor_model_parallel_size, dim=1)[idx] + + losses = output_tensor.view(-1).float() + loss_mask = loss_mask.reshape(-1).float() + loss = torch.sum(losses * loss_mask) + + alpha = args.loss_alpha + if not args.freeze_router and param_loss is not None: + param_loss_item = param_loss[0] * num_tokens * alpha + # add param loss to lm loss + loss += param_loss_item + else: + param_loss_item = None + + if tp_reduce or is_sequence_parallel: + # Losses on parallel tensors require extra all-reduce to sync across MP ranks. + torch.distributed.all_reduce(loss, group=parallel_state.get_tensor_model_parallel_group()) + + if param_loss_item is not None: + return loss, param_loss_item + else: + return loss + + +def loss_func( + loss_mask: torch.Tensor, + output_tensor: torch.Tensor, + model: GPTModel, + selected_budget: float = None, +): + """Loss function (with KD Loss support). + + Args: + loss_mask (Tensor): Used to mask out some portions of the loss + output_tensor (Tensor): The tensor with the losses + model (GPTModel): The model (can be wrapped) + selected_budget (float): The budget value used for this forward pass + """ + args = get_args() + + # Unwrap for both Distillation and LANA + model = unwrap_model(model) + + # Standard lm loss + out_mask_loss = _mask_loss(output_tensor, loss_mask) + + if isinstance(out_mask_loss, tuple): + loss_lm, param_loss_item = out_mask_loss + else: + # assert args.freeze_router, "Param loss None is not supported without freezing router" + loss_lm = out_mask_loss + param_loss_item = torch.tensor(0.0, device=loss_lm.device, dtype=loss_lm.dtype) + + loss = loss_lm + num_tokens = loss_mask.sum().clone().detach().to(torch.int) + # Protect against division by zero when all tokens are masked. + num_tokens = torch.clamp(num_tokens, min=1) + # Report (value, num_tokens) as local-rank values; the training loop performs the + # DP+CP all-reduce on report-dict tuples (training.py: token-weighted reduction). + report = { + 'lm loss': ((loss_lm.detach() - param_loss_item.detach()).view(1), num_tokens), + 'param loss item': (param_loss_item.detach().view(1), num_tokens), + } + + # Add per-model LM loss breakdown for logging only when KD is NOT active + kd_active = model.training and args.export_kd_teacher_load + if not kd_active: + try: + is_full_model = (param_loss_item is None) or (param_loss_item.detach().abs() == 0) + except Exception: + is_full_model = False + zero_num = torch.zeros_like(report['lm loss'][0]) + zero_den = torch.zeros_like(num_tokens) + if is_full_model: + report['lm loss (full)'] = report['lm loss'] + report['lm loss (budget)'] = (zero_num, zero_den) + else: + report['lm loss (budget)'] = report['lm loss'] + report['lm loss (full)'] = (zero_num, zero_den) + + if model.training and args.export_kd_teacher_load: + # [ModelOpt]: Handle knowledge distillation. + # The installed balancer with skip_lm_loss=True drops student_loss (param_loss) from + # the total. Add loss_lm back manually to restore the router gradient signal. + losses = model.compute_kd_loss( + student_loss=loss_lm, loss_reduction_fn=lambda x: _mask_loss(x, loss_mask) + ) + loss = losses["kd_loss"] + param_loss_item + # All-gather logits_loss across DP ranks so we can mask by selected_budget below. + logits_loss = losses["logits_loss"].detach() + dp_world_size = torch.distributed.get_world_size( + group=parallel_state.get_data_parallel_group() + ) + logits_loss_gathered = [torch.zeros_like(logits_loss) for _ in range(dp_world_size)] + torch.distributed.all_gather( + logits_loss_gathered, logits_loss, group=parallel_state.get_data_parallel_group() + ) + logits_loss_gathered = torch.stack(logits_loss_gathered) + + total_loss_report = losses["kd_loss"].detach() + param_loss_item.detach() + report["total loss"] = (total_loss_report, num_tokens) + + # Log KD loss split into full vs budget similar to LM loss breakdown. + try: + is_full_model_kd = (param_loss_item is None) or (param_loss_item.detach().abs() == 0) + except Exception: + is_full_model_kd = False + zero_num_kd = torch.zeros_like(total_loss_report) + zero_den_kd = torch.zeros_like(num_tokens) + if is_full_model_kd: + report["kd loss (full)"] = (total_loss_report, num_tokens) + report["kd loss (budget)"] = (zero_num_kd, zero_den_kd) + else: + report["kd loss (budget)"] = (total_loss_report, num_tokens) + report["kd loss (full)"] = (zero_num_kd, zero_den_kd) + report["logits distillation loss"] = (losses["logits_loss"].detach(), num_tokens) + report["intermediate distillation loss"] = ( + losses["intermediate_loss"].detach(), + num_tokens, + ) + + local_budget = torch.tensor( + [selected_budget], dtype=torch.float32, device=logits_loss.device + ) + budgets_gathered = [torch.zeros_like(local_budget) for _ in range(dp_world_size)] + torch.distributed.all_gather( + budgets_gathered, local_budget, group=parallel_state.get_data_parallel_group() + ) + budgets_gathered = torch.cat(budgets_gathered) + + # Create a binary mask where gathered budgets are equal to selected_budget (with 1e-6 tolerance) + budget_mask = (budgets_gathered - selected_budget).abs() < 1e-6 + logits_loss_gathered_selected = logits_loss_gathered[budget_mask].sum() / budget_mask.sum() + budget_num_tokens = ( + num_tokens.float() * budget_mask.sum() / budget_mask.shape[0] / budget_mask.sum() + ) + + corrected_budget_list = list(set(args.budget_list)) + + for temp_budget in corrected_budget_list: + report[f"logits distillation loss {temp_budget:.3f}"] = ( + torch.tensor(0.0, device=logits_loss.device, dtype=torch.float32), + torch.tensor(0.0, device=logits_loss.device, dtype=torch.float32), + ) + index_of_selected_budget = corrected_budget_list.index(selected_budget) + all_budget_logit = torch.zeros( + len(corrected_budget_list), device=logits_loss.device, dtype=logits_loss.dtype + ) + all_budget_tokens = torch.zeros( + len(corrected_budget_list), device=logits_loss.device, dtype=logits_loss.dtype + ) + + all_budget_logit[index_of_selected_budget] = logits_loss_gathered_selected + all_budget_tokens[index_of_selected_budget] = budget_num_tokens + + for i in range(len(corrected_budget_list)): + report[f"logits distillation loss {corrected_budget_list[i]:.3f}"] = ( + all_budget_logit[i], + all_budget_tokens[i], + ) + + # Convert all items in report dict to a single (value, num_tokens) tensor. + for key, val in report.items(): + assert isinstance(val, tuple), "Value is not a tuple" + report[key] = torch.tensor( + [val[0], val[1].view(1)], device=loss_lm.device, dtype=loss_lm.dtype + ) + + return loss, num_tokens, report diff --git a/megatron/elastification/memory_config.py b/megatron/elastification/memory_config.py new file mode 100644 index 00000000000..3ac5b7132e6 --- /dev/null +++ b/megatron/elastification/memory_config.py @@ -0,0 +1,140 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +""" +memory_config.py — MemoryConfig dataclass and loader for Flextron budget calculations. + +Usage +----- +From CLI args (in training/eval scripts): + cfg = load_memory_config(args) + total_gb = get_memory_footprint(..., memory_config=cfg) + +Directly (in tests or notebooks): + cfg = MemoryConfig(bpe_kv_cache=1, param_budget_target='active') +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +import yaml + +# Path to the bundled presets file (same directory as this module). +_DEFAULT_PROFILES_PATH = os.path.join(os.path.dirname(__file__), "memory_profiles.yaml") + + +@dataclass +class MemoryConfig: + """ + Bytes-per-element for each memory component and param budget supervision target. + + Attributes + ---------- + bpe_params : float + Bytes per weight parameter element (2 = BF16, 1 = FP8/INT8, 0.5625 = FP4). + bpe_kv_cache : float + Bytes per KV-cache element. + bpe_ssm_cache : float + Bytes per SSM-state element (covers both conv_state and ssm_state). + bpe_max_buffer : float + Bytes per MoE dispatch buffer element. + param_budget_target : str + Whether the param-budget loss supervises on ``'active'`` (top-k experts only) + or ``'total'`` (all parameters including non-active experts) parameter count. + """ + + bpe_params: float = 2.0 + bpe_kv_cache: float = 2.0 + bpe_ssm_cache: float = 2.0 + bpe_max_buffer: float = 2.0 + param_budget_target: str = "active" # "active" | "total" + + def __post_init__(self): + valid_targets = {"active", "total"} + if self.param_budget_target not in valid_targets: + raise ValueError( + f"param_budget_target must be one of {valid_targets}, " + f"got '{self.param_budget_target}'" + ) + + +def load_memory_config(args) -> MemoryConfig: + """ + Build a MemoryConfig from parsed CLI args. + + Resolution order (highest wins): + 1. Individual override args (--bpe-params, --bpe-kv-cache, …) + 2. Named preset from YAML (--memory-profile ) + 3. Built-in defaults (BF16 everywhere, active param target) + + Parameters + ---------- + args : argparse.Namespace + Parsed arguments. Relevant attributes (all optional): + memory_profile str — preset name (default: 'bf16') + memory_profile_path str — path to YAML profiles file + bpe_params float — override + bpe_kv_cache float — override + bpe_ssm_cache float — override + bpe_max_buffer float — override + param_budget_target str — override ('active' | 'total') + """ + cfg = MemoryConfig() + + # ── Load preset from YAML ────────────────────────────────────────────── + profile_name = getattr(args, "memory_profile", "bf16") or "bf16" + profile_path = getattr(args, "memory_profile_path", None) or _DEFAULT_PROFILES_PATH + print(f"[memory_config] profile='{profile_name}' path={profile_path}") + + if not os.path.isfile(profile_path): + raise FileNotFoundError(f"Memory profiles file not found: {profile_path}") + + with open(profile_path) as f: + profiles = yaml.safe_load(f) + + presets = profiles.get("presets", {}) + if profile_name not in presets: + available = list(presets.keys()) + raise ValueError( + f"Memory profile '{profile_name}' not found in {profile_path}. " + f"Available: {available}" + ) + + preset = presets[profile_name] + cfg.bpe_params = float(preset.get("params", cfg.bpe_params)) + cfg.bpe_kv_cache = float(preset.get("kv_cache", cfg.bpe_kv_cache)) + cfg.bpe_ssm_cache = float(preset.get("ssm_cache", cfg.bpe_ssm_cache)) + cfg.bpe_max_buffer = float(preset.get("max_buffer", cfg.bpe_max_buffer)) + cfg.param_budget_target = preset.get("param_budget_target", cfg.param_budget_target) + print( + f"[memory_config] after preset : bpe_params={cfg.bpe_params} bpe_kv_cache={cfg.bpe_kv_cache} " + f"bpe_ssm_cache={cfg.bpe_ssm_cache} bpe_max_buffer={cfg.bpe_max_buffer} " + f"param_budget_target={cfg.param_budget_target}" + ) + + # ── Apply individual CLI overrides (take priority over preset) ───────── + if getattr(args, "bpe_params", None) is not None: + print(f"[memory_config] override bpe_params: {cfg.bpe_params} -> {args.bpe_params}") + cfg.bpe_params = float(args.bpe_params) + if getattr(args, "bpe_kv_cache", None) is not None: + print(f"[memory_config] override bpe_kv_cache: {cfg.bpe_kv_cache} -> {args.bpe_kv_cache}") + cfg.bpe_kv_cache = float(args.bpe_kv_cache) + if getattr(args, "bpe_ssm_cache", None) is not None: + print( + f"[memory_config] override bpe_ssm_cache: {cfg.bpe_ssm_cache} -> {args.bpe_ssm_cache}" + ) + cfg.bpe_ssm_cache = float(args.bpe_ssm_cache) + if getattr(args, "bpe_max_buffer", None) is not None: + print( + f"[memory_config] override bpe_max_buffer: {cfg.bpe_max_buffer} -> {args.bpe_max_buffer}" + ) + cfg.bpe_max_buffer = float(args.bpe_max_buffer) + if getattr(args, "param_budget_target", None) is not None: + print( + f"[memory_config] override param_budget_target: {cfg.param_budget_target} -> {args.param_budget_target}" + ) + cfg.param_budget_target = args.param_budget_target + + print(f"[memory_config] final : {cfg}") + return cfg diff --git a/megatron/elastification/memory_profiles.yaml b/megatron/elastification/memory_profiles.yaml new file mode 100644 index 00000000000..f060ce5d0cb --- /dev/null +++ b/megatron/elastification/memory_profiles.yaml @@ -0,0 +1,59 @@ +# memory_profiles.yaml +# +# Named memory quantization profiles for Flextron budget calculations. +# Each preset specifies bytes-per-element (bpe) for each memory component +# and whether the param budget loss targets total or active parameters. +# +# Select a preset via: --memory-profile +# Override a single value via: --bpe-kv-cache 1 (takes priority over preset) +# Use a custom file via: --memory-profile-path /path/to/file.yaml + +presets: + + # ── Standard BF16 inference (current implementation) ────────────────────── + bf16: + params: 2 # BF16 weights + kv_cache: 2 # BF16 KV cache + ssm_cache: 2 # BF16 SSM state (conv + ssm) + max_buffer: 2 # BF16 MoE dispatch buffer + param_budget_target: active + + # ── FP8 KV cache, BF16 weights (common serving optimisation) ────────────── + fp8_kv: + params: 2 # BF16 weights + kv_cache: 1 # FP8 KV cache + ssm_cache: 2 # BF16 SSM state + max_buffer: 2 # BF16 MoE dispatch buffer + param_budget_target: active + + # ── FP8 KV + SSM cache ──────────────────────────────────────────────────── + fp8_kv_ssm: + params: 2 # BF16 weights + kv_cache: 1 # FP8 KV cache + ssm_cache: 1 # FP8/INT8 SSM state + max_buffer: 2 # BF16 MoE dispatch buffer + param_budget_target: active + + # ── Fully quantised FP8 inference ───────────────────────────────────────── + fp8_all: + params: 1 # FP8 weights + kv_cache: 1 # FP8 KV cache + ssm_cache: 1 # FP8 SSM state + max_buffer: 1 # FP8 MoE dispatch buffer + param_budget_target: active + + # ── INT8 weights + INT8 caches ──────────────────────────────────────────── + int8: + params: 1 # INT8 weights + kv_cache: 1 # INT8 KV cache + ssm_cache: 1 # INT8 SSM state + max_buffer: 2 # BF16 MoE dispatch buffer (activation, usually not quantised) + param_budget_target: active + + # ── FP4 (speculative / future) ──────────────────────────────────────────── + fp4: + params: 0.5625 # FP4 weights + kv_cache: 0.5625 # FP4 KV cache + ssm_cache: 1 # INT8 SSM state + max_buffer: 2 # BF16 MoE dispatch buffer + param_budget_target: active diff --git a/megatron/elastification/pretrain_hybrid_flex.py b/megatron/elastification/pretrain_hybrid_flex.py new file mode 100644 index 00000000000..64b909d72d2 --- /dev/null +++ b/megatron/elastification/pretrain_hybrid_flex.py @@ -0,0 +1,542 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +"""Pretrain and SFT Mamba.""" + +import os +from functools import partial +from typing import List, Optional, Tuple, Union + +import torch + +from megatron.core import mpu, parallel_state +from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder +from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset +from megatron.core.enums import ModelType +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.num_microbatches_calculator import ( + get_current_global_batch_size, + get_micro_batch_size, +) +from megatron.core.parallel_state import ( + get_context_parallel_rank, + get_context_parallel_world_size, + get_data_parallel_rank, + get_data_parallel_world_size, + get_pipeline_model_parallel_rank, + get_pipeline_model_parallel_world_size, + get_tensor_model_parallel_group, + get_tensor_model_parallel_rank, +) +from megatron.core.rerun_state_machine import get_rerun_state_machine +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.spec_utils import import_module +from megatron.core.utils import StragglerDetector +from megatron.elastification.arguments import add_flextron_args +from megatron.training import ( + get_args, + get_timers, + get_tokenizer, + inprocess_restart, + pretrain, + print_rank_0, +) +from megatron.training.argument_utils import pretrain_cfg_container_from_args +from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args +from megatron.training.datasets.sft_dataset import SFTDataset +from megatron.training.utils import ( + get_batch_on_this_cp_rank, + get_batch_on_this_tp_rank, + get_blend_and_blend_per_split, +) + +# modelopt distillation +try: + from megatron.elastification.loss_func import loss_func as loss_func_modelopt + from megatron.post_training.arguments import add_modelopt_args + from megatron.post_training.model_builder import ( + modelopt_gpt_mamba_builder as model_provider_modelopt, + ) + has_nvidia_modelopt = True +except ImportError: + print_rank_0("ModelOpt is not installed. Please install it using `pip install nvidia-modelopt`") + has_nvidia_modelopt = False +print_rank_0("has_nvidia_modelopt is {}".format(has_nvidia_modelopt)) +import numpy as np + +try: + # Register the TE CUDA kernels + import transformer_engine # pylint: disable=unused-import + + # Alias the PyTorch wrapper so we can call tex.* APIs + import transformer_engine_torch as tex +except ImportError: + # TE isn’t installed or the torch wrapper is missing + tex = None + +from megatron.core.utils import is_te_min_version + +_global_choice_counter = 0 +_logged_params_norm = False + +stimer = StragglerDetector() + +def count_parameters_in_layer(model, layer_name): + num_params = 0 + for name, param in model.named_parameters(): + if layer_name in name: + num_params += param.numel() + print_rank_0(f" - {name}: {param.numel()}") + return num_params + + +def model_provider(pre_process=True, post_process=True, vp_stage: Optional[int] = None, config = None, pg_collection = None) -> HybridModel: + """Builds the model. + + Args: + pre_process (bool, optional): Set to true if you need to compute embeddings. Defaults to True. + post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True. + + + Returns: + HybridModel: The returned model + """ + args = get_args() + if has_nvidia_modelopt: + + model = model_provider_modelopt(args, pre_process, post_process, vp_stage=vp_stage, config=config, pg_collection=pg_collection) + from megatron.elastification.flextron_utils import ( + inject_flextron_forward_logic, + setup_flextron_model, + ) + setup_flextron_model(model) + inject_flextron_forward_logic(model) + + if args.freeze_model: + for name, param in model.named_parameters(): + if 'gate' not in name: + param.requires_grad = False + + if args.freeze_router: + for name, param in model.named_parameters(): + if 'gate' in name: + param.requires_grad = False + + return model + + print_rank_0('building Mamba model ...') + config = core_transformer_config_from_args(args, TransformerConfig) + + assert args.use_legacy_models == False, "Mamba only supported in Mcore!" + + if args.spec is not None: + hybrid_stack_spec = import_module(args.spec) + else: + raise ValueError("You must provide a valid Mamba layer spec!") + + model = HybridModel( + config=config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + hybrid_layer_pattern=args.hybrid_layer_pattern, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + rotary_percent=args.rotary_percent, + rotary_base=args.rotary_base, + vp_stage=vp_stage + ) + from megatron.elastification.flextron_utils import ( + inject_flextron_forward_logic, + setup_flextron_model, + ) + setup_flextron_model(model) + inject_flextron_forward_logic(model) + + for l in range(model.decoder.num_layers_per_pipeline_rank): + layer_params = count_parameters_in_layer(model, f'decoder.layers.{l}.') + print_rank_0(f" == params layer {l}: {layer_params}") + + return model + + +def get_batch(data_iterator): + """Generate a batch.""" + + # TODO: this is pretty hacky, find a better way + if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()): + return None, None, None, None, None, None, None + + # get batches based on the TP rank you are on + batch = get_batch_on_this_tp_rank(data_iterator) + + cu_seqlens = batch['cu_seqlens'] + if cu_seqlens is None: + # slice batch along sequence dimension for context parallelism + batch = get_batch_on_this_cp_rank(batch) # The implementation of this function is in MCore + else: # Packed THD format + assert ( + cu_seqlens.dim() == 2 and cu_seqlens.shape[0] == 1 + ), "micro-batch-size must be 1 for packing" + cu_seqlens = cu_seqlens[0] + batch['cu_seqlens'] = cu_seqlens + + max_seqlen = batch['max_seqlen'] + assert max_seqlen.dim() == 1 + # TODO(duncan): can this be kept as a 0-D tensor? + batch['max_seqlen'] = int(max_seqlen[0].item()) + + cp_size = get_context_parallel_world_size() + if cp_size > 1: # slice batch along sequence dimension for context parallelism + assert tex is not None and is_te_min_version("1.10.0"), ( + "Please update Transformer Engine to >= 1.10 to use " + "Context Parallel with THD format data" + ) + cp_rank = get_context_parallel_rank() + index = tex.thd_get_partitioned_indices( + cu_seqlens, + batch['tokens'].size(1), + cp_size, + cp_rank, + ) + for key, data in batch.items(): + if key in {'attention_mask', 'cu_seqlens', 'max_seqlen'}: + continue + batch[key] = data.index_select(1, index) + + return ( + batch.get('tokens'), + batch.get('labels'), + batch.get('loss_mask'), + batch.get('attention_mask'), + batch.get('position_ids'), + batch.get('cu_seqlens'), + batch.get('max_seqlen'), + ) + + +# define spiky loss as a loss that's 10x the max loss observed +SPIKY_LOSS_FACTOR = 10 + + +def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optional[HybridModel] = None, selected_budget=None): + """Loss function. + + Args: + loss_mask (torch.Tensor): Used to mask out some portions of the loss + output_tensor (torch.Tensor): The tensor with the losses + + Returns: + the loss scalar for this micro-batch + the number of non-padded tokens in this microbatch + a dict containing reporting metrics on the loss and number of tokens across + the data parallel ranks + """ + args = get_args() + if has_nvidia_modelopt: + return loss_func_modelopt(loss_mask, output_tensor, model=model, selected_budget=selected_budget) + + alpha = args.loss_alpha + + (output_tensor, (param_loss, extra_reporting_dict)) = output_tensor + + if param_loss is not None: + if param_loss > 0: + param_loss_report = param_loss.detach().clone() + else: + param_loss_report = param_loss.detach().clone() + param_loss = -args.router_beta * param_loss + + losses = output_tensor.view(-1).float() + loss_mask = loss_mask.view(-1).float() + loss = torch.sum(losses * loss_mask) + + # Check individual rank losses are not NaN prior to DP all-reduce. + rerun_state_machine = get_rerun_state_machine() + if args.check_for_nan_in_loss_and_grad: + rerun_state_machine.validate_result( + result=loss, + rejection_func=torch.isnan, + message="found NaN in local forward loss calculation", + tolerance=0.0, # forward pass calculations are deterministic + fatal=True, + ) + rerun_state_machine.validate_result( + result=loss, + rejection_func=torch.isinf, + message="found Inf in local forward loss calculation", + tolerance=0.0, # forward pass calculations are deterministic + fatal=True, + ) + # Check for spiky loss + if args.check_for_spiky_loss: + rerun_state_machine.validate_result( + result=loss, + rejection_func=partial( + rerun_state_machine.is_unexpectedly_large, + threshold=SPIKY_LOSS_FACTOR, + context="loss", + ), + message="Spiky loss", + tolerance=0.0, # forward pass calculations are deterministic + fatal=False, + ) + + + num_tokens = loss_mask.sum().clone().detach().to(torch.int) + + if param_loss is not None: + param_loss *= num_tokens * alpha + if param_loss < 0: + param_loss = -args.router_beta * param_loss + + param_loss_report = torch.cat([param_loss.clone().detach().view(1), num_tokens.view(1)]) + lm_loss_report = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) + loss += param_loss[0] + + # Protect against division by zero when all tokens are masked. + num_tokens = torch.clamp(num_tokens, min=1) + reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) + + if param_loss is not None: + return (loss, num_tokens, {'lm loss': lm_loss_report, + 'param loss': param_loss_report, + 'total loss': reporting_loss}) + else: + return (loss, num_tokens, {'lm loss': reporting_loss}) + +def get_grad_acc_based_random_choice(args, choices=None, prob=None, base_seed=42): + + dp_size = get_data_parallel_world_size() + grad_accumulation_steps = get_current_global_batch_size() // (get_micro_batch_size() * dp_size) + global _global_choice_counter + + # DP-specific seeding + rng = np.random.RandomState(base_seed + _global_choice_counter + grad_accumulation_steps*args.curr_iteration*10) + if choices is None: + choice = rng.uniform(0, 1) + else: + if prob is None: + choice = rng.choice(choices) + else: + choice = rng.choice(choices, p=prob) + _global_choice_counter += 1 + _global_choice_counter %= grad_accumulation_steps + return choice + +def forward_step(data_iterator, model: HybridModel): + """Forward training step. + + Args: + data_iterator : Input data iterator + model (HybridModel): The GPT Model + """ + args = get_args() + timers = get_timers() + + # One-time per-component params-norm breakdown (mirrors calc_params_l2_norm). + global _logged_params_norm + if not _logged_params_norm: + _logged_params_norm = True + from collections import defaultdict + groups = defaultdict(float) + trainable_sq = frozen_sq = total_sq = 0.0 + for name, param in model.named_parameters(): + # Use fp32 main_param when distributed optimizer is active (matches WandB metric). + # Note: main_param can be None for some params under DistOpt; getattr's default + # only applies when the attr is missing, so handle the None case explicitly. + main = getattr(param, 'main_param', None) + p = (main if main is not None else param.detach()).float() + norm_sq = p.norm(2).item() ** 2 + # Strip DDP 'module.' wrappers to get the logical top-level name. + clean = name + while clean.startswith('module.'): + clean = clean[len('module.'):] + top = clean.split('.')[0] + groups[top] += norm_sq + total_sq += norm_sq + if param.requires_grad: + trainable_sq += norm_sq + else: + frozen_sq += norm_sq + print_rank_0( + f"[PARAMS_NORM] total={total_sq**0.5:.2f} " + f"trainable={trainable_sq**0.5:.2f} frozen={frozen_sq**0.5:.2f}" + ) + for grp, sq in sorted(groups.items()): + print_rank_0(f"[PARAMS_NORM] {grp}: {sq**0.5:.2f}") + + # Get the batch. + timers('batch-generator', log_level=2).start() + global stimer + with stimer(bdata=True): + ( + tokens, + labels, + loss_mask, + attention_mask, + position_ids, + cu_seqlens, + max_seqlen, + ) = get_batch(data_iterator) + timers('batch-generator').stop() + + if get_grad_acc_based_random_choice(args=args) < args.original_model_sample_prob: + # Funnel "full-model sample" through the regular router-driven path + # with budget=1.0. flextron_forward unconditionally reads + # flextron_kwargs['budget'], so an empty dict would KeyError here. + flextron_kwargs = {'budget': 1.0} + selected_budget = 1.0 + else: + if args.budget_probs is None: + budget_probs = [1.0 for _ in args.budget_list] + else: + budget_probs = args.budget_probs + + assert len(args.budget_list) == len(budget_probs), "budget_list and budget_probs must have the same length" + budget_probs = [float(p) for p in budget_probs] + budget_probs = [p / sum(budget_probs) for p in budget_probs] + selected_budget = get_grad_acc_based_random_choice(args=args, choices=args.budget_list, prob=budget_probs) + flextron_kwargs = {'budget': selected_budget} + + with stimer: + output_tensor = model(tokens, position_ids, attention_mask, + labels=labels, **flextron_kwargs) + + # [ModelOpt]: model is needed to access ModelOpt distillation losses + return output_tensor, partial(loss_func, loss_mask, model=model, selected_budget=selected_budget) + + + +def is_dataset_built_on_rank(vp_stage=None): + ignore_virtual = True + if vp_stage is not None: + ignore_virtual = False + return ( + mpu.is_pipeline_first_stage(ignore_virtual=ignore_virtual, vp_stage=vp_stage) + or mpu.is_pipeline_last_stage(ignore_virtual=ignore_virtual, vp_stage=vp_stage) + ) and mpu.get_tensor_model_parallel_rank() == 0 + + +def core_gpt_dataset_config_from_args(args): + tokenizer = get_tokenizer() + + # Sometimes --data-path is too long, instead we parse it from a file. + blend: Optional[Tuple[List[str], Optional[List[float]]]] + blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]] + blend, blend_per_split = get_blend_and_blend_per_split(args) + + return GPTDatasetConfig( + random_seed=args.seed, + sequence_length=args.seq_length, + blend=blend, + blend_per_split=blend_per_split, + split=args.split, + num_dataset_builder_threads=args.num_dataset_builder_threads, + path_to_cache=args.data_cache_path, + mmap_bin_files=args.mmap_bin_files, + tokenizer=tokenizer, + reset_position_ids=args.reset_position_ids, + reset_attention_mask=args.reset_attention_mask, + eod_mask_loss=args.eod_mask_loss, + create_attention_mask=args.create_attention_mask_in_dataloader, + object_storage_cache_path=args.object_storage_cache_path, + mid_level_dataset_surplus=args.mid_level_dataset_surplus, + ) + + +def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None): + """Build the train test and validation datasets. + + Args: + train_val_test_num_samples : A list containing the number of samples in train test and validation. + """ + args = get_args() + + config = core_gpt_dataset_config_from_args(args) + + if args.sft: + dataset_type = SFTDataset + else: + if args.mock_data: + dataset_type = MockGPTDataset + else: + dataset_type = GPTDataset + + print_rank_0("> building train, validation, and test datasets for GPT ...") + + train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( + dataset_type, + train_val_test_num_samples, + partial(is_dataset_built_on_rank, vp_stage=vp_stage), + config + ).build() + + print_rank_0("> finished creating GPT datasets ...") + + return train_ds, valid_ds, test_ds + + +def mamba_flex_extra_args_provider(parser): + """Add Flextron CLI if not already registered by ``add_megatron_arguments``, then ModelOpt.""" + if not any(getattr(action, "dest", None) == "flextron" for action in parser._actions): + parser = add_flextron_args(parser) + if has_nvidia_modelopt: + parser = add_modelopt_args(parser) + return parser + + +if __name__ == "__main__": + + # Temporary for transition to core datasets + train_valid_test_datasets_provider.is_distributed = True + + # Optionally enable inprocess restart on pretrain + pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain) + + # Restore router LR multiplier (Bug 4 fix): monkey-patch get_megatron_optimizer_config + # to inject a per-parameter LR override for router params via config_overrides. + # Main branch removed the scale_lr_cond parameter from pretrain(); this achieves the same. + import megatron.training.training as _mtt + from megatron.core.optimizer.optimizer_config import ParamKey, ParamWithNamePredicate + from megatron.core.optimizer_param_scheduler import ParamGroupOverride + + _orig_get_opt_cfg = _mtt.get_megatron_optimizer_config + + def _patched_get_opt_cfg(args): + config, config_overrides = _orig_get_opt_cfg(args) + lr_mult = getattr(args, 'lr_mult_router', 1.0) + if lr_mult != 1.0: + router_key = ParamKey( + with_name_predicate=ParamWithNamePredicate( + name="router_pp", + fn=lambda p, name: 'router_pp' in name, + ) + ) + router_override = ParamGroupOverride( + max_lr=args.lr * lr_mult, + min_lr=args.min_lr * lr_mult, + ) + config_overrides = {**(config_overrides or {}), router_key: router_override} + return config, config_overrides + + _mtt.get_megatron_optimizer_config = _patched_get_opt_cfg + + # `pretrain()` no longer accepts extra_args_provider / args_defaults; parse + # args up-front instead (see pretrain_mamba.py for the same pattern). + args = parse_and_validate_args( + extra_args_provider=mamba_flex_extra_args_provider, + args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, + ) + + full_config = pretrain_cfg_container_from_args(args) + pretrain(full_config, + train_valid_test_datasets_provider, + model_provider, + ModelType.encoder_or_decoder, + forward_step, + store=store, + ) diff --git a/megatron/elastification/router/__init__.py b/megatron/elastification/router/__init__.py new file mode 100644 index 00000000000..26496bfed70 --- /dev/null +++ b/megatron/elastification/router/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. diff --git a/megatron/elastification/router/flex_budget_utils.py b/megatron/elastification/router/flex_budget_utils.py new file mode 100644 index 00000000000..ba6537196f1 --- /dev/null +++ b/megatron/elastification/router/flex_budget_utils.py @@ -0,0 +1,393 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from typing import Dict, List, Optional, Tuple, Union + +import torch + + +def get_num_parameters( + hybrid_pattern: str = None, + mamba_num_heads: int = 0, + mamba_d_head: int = 0, + mamba_d_state: int = 0, + num_attention_heads: int = 0, + num_query_groups: int = 0, + ffn_hidden_size: int = 0, + hidden_size: int = 0, + kv_channels: int = 0, + vocab_size: int = 0, + tied_vocab: bool = False, + num_experts: int = 0, + shared_expert_intermediate_size: int = 0, + moe_router_topk: int = 0, +) -> int: + + norm_multiplier = 1 + + embedding = vocab_size * hidden_size + final_layernorm = hidden_size * 1 + output_layer = 0 if tied_vocab else (vocab_size * hidden_size) + if isinstance(ffn_hidden_size, int): + flex_hetero_ffn = False + else: + flex_hetero_ffn = ffn_hidden_size.shape[0] != 1 + + if isinstance(mamba_num_heads, int): + flex_hetero_mamba = False + else: + flex_hetero_mamba = mamba_num_heads.shape[0] != 1 + + # Per-layer attention head counts arise only from layer skipping; head + # elasticity itself is no longer supported. + if isinstance(num_attention_heads, int): + per_layer_attn_heads = False + else: + per_layer_attn_heads = num_attention_heads.shape[0] != 1 + + if isinstance(num_experts, int): + flex_hetero_moe_expert = False + else: + flex_hetero_moe_expert = num_experts.shape[0] != 1 + + # MOE + + if flex_hetero_ffn or flex_hetero_moe_expert: + if flex_hetero_ffn and not flex_hetero_moe_expert: + num_experts = [num_experts] * ffn_hidden_size.shape[0] + if flex_hetero_moe_expert and not flex_hetero_ffn: + ffn_hidden_size = [ffn_hidden_size] * num_experts.shape[0] + + moe_all = [] + moe_active = [] + for i in range(len(num_experts)): + pre_moe_ln = norm_multiplier * hidden_size + linear_fc1 = ffn_hidden_size[i] * ( + hidden_size * num_experts[i] + shared_expert_intermediate_size + ) + linear_fc2 = ffn_hidden_size[i] * ( + hidden_size * num_experts[i] + shared_expert_intermediate_size + ) + linear_fc1_active = ffn_hidden_size[i] * ( + hidden_size * moe_router_topk + shared_expert_intermediate_size + ) + linear_fc2_active = ffn_hidden_size[i] * ( + hidden_size * moe_router_topk + shared_expert_intermediate_size + ) + moe_all.append(pre_moe_ln + linear_fc1 + linear_fc2) + moe_active.append(pre_moe_ln + linear_fc1_active + linear_fc2_active) + else: + pre_mlp_ln = norm_multiplier * hidden_size + linear_fc1 = ffn_hidden_size * (hidden_size * num_experts + shared_expert_intermediate_size) + linear_fc2 = ffn_hidden_size * (hidden_size * num_experts + shared_expert_intermediate_size) + linear_fc1_active = ffn_hidden_size * ( + hidden_size * moe_router_topk + shared_expert_intermediate_size + ) + linear_fc2_active = ffn_hidden_size * ( + hidden_size * moe_router_topk + shared_expert_intermediate_size + ) + moe_all = pre_mlp_ln + linear_fc1 + linear_fc2 + moe_active = pre_mlp_ln + linear_fc1_active + linear_fc2_active + + # ATT + if per_layer_attn_heads: + att = [] + for i in range(num_attention_heads.shape[0]): + input_ln = norm_multiplier * hidden_size + linear_proj = num_attention_heads[i] * kv_channels * hidden_size + linear_qkv = (num_attention_heads[i] + 2 * num_query_groups) * kv_channels * hidden_size + att.append(input_ln + linear_proj + linear_qkv) + else: + input_ln = norm_multiplier * hidden_size + linear_proj = num_attention_heads * kv_channels * hidden_size + linear_qkv = (num_attention_heads + 2 * num_query_groups) * kv_channels * hidden_size + att = input_ln + linear_proj + linear_qkv + + # Mamba + def mamba_params(mamba_nheads): + d_inner = mamba_nheads * mamba_d_head + ngroups = 8 + + def get_conv_params(kernel_size, stride): + cdim = d_inner + 2 * ngroups * mamba_d_state + cbias = cdim + cweight = cdim * stride * kernel_size + return cbias + cweight + + mamba_dt_bias = mamba_nheads + mamba_A_log = mamba_nheads + # self.d_inner_local if self.D_has_hdim else self.nheads_local, + mamba_D = mamba_nheads + mamba_input_ln = norm_multiplier * hidden_size + mamba_in_proj = hidden_size * (d_inner * 2 + 2 * ngroups * mamba_d_state + mamba_nheads) + mamba_conv = get_conv_params(4, 1) + mamba_norm = d_inner + mamba_out_proj = d_inner * hidden_size + return ( + mamba_dt_bias + + mamba_A_log + + mamba_D + + mamba_input_ln + + mamba_in_proj + + mamba_conv + + mamba_norm + + mamba_out_proj + ) + + all_params = 0 + active_params = 0 + for i, c in enumerate(hybrid_pattern): + + if c == 'M': + if flex_hetero_mamba: + mamba_idx = hybrid_pattern[: i + 1].count('M') - 1 + all_params += mamba_params(mamba_num_heads[mamba_idx]) + active_params += mamba_params(mamba_num_heads[mamba_idx]) + else: + all_params += mamba_params(mamba_num_heads) + active_params += mamba_params(mamba_num_heads) + elif c == '*': + if per_layer_attn_heads: + head_idx = hybrid_pattern[: i + 1].count('*') - 1 + all_params += att[head_idx] + active_params += att[head_idx] + else: + all_params += att + active_params += att + elif c == 'E': + if flex_hetero_ffn or flex_hetero_moe_expert: + # Count how many 'E' characters appear before and including layer i + moe_idx = hybrid_pattern[: i + 1].count('E') - 1 + all_params += moe_all[moe_idx] + active_params += moe_active[moe_idx] + else: + all_params += moe_all + active_params += moe_active + elif c == '|': + pass + else: + raise RuntimeError(f'Unknown layer type: {c}') + + return ( + embedding + all_params + final_layernorm + output_layer, + embedding + active_params + final_layernorm + output_layer, + ) + + +def get_kv_cache_size( + hybrid_pattern: str = None, + num_attention_heads=None, + num_query_groups=None, + kv_channels=None, + mem_infer_seq_len: int = 0, + mem_batch_size: int = 0, +) -> Union[int, torch.Tensor]: + + # Per-layer attention head counts arise only from layer skipping; head + # elasticity itself is no longer supported. + if isinstance(num_attention_heads, int): + per_layer_attn_heads = False + else: + per_layer_attn_heads = num_attention_heads.shape[0] != 1 + + if per_layer_attn_heads: + kv_cache_size = 0 + head_idx = 0 + + for c in hybrid_pattern: + if c == '*': + current_heads = num_attention_heads[head_idx] + + kv_cache_size_per_layer = ( + 2.0 + * mem_batch_size + * mem_infer_seq_len + * num_query_groups + * current_heads + * kv_channels + / current_heads.detach().item() + ) + kv_cache_size += kv_cache_size_per_layer + head_idx += 1 + + else: + num_attention_layers = hybrid_pattern.count('*') + divider = ( + num_attention_heads.detach().item() + if isinstance(num_attention_heads, torch.Tensor) + else num_attention_heads + ) + kv_cache_size = ( + 2.0 + * mem_batch_size + * mem_infer_seq_len + * num_query_groups + * num_attention_heads + * kv_channels + * num_attention_layers + / divider + ) + + return kv_cache_size + + +def get_mamba_ssm_cache_size( + hybrid_pattern: str = None, + mamba_num_heads: int = 0, + mamba_d_head: int = 0, + mamba_d_state: int = 0, + mem_batch_size: int = 0, +) -> int: + + if isinstance(mamba_num_heads, int): + flex_hetero_mamba = False + else: + flex_hetero_mamba = mamba_num_heads.shape[0] != 1 + + if flex_hetero_mamba: + ssm_cache_size = 0 + mamba_idx = 0 + for c in hybrid_pattern: + if c == 'M': + current_mamba_num_heads = mamba_num_heads[mamba_idx] + ssm_cache_size += ( + mem_batch_size * current_mamba_num_heads * mamba_d_head * mamba_d_state + ) + mamba_idx += 1 + + else: + num_mamba_layers = hybrid_pattern.count('M') + ssm_cache_size = ( + mem_batch_size * mamba_num_heads * mamba_d_head * mamba_d_state * num_mamba_layers + ) + + return ssm_cache_size + + +def get_max_buffer_size( + hybrid_pattern: str = None, + moe_num_experts: int = 0, + shared_expert_intermediate_size: int = 0, + ffn_hidden_size: int = 0, + moe_router_topk: int = 0, + mem_batch_size: int = 0, + prefill_chunk_size: int = 0, +) -> int: + + if isinstance(moe_num_experts, int) or moe_num_experts.shape[0] == 1: + moe_num_experts = ( + torch.tensor([moe_num_experts] * hybrid_pattern.count('E')) + .to(torch.cuda.current_device()) + .float() + ) + + if isinstance(ffn_hidden_size, int) or ffn_hidden_size.shape[0] == 1: + ffn_hidden_size = ( + torch.tensor([ffn_hidden_size] * hybrid_pattern.count('E')) + .to(torch.cuda.current_device()) + .float() + ) + + max_buffer_list = [] + moe_idx = 0 + for char in hybrid_pattern: + if char == 'E': + current_moe_num_experts = moe_num_experts[moe_idx] + current_ffn_hidden_size = ffn_hidden_size[moe_idx] + max_buffer_list.append( + shared_expert_intermediate_size + current_ffn_hidden_size * moe_router_topk + ) + moe_idx += 1 + + max_buffer = torch.stack(max_buffer_list) + max_buffer_softmax = torch.nn.functional.softmax(max_buffer, dim=0) + max_buffer = (max_buffer_softmax * max_buffer).sum().unsqueeze(0) + max_buffer *= mem_batch_size * prefill_chunk_size + + return max_buffer + + +def get_memory_footprint( + hybrid_pattern: str = None, + mamba_num_heads: int = 0, + mamba_d_head: int = 80, + mamba_d_state: int = 128, + num_attention_heads: int = 0, + num_query_groups: int = 8, + ffn_hidden_size: int = 0, + hidden_size: int = 0, + kv_channels: int = 128, + vocab_size: int = 131072, + tied_vocab: bool = False, + mem_infer_seq_len: int = 131072, + mem_batch_size: int = 1, + prefill_chunk_size: int = 16384, + moe_num_experts: int = 0, + shared_expert_intermediate_size: int = 0, + moe_router_topk: int = 0, + memory_config=None, +): + """ + Returns total inference memory footprint in GB. + + Parameters + ---------- + memory_config : MemoryConfig, optional + Bytes-per-element values and param budget target. When None, defaults + to BF16 for all components (bpe=2). Pass a MemoryConfig built via + ``load_memory_config(args)`` to select a quantisation profile. + """ + from megatron.elastification.memory_config import MemoryConfig + + if memory_config is None: + memory_config = MemoryConfig() # BF16 defaults + + # Select all-param or active-param count based on param_budget_target + param_idx = 1 if memory_config.param_budget_target == "active" else 0 + + mem_params = ( + memory_config.bpe_params + * get_num_parameters( + hybrid_pattern=hybrid_pattern, + mamba_num_heads=mamba_num_heads, + mamba_d_head=mamba_d_head, + mamba_d_state=mamba_d_state, + num_attention_heads=num_attention_heads, + num_query_groups=num_query_groups, + ffn_hidden_size=ffn_hidden_size, + hidden_size=hidden_size, + kv_channels=kv_channels, + vocab_size=vocab_size, + tied_vocab=tied_vocab, + num_experts=moe_num_experts, + shared_expert_intermediate_size=shared_expert_intermediate_size, + moe_router_topk=moe_router_topk, + )[param_idx] + ) + + mem_kv_cache = memory_config.bpe_kv_cache * get_kv_cache_size( + hybrid_pattern=hybrid_pattern, + num_attention_heads=num_attention_heads, + num_query_groups=num_query_groups, + kv_channels=kv_channels, + mem_infer_seq_len=mem_infer_seq_len, + mem_batch_size=mem_batch_size, + ) + + mem_max_buffer = memory_config.bpe_max_buffer * get_max_buffer_size( + hybrid_pattern=hybrid_pattern, + moe_num_experts=moe_num_experts, + shared_expert_intermediate_size=shared_expert_intermediate_size, + ffn_hidden_size=ffn_hidden_size, + moe_router_topk=moe_router_topk, + mem_batch_size=mem_batch_size, + prefill_chunk_size=prefill_chunk_size, + ) + + mem_mamba_ssm_cache = memory_config.bpe_ssm_cache * get_mamba_ssm_cache_size( + hybrid_pattern=hybrid_pattern, + mamba_num_heads=mamba_num_heads, + mamba_d_head=mamba_d_head, + mamba_d_state=mamba_d_state, + mem_batch_size=mem_batch_size, + ) + return (mem_params + mem_kv_cache + mem_max_buffer + mem_mamba_ssm_cache) / 1024 / 1024 / 1024 diff --git a/megatron/elastification/router/hybrid_flex_router.py b/megatron/elastification/router/hybrid_flex_router.py new file mode 100644 index 00000000000..915513af820 --- /dev/null +++ b/megatron/elastification/router/hybrid_flex_router.py @@ -0,0 +1,622 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import random + +import numpy as np +import torch +import torch.nn as nn +from torch.nn import functional as F + +from megatron.core import parallel_state +from megatron.core.num_microbatches_calculator import ( + get_current_global_batch_size, + get_micro_batch_size, +) +from megatron.core.parallel_state import ( + get_data_parallel_rank, + get_data_parallel_world_size, + get_pipeline_model_parallel_rank, + get_pipeline_model_parallel_world_size, + get_tensor_model_parallel_group, + get_tensor_model_parallel_rank, +) + +# Remove top-level import to avoid circular imports +# from megatron.training import get_args, print_rank_0 +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import init_method_normal + +# Import TE parallel linear layers +try: + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TERowParallelLinear, + ) + + HAVE_TE = True +except ImportError: + HAVE_TE = False + # Fallback to regular tensor parallel layers + from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear + + +# Router implementation for pre-gating router. +# Use router to determine #heads, MLP sizes, and layers to skip (Router_v2) +# Only takes the budget as input +class FlextronRouter(MegatronModule): + def __init__(self, config: TransformerConfig): + super().__init__(config=config) + + self.config = config + self.input_dim = len(self.config.budget_list) + self.n_dim = self.config.router_inter_dim + self.budget_map = { + item: torch.tensor(idx) for idx, item in enumerate(self.config.budget_list) + } + + # Initialize DP-aware Gumbel softmax + self._init_dp_gumbel_softmax() + + # Create init method for router layers + self.init_method = init_method_normal(self.config.router_std) + + self.add_router_for_mlp() + self.add_router_for_emb() + self.add_router_for_mamba() + self.add_router_for_moe_expert() + if self.config.add_skipping: + self.add_router_for_skipping() + + # Synchronize router weights across all pipeline parallel ranks + self._sync_router_weights() + self._mark_router_params_for_pp_sync() + self.hard_sample_th = config.hard_sample_th + + self.add_scaler_schedule() + + self.dp_size = get_data_parallel_world_size() + self.grad_accumulation_steps = get_current_global_batch_size() // ( + get_micro_batch_size() * self.dp_size + ) + self.fwd_pass_count = 0 + + def _init_dp_gumbel_softmax(self, base_seed=42): + """Initialize DP-aware Gumbel softmax functionality""" + self.dp_rank = get_data_parallel_rank() + self.gumbel_base_seed = base_seed + + def _sync_router_weights(self): + """ + Synchronize router weights across all pipeline parallel groups by broadcasting + from global rank 0 to all other ranks. + """ + if not torch.distributed.is_initialized(): + return + + # Get global rank 0 as the source + source_rank = 0 + + # Broadcast all router parameters from rank 0 + for name, param in self.named_parameters(): + if param is not None: + torch.distributed.broadcast(param.data, src=source_rank) + + def _mark_router_params_for_pp_sync(self): + """ + Mark all router parameters to be synchronized across pipeline parallel ranks. + This ensures they get handled by the main gradient synchronization system. + """ + for param in self.parameters(): + if param.requires_grad: + # Mark parameter for pipeline parallel synchronization + setattr(param, 'flextron_router_pp_sync', True) + + def _dp_gumbel_softmax(self, logits, tau=1.0, hard=False, curr_iteration=0): + """DP-aware Gumbel softmax that uses different random seeds per DP rank and iteration""" + # Create unique seed for this iteration and DP rank + + seed = ( + self.gumbel_base_seed + + (self.dp_rank + self.fwd_pass_count * self.dp_size) % self.config.router_gbs + + curr_iteration * 1000 + ) + # torch.manual_seed seeds both CPU and CUDA generators globally, so we + # must save/restore both - otherwise the CUDA RNG leaks the deterministic + # state we set here into other CUDA random ops elsewhere in the model. + cpu_state = torch.get_rng_state() + cuda_state = ( + torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None + ) + torch.manual_seed(seed) + + try: + return F.gumbel_softmax(logits, tau=tau, hard=hard) + finally: + torch.set_rng_state(cpu_state) + if cuda_state is not None: + torch.cuda.set_rng_state_all(cuda_state) + + def _create_linear_layer(self, input_size, output_size, bias=False, is_first_layer=True): + """Helper method to create appropriate linear layer (TE or fallback)""" + if HAVE_TE: + if is_first_layer: + # First layer: TEColumnParallelLinear + return TEColumnParallelLinear( + input_size=input_size, + output_size=output_size, + config=self.config, + init_method=self.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + ) + else: + # Second layer: TERowParallelLinear + return TERowParallelLinear( + input_size=input_size, + output_size=output_size, + config=self.config, + init_method=self.init_method, + bias=bias, + input_is_parallel=True, + skip_bias_add=False, + is_expert=False, + ) + else: + # Fallback to regular tensor parallel layers + if is_first_layer: + return ColumnParallelLinear( + input_size=input_size, + output_size=output_size, + config=self.config, + init_method=self.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + ) + else: + return RowParallelLinear( + input_size=input_size, + output_size=output_size, + config=self.config, + init_method=self.init_method, + bias=bias, + input_is_parallel=True, + skip_bias_add=False, + is_expert=False, + ) + + def add_router_for_mlp(self): + mlp_list = self.config.mlp_int_list + if self.config.flex_hetero_ffn: + num_mlp = self.config.hybrid_layer_pattern.count("E") + gate_mlp_layer_list = [ + self._create_linear_layer( + self.input_dim, self.n_dim, bias=False, is_first_layer=True + ), + nn.LeakyReLU(0.1), + self._create_linear_layer( + self.n_dim, len(mlp_list) * num_mlp, bias=False, is_first_layer=False + ), + ] + # Set bias for the last layer + if ( + hasattr(gate_mlp_layer_list[-1], 'bias') + and gate_mlp_layer_list[-1].bias is not None + ): + last_layer_bias = [0.00 for _ in range(len(mlp_list))] + last_layer_bias[-1] = 1.00 + gate_mlp_layer_list[-1].bias.data = torch.tensor( + last_layer_bias, + dtype=gate_mlp_layer_list[-1].weight.dtype, + device=gate_mlp_layer_list[-1].weight.device, + ).repeat(num_mlp) + else: + gate_mlp_layer_list = [ + self._create_linear_layer( + self.input_dim, self.n_dim, bias=False, is_first_layer=True + ), + nn.LeakyReLU(0.1), + self._create_linear_layer( + self.n_dim, len(mlp_list), bias=False, is_first_layer=False + ), + ] + self.gate_mlp = nn.Sequential(*gate_mlp_layer_list) + + def add_router_for_moe_expert(self): + moe_expert_list = self.config.moe_expert_int_list + if self.config.flex_hetero_moe_expert: + num_moe_expert = self.config.hybrid_layer_pattern.count("E") + gate_moe_expert_layer_list = [ + self._create_linear_layer( + self.input_dim, self.n_dim, bias=False, is_first_layer=True + ), + nn.LeakyReLU(0.1), + self._create_linear_layer( + self.n_dim, + len(moe_expert_list) * num_moe_expert, + bias=False, + is_first_layer=False, + ), + ] + else: + gate_moe_expert_layer_list = [ + self._create_linear_layer( + self.input_dim, self.n_dim, bias=False, is_first_layer=True + ), + nn.LeakyReLU(0.1), + self._create_linear_layer( + self.n_dim, len(moe_expert_list), bias=False, is_first_layer=False + ), + ] + self.gate_moe_expert = nn.Sequential(*gate_moe_expert_layer_list) + + def add_router_for_emb(self): + emb_list = self.config.emb_int_list + gate_emb_layer_list = [ + self._create_linear_layer(self.input_dim, self.n_dim, bias=False, is_first_layer=True), + nn.LeakyReLU(0.1), + self._create_linear_layer(self.n_dim, len(emb_list), bias=False, is_first_layer=False), + ] + self.gate_emb = nn.Sequential(*gate_emb_layer_list) + + def add_router_for_skipping(self): + + self.output_dim = int(len(self.config.layer_ranking_list) + 1) + + gate_skip_mlp_layer_list = [ + self._create_linear_layer(self.input_dim, self.n_dim, bias=False, is_first_layer=True), + nn.LeakyReLU(0.1), + self._create_linear_layer( + self.n_dim, self.output_dim, bias=False, is_first_layer=False + ), + ] + + self.gate_skip_layer = nn.Sequential(*gate_skip_mlp_layer_list) + + def add_router_for_mamba(self): + mamba_list = self.config.mamba_int_list + if self.config.flex_hetero_mamba: + num_mamba = self.config.hybrid_layer_pattern.count("M") + gate_mamba_layer_list = [ + self._create_linear_layer( + self.input_dim, self.n_dim, bias=False, is_first_layer=True + ), + nn.LeakyReLU(0.1), + self._create_linear_layer( + self.n_dim, len(mamba_list) * num_mamba, bias=False, is_first_layer=False + ), + ] + # Set bias for the last layer + if ( + hasattr(gate_mamba_layer_list[-1], 'bias') + and gate_mamba_layer_list[-1].bias is not None + ): + last_layer_bias = [0.00 for _ in range(len(mamba_list))] + last_layer_bias[-1] = 1.00 + gate_mamba_layer_list[-1].bias.data = torch.tensor(last_layer_bias).repeat( + num_mamba + ) + else: + gate_mamba_layer_list = [ + self._create_linear_layer( + self.input_dim, self.n_dim, bias=False, is_first_layer=True + ), + nn.LeakyReLU(0.1), + self._create_linear_layer( + self.n_dim, len(mamba_list), bias=False, is_first_layer=False + ), + ] + self.gate_mamba = nn.Sequential(*gate_mamba_layer_list) + + def mamba_forward(self, args, budget_tensor, device, dtype, tau, hard_sample): + + # TODO @ataghibakhsh: check router out of sync on TP ranks + + router_mamba_logits1 = self.gate_mamba[0](budget_tensor) + router_mamba_logits2 = self.gate_mamba[1](router_mamba_logits1[0]) + router_mamba_logits = self.gate_mamba[2](router_mamba_logits2)[0].flatten() + # torch.distributed.all_reduce(router_mamba_logits, group=get_tensor_model_parallel_group(), op=torch.distributed.ReduceOp.AVG) + if self.scaler is not None: + scale = self.scaler[args.curr_iteration].to(device=device, dtype=dtype) + + if self.config.flex_hetero_mamba: + mamba_n = len(self.config.mamba_int_list) + router_mamba_logits = router_mamba_logits.reshape(-1, mamba_n) + if self.config.normalize_router_logits: + router_mamba_logits = ( + scale + * router_mamba_logits + / router_mamba_logits.std(dim=1, keepdim=True).clamp(min=1e-6) + ) + else: + router_mamba_logits = scale * router_mamba_logits + router_mamba_logits = self._dp_gumbel_softmax( + router_mamba_logits, tau=tau, hard=hard_sample, curr_iteration=args.curr_iteration + ) + _, choices_mamba = torch.topk(router_mamba_logits, 1, dim=-1) + return ( + router_mamba_logits, + [self.config.mamba_int_list[i] for i in choices_mamba.flatten().tolist()], + ) + else: + if self.config.normalize_router_logits: + # Std-normalize only when there's actually >1 choice; with a + # single choice the std is 0 and the routing is trivial, so we + # skip both the scale and the normalization (consistent with + # the no-op semantics of a single-choice axis). + if len(self.config.mamba_int_list) > 1: + router_mamba_logits = ( + scale + * router_mamba_logits + / router_mamba_logits.std(dim=0, keepdim=True).clamp(min=1e-6) + ) + else: + router_mamba_logits = scale * router_mamba_logits + router_mamba_logits = self._dp_gumbel_softmax( + router_mamba_logits, tau=tau, hard=hard_sample, curr_iteration=args.curr_iteration + ) + _, choices_mamba = torch.topk(router_mamba_logits, 1, dim=-1) + return (router_mamba_logits, self.config.mamba_int_list[choices_mamba.item()]) + + def mlp_forward(self, args, budget_tensor, device, dtype, tau, hard_sample): + + # TODO @ataghibakhsh: check router out of sync on TP ranks + router_mlp_logits1 = self.gate_mlp[0](budget_tensor) + router_mlp_logits2 = self.gate_mlp[1](router_mlp_logits1[0]) + router_mlp_logits = self.gate_mlp[2](router_mlp_logits2)[0].flatten() + # torch.distributed.all_reduce(router_mlp_logits, group=get_tensor_model_parallel_group(), op=torch.distributed.ReduceOp.AVG) + if self.scaler is not None: + scale = self.scaler[args.curr_iteration].to(device=device, dtype=dtype) + if self.config.flex_hetero_ffn: + mlp_n = len(self.config.mlp_int_list) + router_mlp_logits = router_mlp_logits.reshape(-1, mlp_n) + if self.config.normalize_router_logits: + router_mlp_logits = ( + scale + * router_mlp_logits + / router_mlp_logits.std(dim=1, keepdim=True).clamp(min=1e-6) + ) + else: + router_mlp_logits = scale * router_mlp_logits + router_mlp_logits = self._dp_gumbel_softmax( + router_mlp_logits, tau=tau, hard=hard_sample, curr_iteration=args.curr_iteration + ) + _, choices_mlp = torch.topk(router_mlp_logits, 1, dim=-1) + return ( + router_mlp_logits, + [self.config.mlp_int_list[i] for i in choices_mlp.flatten().tolist()], + ) + else: + if self.config.normalize_router_logits: + # Std-normalize only when there's actually >1 choice; with a + # single choice the std is 0 and the routing is trivial, so we + # skip both the scale and the normalization (consistent with + # the no-op semantics of a single-choice axis). + if len(self.config.mlp_int_list) > 1: + router_mlp_logits = ( + scale + * router_mlp_logits + / router_mlp_logits.std(dim=0, keepdim=True).clamp(min=1e-6) + ) + else: + router_mlp_logits = scale * router_mlp_logits + router_mlp_logits = self._dp_gumbel_softmax( + router_mlp_logits, tau=tau, hard=hard_sample, curr_iteration=args.curr_iteration + ) + _, choices_mlp = torch.topk(router_mlp_logits, 1, dim=-1) + return (router_mlp_logits, self.config.mlp_int_list[choices_mlp.item()]) + + def moe_expert_forward(self, args, budget_tensor, device, dtype, tau, hard_sample): + router_moe_expert_logits1 = self.gate_moe_expert[0](budget_tensor) + router_moe_expert_logits2 = self.gate_moe_expert[1](router_moe_expert_logits1[0]) + router_moe_expert_logits = self.gate_moe_expert[2](router_moe_expert_logits2)[0].flatten() + # torch.distributed.all_reduce(router_moe_expert_logits, group=get_tensor_model_parallel_group(), op=torch.distributed.ReduceOp.AVG) + if self.scaler is not None: + scale = self.scaler[args.curr_iteration].to(device=device, dtype=dtype) + if self.config.flex_hetero_moe_expert: + moe_expert_n = len(self.config.moe_expert_int_list) + router_moe_expert_logits = router_moe_expert_logits.reshape(-1, moe_expert_n) + if self.config.normalize_router_logits: + router_moe_expert_logits = ( + scale + * router_moe_expert_logits + / router_moe_expert_logits.std(dim=1, keepdim=True).clamp(min=1e-6) + ) + else: + router_moe_expert_logits = scale * router_moe_expert_logits + router_moe_expert_logits = self._dp_gumbel_softmax( + router_moe_expert_logits, + tau=tau, + hard=hard_sample, + curr_iteration=args.curr_iteration, + ) + _, choices_moe_expert = torch.topk(router_moe_expert_logits, 1, dim=-1) + return ( + router_moe_expert_logits, + [self.config.moe_expert_int_list[i] for i in choices_moe_expert.flatten().tolist()], + ) + else: + if self.config.normalize_router_logits: + # Std-normalize only when there's actually >1 choice; with a + # single choice the std is 0 and the routing is trivial, so we + # skip both the scale and the normalization (consistent with + # the no-op semantics of a single-choice axis). + if len(self.config.moe_expert_int_list) > 1: + router_moe_expert_logits = ( + scale + * router_moe_expert_logits + / router_moe_expert_logits.std(dim=0, keepdim=True).clamp(min=1e-6) + ) + else: + router_moe_expert_logits = scale * router_moe_expert_logits + router_moe_expert_logits = self._dp_gumbel_softmax( + router_moe_expert_logits, + tau=tau, + hard=hard_sample, + curr_iteration=args.curr_iteration, + ) + _, choices_moe_expert = torch.topk(router_moe_expert_logits, 1, dim=-1) + return ( + router_moe_expert_logits, + self.config.moe_expert_int_list[choices_moe_expert.item()], + ) + + def emb_forward(self, args, budget_tensor, device, dtype, tau, hard_sample): + + router_emb_logits1 = self.gate_emb[0](budget_tensor) + router_emb_logits2 = self.gate_emb[1](router_emb_logits1[0]) + router_emb_logits = self.gate_emb[2](router_emb_logits2)[0].flatten() + # torch.distributed.all_reduce(router_emb_logits, group=get_tensor_model_parallel_group(), op=torch.distributed.ReduceOp.AVG) + if self.scaler is not None: + scale = self.scaler[args.curr_iteration].to(device=device, dtype=dtype) + router_emb_logits = scale * router_emb_logits + + # router_emb_logits = F.gumbel_softmax(router_emb_logits, tau=tau, hard=hard_sample) + router_emb_logits = self._dp_gumbel_softmax( + router_emb_logits, tau=tau, hard=hard_sample, curr_iteration=args.curr_iteration + ) + _, choices_emb = torch.topk(router_emb_logits, 1, dim=-1) + + return (router_emb_logits, self.config.emb_int_list[choices_emb.item()]) + + def skipping_forward(self, args, budget_tensor, device, dtype, tau, hard_sample): + + # for layer skipping, skipping MLP layers + router_skip_layer_logits1 = self.gate_skip_layer[0](budget_tensor) + router_skip_layer_logits2 = self.gate_skip_layer[1](router_skip_layer_logits1[0]) + router_skip_layer_logits = self.gate_skip_layer[2](router_skip_layer_logits2)[0].flatten() + # torch.distributed.all_reduce(router_skip_layer_logits, group=get_tensor_model_parallel_group(), op=torch.distributed.ReduceOp.AVG) + router_skip_layer_logits = torch.repeat_interleave( + router_skip_layer_logits, repeats=1, dim=0 + ) + if self.scaler is not None: + router_skip_layer_logits = router_skip_layer_logits * self.scaler[ + args.curr_iteration + ].to(device=device, dtype=dtype) + + # router_skip_layer_logits = F.gumbel_softmax(router_skip_layer_logits, tau=tau, hard=hard_sample) + router_skip_layer_logits = self._dp_gumbel_softmax( + router_skip_layer_logits, tau=tau, hard=hard_sample, curr_iteration=args.curr_iteration + ) + _, choices_skip_layer = torch.topk(router_skip_layer_logits, 1, dim=-1) + if choices_skip_layer.item() != 0: + selected_to_drop = self.config.layer_ranking_list[: choices_skip_layer.item()] + choices_skip_layer = torch.zeros(self.config.num_layers).to(device=device, dtype=dtype) + choices_skip_layer[selected_to_drop] = 1 + else: + choices_skip_layer = torch.zeros(self.config.num_layers).to(device=device, dtype=dtype) + return (router_skip_layer_logits, choices_skip_layer) + + def get_curr_tau(self, curr_iteration): + tau = self.config.tau_init * torch.pow(torch.tensor(self.config.tau_decay), curr_iteration) + return tau + + def add_scaler_schedule(self): + + if ( + self.config.linear_scaler_start is not None + and self.config.linear_scaler_end is not None + ): + from megatron.training import get_args + + args = get_args() + self.scaler = torch.linspace( + start=self.config.linear_scaler_start, + end=self.config.linear_scaler_end, + steps=( + args.train_iters + if args.train_iters is not None + else (args.train_samples // args.global_batch_size) + ), + ) + else: + self.scaler = None + + def forward(self, budget): + + from megatron.training import get_args + + args = get_args() + + hard_sample = random.random() > self.hard_sample_th + + tau = self.get_curr_tau(args.curr_iteration) + + device, dtype = next(self.parameters()).device, next(self.parameters()).dtype + + if budget in self.budget_map.keys(): + budget_tensor = torch.nn.functional.one_hot( + self.budget_map[budget], len(self.config.budget_list) + ).to(device=device, dtype=dtype) + elif budget == 1.0: + # Requested full model but 1.0 isn't a trained budget — fall back + # to the largest configured budget. Using max() instead of [0] + # makes this independent of budget_list ordering. + budget_tensor = torch.nn.functional.one_hot( + self.budget_map[max(self.budget_map.keys())], + len(self.config.budget_list), + ).to(device=device, dtype=dtype) + else: + # budget_list is enforced descending by sort_budget_list_descending + # at config-injection time. We re-sort ascending locally for + # bucketize, then flip(0) below to land back in the descending + # one-hot coordinate system the router was trained against. + budget_values = torch.tensor( + sorted(self.config.budget_list), device=device, dtype=dtype + ) + budget_t = torch.as_tensor(budget, device=device, dtype=dtype) + + # idx2 = first index where budget_values[idx] > budget (right=False gives >= behavior with floats) + idx2 = torch.bucketize(budget_t, budget_values, right=False) + # Clamp to valid interior so we always have a left neighbor + idx2 = idx2.clamp(min=1, max=len(self.config.budget_list) - 1) + idx1 = idx2 - 1 + + b1 = budget_values.index_select(0, idx1.to(torch.long)) + b2 = budget_values.index_select(0, idx2.to(torch.long)) + denom = b2 - b1 # .clamp_min(1e-12) + weight = (budget_t - b1) / denom # in [0,1] when budget is between b1 and b2 + + num_classes = len(self.config.budget_list) + one_hot_1 = torch.nn.functional.one_hot( + idx1.to(torch.long), num_classes=num_classes + ).to(device=device, dtype=dtype) + one_hot_2 = torch.nn.functional.one_hot( + idx2.to(torch.long), num_classes=num_classes + ).to(device=device, dtype=dtype) + + # If weight is scalar, broadcasting works; if vector, it blends per-sample + budget_tensor = (1 - weight).unsqueeze(-1) * one_hot_1 + weight.unsqueeze( + -1 + ) * one_hot_2 + budget_tensor = budget_tensor.squeeze(0).flip(0) + + budget_tensor = budget_tensor.unsqueeze(0) + mlp_forward_outputs = self.mlp_forward(args, budget_tensor, device, dtype, tau, hard_sample) + mamba_forward_outputs = self.mamba_forward( + args, budget_tensor, device, dtype, tau, hard_sample + ) + moe_expert_forward_outputs = self.moe_expert_forward( + args, budget_tensor, device, dtype, tau, hard_sample + ) + + if self.config.add_skipping: + skipping_forward_outputs = self.skipping_forward( + args, budget_tensor, device, dtype, tau, hard_sample + ) + else: + skipping_forward_outputs = None + + emb_forward_outputs = self.emb_forward(args, budget_tensor, device, dtype, tau, hard_sample) + self.fwd_pass_count += 1 + return ( + mlp_forward_outputs, + skipping_forward_outputs, + emb_forward_outputs, + mamba_forward_outputs, + moe_expert_forward_outputs, + ) diff --git a/megatron/post_training/model_builder.py b/megatron/post_training/model_builder.py index 383ae6ec8aa..2994ca5fed5 100644 --- a/megatron/post_training/model_builder.py +++ b/megatron/post_training/model_builder.py @@ -104,6 +104,13 @@ def _load_teacher_model_config(checkpoint_path: str) -> Namespace: args_dict = vars(get_args()).copy() del args_dict["kv_channels"] # not recalculated if present + # Setting teacher Flextron fields to false if training with Flextron, can be overridden + if "flextron" in args_dict: + config["flextron"] = False + if "enable_router" in args_dict: + config["enable_router"] = False + if "freeze_model" in args_dict: + config["freeze_model"] = False args_dict.update(config) # Backward compat: old checkpoints have hybrid_override_pattern but not hybrid_layer_pattern diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/golden_values_dev_dgx_h100.json new file mode 100644 index 00000000000..3550f734923 --- /dev/null +++ b/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/golden_values_dev_dgx_h100.json @@ -0,0 +1,137 @@ +{ + "lm loss": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 10.96772, + "2": 10.94878, + "3": 10.13223, + "4": 10.10273, + "5": 9.77116, + "6": 9.74409, + "7": 9.67572, + "8": 9.54443, + "9": 9.47498, + "10": 9.46493, + "11": 9.0426, + "12": 9.29951, + "13": 9.30226, + "14": 9.27311, + "15": 9.10353, + "16": 8.79531, + "17": 8.51664, + "18": 8.77952, + "19": 8.71592, + "20": 8.71692 + } + }, + "num-zeros": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 47889940.0, + "2": 49012016.0, + "3": 48644404.0, + "4": 48106204.0, + "5": 58246720.0, + "6": 63126404.0, + "7": 77135048.0, + "8": 67132312.0, + "9": 67750168.0, + "10": 64553460.0, + "11": 77738728.0, + "12": 64587012.0, + "13": 64661168.0, + "14": 61613476.0, + "15": 61997576.0, + "16": 71691368.0, + "17": 69936240.0, + "18": 70742744.0, + "19": 67630016.0, + "20": 67347280.0 + } + }, + "mem-allocated-bytes": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 1872031232.0, + "2": 1872031232.0, + "3": 1872031232.0, + "4": 1872031232.0, + "5": 1872031232.0, + "6": 1872031232.0, + "7": 1872031232.0, + "8": 1872031232.0, + "9": 1872031232.0, + "10": 1872031232.0, + "11": 1872031232.0, + "12": 1872031232.0, + "13": 1872031232.0, + "14": 1872031232.0, + "15": 1872031232.0, + "16": 1872031232.0, + "17": 1872031232.0, + "18": 1872031232.0, + "19": 1872031232.0, + "20": 1872031232.0 + } + }, + "mem-max-allocated-bytes": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 2476529664.0, + "2": 2827837440.0, + "3": 2827837440.0, + "4": 2847291904.0, + "5": 2886460416.0, + "6": 2886460416.0, + "7": 2938775552.0, + "8": 2938775552.0, + "9": 2938775552.0, + "10": 2938775552.0, + "11": 2938775552.0, + "12": 2938775552.0, + "13": 2938775552.0, + "14": 2938775552.0, + "15": 2938775552.0, + "16": 2938775552.0, + "17": 2938775552.0, + "18": 2938775552.0, + "19": 2938775552.0, + "20": 2938775552.0 + } + }, + "iteration-time": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": "nan", + "2": 43.50173, + "3": 1.91918, + "4": 0.33128, + "5": 0.32613, + "6": 0.30839, + "7": 0.30715, + "8": 0.29241, + "9": 0.28352, + "10": 0.2695, + "11": 0.27092, + "12": 0.27945, + "13": 0.26134, + "14": 0.28396, + "15": 0.2654, + "16": 0.2759, + "17": 0.26684, + "18": 0.25266, + "19": 0.25121, + "20": 0.25109 + } + } +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/model_config.yaml new file mode 100644 index 00000000000..493aaa31b16 --- /dev/null +++ b/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/model_config.yaml @@ -0,0 +1,127 @@ +ENV_VARS: + CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_NONDETERMINISTIC_ALGO: 1 + NCCL_ALGO: Ring + CUBLAS_WORKSPACE_CONFIG: :4096:8 + TORCH_COMPILE_DISABLE: 1 + TORCHDYNAMO_DISABLE: 1 + TORCH_INDUCTOR_DISABLE: 1 +MODEL_ARGS: + # ── Hybrid Mamba / Attention / MoE backbone (scaled down from train_flextron.sh) ── + --num-layers: 12 + --hidden-size: 1024 + --ffn-hidden-size: 512 + --num-attention-heads: 32 + --group-query-attention: true + --num-query-groups: 2 + --kv-channels: 128 + --mamba-num-heads: 32 + --mamba-head-dim: 64 + --hybrid-override-pattern: MEM*EMEM*EME* + --is-hybrid-model: true + --position-embedding-type: none + --normalization: RMSNorm + --squared-relu: true + --use-fused-weighted-squared-relu: true + --untie-embeddings-and-output-weights: true + --disable-bias-linear: true + --init-method-std: 0.0173 + --attention-dropout: 0.0 + --hidden-dropout: 0.0 + # ── MoE ─────────────────────────────────────────────────────────────────── + --num-experts: 32 + --moe-router-topk: 2 + --moe-router-score-function: sigmoid + --moe-router-enable-expert-bias: true + --moe-router-topk-scaling-factor: 2.5 + --moe-router-dtype: fp32 + --moe-router-load-balancing-type: none + --moe-aux-loss-coeff: 1.0e-4 + --moe-shared-expert-intermediate-size: 512 + --moe-shared-expert-overlap: true + --moe-token-dispatcher-type: alltoall + --moe-grouped-gemm: true + --moe-permute-fusion: true + --cross-entropy-loss-fusion: true + --cross-entropy-fusion-impl: native + # ── Parallelism (8 GPUs total: TP=2, EP=2, PP=1, CP=1 → DP=4) ───────────── + --tensor-model-parallel-size: 2 + --pipeline-model-parallel-size: 1 + --expert-model-parallel-size: 2 + --expert-tensor-parallel-size: 1 + --context-parallel-size: 1 + --sequence-parallel: true + --attention-backend: flash + # ── Data / tokenizer (common_pile + GPT2 BPE, matches hybrid tests) ─────── + --seq-length: 2048 + --max-position-embeddings: 2048 + --micro-batch-size: 1 + --global-batch-size: 8 + --train-iters: 20 + --data-path: ${DATA_PATH}/text/common_pile/v01_filtered_data/my-gpt3_00_text_document + --vocab-file: ${DATA_PATH}/text/common_pile/v01_filtered_data/bpe/vocab.json + --merge-file: ${DATA_PATH}/text/common_pile/v01_filtered_data/bpe/merges.txt + --split: 949,50,1 + --data-cache-path: ${DATA_CACHE_PATH} + --no-mmap-bin-files: true + --no-create-attention-mask-in-dataloader: true + # ── Optimizer / schedule ────────────────────────────────────────────────── + --lr: 1.0e-4 + --min-lr: 1.0e-5 + --lr-decay-style: cosine + --lr-warmup-fraction: 0.01 + --weight-decay: 0.0 + --clip-grad: 1.0 + --adam-beta1: 0.9 + --adam-beta2: 0.98 + --use-distributed-optimizer: true + --bf16: true + # ── Logging / checkpointing ─────────────────────────────────────────────── + --log-interval: 1 + --log-params-norm: true + --log-num-zeros-in-grad: true + --log-validation-ppl-to-tensorboard: true + --log-timers-to-tensorboard: true + --log-memory-to-tensorboard: true + --tensorboard-dir: ${TENSORBOARD_PATH} + --save-interval: 10000 + --save: ${CHECKPOINT_SAVE_PATH} + --ckpt-format: torch_dist + --ckpt-fully-parallel-save: true + --ckpt-fully-parallel-load: true + --dist-ckpt-strictness: ignore_all + --eval-interval: 1000 + --eval-iters: 4 + # ── Core model plumbing ─────────────────────────────────────────────────── + --use-mcore-models: true + --transformer-impl: transformer_engine + --export-default-te-spec: true + --export-model-type: HybridModel + --distributed-backend: nccl + --distributed-timeout-minutes: 20 + # ── Flextron (from train_flextron.sh flex_options, int-lists with ÷8) ───── + --flextron: true + --enable-router: true + --binary-mask: true + --soft-mask: true + --hard-sample-th: 0.996 + --router-beta: 1.0 + --original-model-sample-prob: 0.0 + --tau-init: 1.0 + --tau-decay: 0.9997 + --loss-alpha: 1.0 + --lr-mult-router: 100 + --router-gbs: 2 + --router-inter-dim: 256 + --budget-list: "[1.0 0.697]" + --budget-probs: "[1.0 1.0]" + --budget-type: param + --emb-int-list: "[1024 768 512]" + --mlp-int-list: "[512 384 256]" + --mamba-int-list: "[32 24 16]" + --moe-expert-int-list: "[32 24 16]" + --linear-scaler-start: 1.0 + --linear-scaler-end: 10.0 + --slice: true + --router-std: 0.1 +TEST_TYPE: regular diff --git a/tests/test_utils/recipes/h100/flextron.yaml b/tests/test_utils/recipes/h100/flextron.yaml new file mode 100644 index 00000000000..9bd40eaa493 --- /dev/null +++ b/tests/test_utils/recipes/h100/flextron.yaml @@ -0,0 +1,62 @@ +type: basic +format_version: 1 +maintainers: [mcore] +loggers: [stdout] +spec: + name: "{test_case}_{environment}_{platforms}" + model: hybrid + build: mcore-pyt-{environment} + nodes: 1 + gpus: 8 + n_repeat: 1 + platforms: dgx_h100 + script_setup: | + unset https_proxy + echo "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" | tee -a /root/.netrc + + # Checkout latest + cd /opt + rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm + git init + git remote add origin $MCORE_REPO + git fetch origin '+refs/merge-requests/*:refs/remotes/merge-requests/*' + git fetch origin $MCORE_MR_COMMIT + git checkout $MCORE_MR_COMMIT + git rev-parse HEAD + + # Checkout backwards-ref + cd /opt + rm -rf /opt/megatron-lm-legacy; mkdir megatron-lm-legacy; cd megatron-lm-legacy + git init + git remote add origin $MCORE_REPO + git fetch origin $MCORE_BACKWARDS_COMMIT + git checkout $MCORE_BACKWARDS_COMMIT + git rev-parse HEAD + rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ + script: |- + ls + cd /opt/megatron-lm + + ARGUMENTS=( + "DATA_PATH=/mnt/artifacts" + "DATA_CACHE_PATH=/workspace/data/cache" + "OUTPUT_PATH={assets_dir}" + "TENSORBOARD_PATH={assets_dir}/tensorboard" + "CHECKPOINT_SAVE_PATH={artifacts_dir}/checkpoints" + "CHECKPOINT_LOAD_PATH=/mnt/artifacts/model/{name}" + "TRAINING_SCRIPT_PATH=megatron/elastification/pretrain_hybrid_flex.py" + "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml" + "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json" + "N_REPEAT={n_repeat}" + "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}" + "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}" + ) + + bash ./tests/functional_tests/shell_test_utils/run_ci_test.sh ${{ARGUMENTS[@]}} + +products: + - test_case: [hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G] + products: + - environment: [dev] + scope: [nightly] + platforms: [dgx_h100] diff --git a/tests/unit_tests/elastification/__init__.py b/tests/unit_tests/elastification/__init__.py new file mode 100644 index 00000000000..26496bfed70 --- /dev/null +++ b/tests/unit_tests/elastification/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. diff --git a/tests/unit_tests/elastification/test_apply_flextron_elasticity_to_model.py b/tests/unit_tests/elastification/test_apply_flextron_elasticity_to_model.py new file mode 100644 index 00000000000..7f0edd8f6de --- /dev/null +++ b/tests/unit_tests/elastification/test_apply_flextron_elasticity_to_model.py @@ -0,0 +1,236 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for ``apply_flextron_elasticity_to_model``. + +These tests focus on the layer-class-name-based routing logic (which manager +gets attached to which layer type). They use stub nn.Modules so the tests are +pure-Python and run without a GPU or distributed setup. The individual manager +classes are exercised via GPU-backed tests elsewhere. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from megatron.elastification import flextron_elasticity_hooks as hooks_module +from megatron.elastification.flextron_elasticity_hooks import apply_flextron_elasticity_to_model + + +def _make_submod(class_name): + """Build a bare nn.Module with a given __class__.__name__.""" + mod = nn.Module() + mod.__class__ = type(class_name, (nn.Module,), {}) + return mod + + +def _mamba_layer(): + layer = nn.Module() + layer.__class__ = type("MambaLayer", (nn.Module,), {}) + layer.add_module("mixer", _make_submod("MambaMixer")) + return layer + + +def _moe_layer(*, cls="MoETransformerLayer"): + layer = nn.Module() + layer.__class__ = type(cls, (nn.Module,), {}) + layer.add_module("pre_mlp_layernorm", _make_submod("RMSNorm")) + mlp = _make_submod("MoELayer") + mlp.add_module("router", _make_submod("TopKRouter")) + mlp.add_module("experts", _make_submod("TEGroupedMLP")) + layer.add_module("mlp", mlp) + return layer + + +def _attention_layer(): + layer = nn.Module() + layer.__class__ = type("TransformerLayer", (nn.Module,), {}) + attn = _make_submod("SelfAttention") + layer.add_module("self_attention", attn) + return layer + + +class _StubModel(nn.Module): + """Minimal model exposing .decoder.layers and (optionally) .decoder.final_norm.""" + + def __init__(self, layers, with_final_norm=False): + super().__init__() + decoder = nn.Module() + decoder.layers = nn.ModuleList(layers) + if with_final_norm: + decoder.add_module("final_norm", _make_submod("RMSNorm")) + self.decoder = decoder + + +def _make_config(pattern="MEM*", flextron=True): + return SimpleNamespace(hybrid_layer_pattern=pattern, flextron=flextron) + + +@pytest.fixture(autouse=True) +def stub_managers(monkeypatch): + """Stub every add_flextron_* entry point to a call-recorder. + + The real managers attach PyTorch hooks to real submodules; testing the + routing logic does not need that machinery. Each stub returns a Sentinel + whose ``.target`` field points at the module it would have hooked. + """ + calls = { + "transformer_layer": [], + "moe": [], + "topk_router": [], + "grouped_mlp": [], + "mamba": [], + "attention": [], + "stack": [], + } + + def _record(bucket): + def _stub(module, config, layer_idx=None): + entry = SimpleNamespace(target=module, layer_idx=layer_idx, config=config) + calls[bucket].append(entry) + return entry + + return _stub + + def _record_stack(module, config): + entry = SimpleNamespace(target=module, config=config) + calls["stack"].append(entry) + return entry + + monkeypatch.setattr( + hooks_module, "add_flextron_transformer_layer_elasticity", _record("transformer_layer") + ) + monkeypatch.setattr(hooks_module, "add_flextron_moe_elasticity", _record("moe")) + monkeypatch.setattr(hooks_module, "add_flextron_topk_router_elasticity", _record("topk_router")) + monkeypatch.setattr(hooks_module, "add_flextron_grouped_mlp_elasticity", _record("grouped_mlp")) + monkeypatch.setattr(hooks_module, "add_flextron_mamba_elasticity", _record("mamba")) + monkeypatch.setattr(hooks_module, "add_flextron_attention_elasticity", _record("attention")) + monkeypatch.setattr(hooks_module, "add_flextron_stack_elasticity", _record_stack) + + return calls + + +class TestEarlyReturns: + def test_missing_hybrid_pattern_returns_empty(self): + model = _StubModel([_mamba_layer()]) + config = SimpleNamespace() # no hybrid_layer_pattern + assert apply_flextron_elasticity_to_model(model, config) == [] + + def test_empty_hybrid_pattern_returns_empty(self): + model = _StubModel([_mamba_layer()]) + config = _make_config(pattern="") + assert apply_flextron_elasticity_to_model(model, config) == [] + + def test_missing_decoder_returns_empty(self): + model = nn.Module() # no .decoder + config = _make_config() + assert apply_flextron_elasticity_to_model(model, config) == [] + + +class TestLayerRouting: + def test_m_layer_registers_mamba_only(self, stub_managers): + model = _StubModel([_mamba_layer()]) + config = _make_config(pattern="M") + apply_flextron_elasticity_to_model(model, config) + assert len(stub_managers["mamba"]) == 1 + assert stub_managers["mamba"][0].layer_idx == 0 + assert stub_managers["mamba"][0].target.__class__.__name__ == "MambaMixer" + for key in ("transformer_layer", "moe", "topk_router", "grouped_mlp", "attention"): + assert stub_managers[key] == [] + + def test_star_layer_registers_attention_only(self, stub_managers): + model = _StubModel([_attention_layer()]) + config = _make_config(pattern="*") + apply_flextron_elasticity_to_model(model, config) + assert len(stub_managers["attention"]) == 1 + assert stub_managers["attention"][0].target.__class__.__name__ == "SelfAttention" + for key in ("transformer_layer", "moe", "topk_router", "grouped_mlp", "mamba"): + assert stub_managers[key] == [] + + def test_e_layer_registers_all_four_moe_managers(self, stub_managers): + model = _StubModel([_moe_layer()]) + config = _make_config(pattern="E") + apply_flextron_elasticity_to_model(model, config) + assert len(stub_managers["transformer_layer"]) == 1 + assert len(stub_managers["moe"]) == 1 + assert len(stub_managers["topk_router"]) == 1 + assert len(stub_managers["grouped_mlp"]) == 1 + + def test_e_layer_accepts_both_class_names(self, stub_managers): + """Regression: the E-layer hook should fire whether the layer class is + TransformerLayer (modelopt spec) or MoETransformerLayer (default spec).""" + model = _StubModel( + [_moe_layer(cls="TransformerLayer"), _moe_layer(cls="MoETransformerLayer")] + ) + config = _make_config(pattern="EE") + apply_flextron_elasticity_to_model(model, config) + # Both E-layers should have TransformerLayer elasticity attached. + assert len(stub_managers["transformer_layer"]) == 2 + + def test_hybrid_pattern_routes_each_layer(self, stub_managers): + layers = [_mamba_layer(), _moe_layer(), _mamba_layer(), _attention_layer()] + model = _StubModel(layers, with_final_norm=True) + config = _make_config(pattern="MEM*") + apply_flextron_elasticity_to_model(model, config) + + # One mamba manager per M, one attention per *, and all four moe managers per E. + assert len(stub_managers["mamba"]) == 2 + assert len(stub_managers["attention"]) == 1 + assert len(stub_managers["transformer_layer"]) == 1 + assert len(stub_managers["moe"]) == 1 + assert len(stub_managers["topk_router"]) == 1 + assert len(stub_managers["grouped_mlp"]) == 1 + # And a single stack-level manager for the final norm. + assert len(stub_managers["stack"]) == 1 + + +class TestStackManager: + def test_stack_manager_registered_when_final_norm_present(self, stub_managers): + model = _StubModel([_mamba_layer()], with_final_norm=True) + config = _make_config(pattern="M") + apply_flextron_elasticity_to_model(model, config) + assert len(stub_managers["stack"]) == 1 + + def test_stack_manager_skipped_when_no_final_norm(self, stub_managers): + model = _StubModel([_mamba_layer()], with_final_norm=False) + config = _make_config(pattern="M") + apply_flextron_elasticity_to_model(model, config) + assert stub_managers["stack"] == [] + + +class TestMissingSubmodules: + def test_mamba_layer_without_mixer_is_skipped(self, stub_managers): + """M-layer without a MambaMixer submodule should not crash.""" + layer = nn.Module() + layer.__class__ = type("MambaLayer", (nn.Module,), {}) + # intentionally no 'mixer' submodule + model = _StubModel([layer]) + config = _make_config(pattern="M") + apply_flextron_elasticity_to_model(model, config) + assert stub_managers["mamba"] == [] + + def test_attention_layer_without_self_attention_is_skipped(self, stub_managers): + layer = nn.Module() + layer.__class__ = type("TransformerLayer", (nn.Module,), {}) + # no SelfAttention submodule + model = _StubModel([layer]) + config = _make_config(pattern="*") + apply_flextron_elasticity_to_model(model, config) + assert stub_managers["attention"] == [] + + +class TestManagersStoredOnModel: + def test_model_gets_flextron_managers_attribute(self, stub_managers): + model = _StubModel([_mamba_layer()]) + config = _make_config(pattern="M") + returned = apply_flextron_elasticity_to_model(model, config) + assert model._flextron_managers is returned + assert len(returned) == len(stub_managers["mamba"]) + + def test_pattern_shorter_than_layers_only_uses_pattern_length(self, stub_managers): + layers = [_mamba_layer(), _mamba_layer(), _mamba_layer()] + model = _StubModel(layers) + config = _make_config(pattern="M") # only first layer is covered + apply_flextron_elasticity_to_model(model, config) + assert len(stub_managers["mamba"]) == 1 diff --git a/tests/unit_tests/elastification/test_arguments.py b/tests/unit_tests/elastification/test_arguments.py new file mode 100644 index 00000000000..15a726d5498 --- /dev/null +++ b/tests/unit_tests/elastification/test_arguments.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for megatron.elastification.arguments.""" + +from argparse import Namespace + +import pytest + +from megatron.elastification.arguments import ( + convert_per_lists_to_int_lists, + sort_budget_list_descending, + validate_flextron_per_int_lists, +) + + +def _make_config(**overrides): + defaults = dict( + hidden_size=1920, + ffn_hidden_size=960, + num_attention_heads=32, + mamba_num_heads=64, + num_moe_experts=128, + emb_per_list=None, + mlp_per_list=None, + mamba_per_list=None, + moe_expert_per_list=None, + emb_int_list=None, + mlp_int_list=None, + mamba_int_list=None, + moe_expert_int_list=None, + ) + defaults.update(overrides) + return Namespace(**defaults) + + +class TestConvertPerListsToIntLists: + def test_ratio_one_maps_to_full_dim(self): + cfg = _make_config(emb_per_list=[1.0, 0.5]) + convert_per_lists_to_int_lists(cfg) + assert cfg.emb_int_list == [1920, 960] + # After conversion the per-list is cleared. + assert cfg.emb_per_list is None + + def test_floor_rounding(self): + # 0.71429 * 1920 = 1371.4368 -> floor -> 1371 + cfg = _make_config(emb_per_list=[0.71429, 0.51725]) + convert_per_lists_to_int_lists(cfg) + assert cfg.emb_int_list == [1371, 993] + + def test_all_axes_converted_with_correct_ref_dim(self): + cfg = _make_config( + emb_per_list=[1.0], mlp_per_list=[0.5], mamba_per_list=[0.75], moe_expert_per_list=[0.5] + ) + convert_per_lists_to_int_lists(cfg) + assert cfg.emb_int_list == [1920] + assert cfg.mlp_int_list == [480] # 0.5 * 960 + assert cfg.mamba_int_list == [48] # 0.75 * 64 + assert cfg.moe_expert_int_list == [64] # 0.5 * 128 + + def test_axis_with_no_per_list_is_untouched(self): + cfg = _make_config(emb_per_list=[1.0]) # only emb + convert_per_lists_to_int_lists(cfg) + # mlp / mamba / moe_expert should not gain int_list from nothing. + assert cfg.mlp_int_list is None + assert cfg.mamba_int_list is None + assert cfg.moe_expert_int_list is None + + +class TestValidateFlextronPerIntLists: + def _make_args(self, **overrides): + defaults = dict( + emb_per_list=None, + emb_int_list=None, + mlp_per_list=None, + mlp_int_list=None, + mamba_per_list=None, + mamba_int_list=None, + moe_expert_per_list=None, + moe_expert_int_list=None, + ) + defaults.update(overrides) + return Namespace(**defaults) + + def test_unset_axis_defaults_to_full(self): + args = self._make_args() + validate_flextron_per_int_lists(args) + # Each axis defaults to [1.0] on the per-list side. + assert args.emb_per_list == [1.0] + assert args.mlp_per_list == [1.0] + assert args.mamba_per_list == [1.0] + assert args.moe_expert_per_list == [1.0] + + def test_per_list_preserved_when_set(self): + args = self._make_args(emb_per_list=[1.0, 0.5]) + validate_flextron_per_int_lists(args) + assert args.emb_per_list == [1.0, 0.5] + + def test_int_list_preserved_when_set(self): + args = self._make_args(emb_int_list=[1920, 960]) + validate_flextron_per_int_lists(args) + # int_list was explicitly set: per_list stays None (not defaulted to [1.0]). + assert args.emb_per_list is None + assert args.emb_int_list == [1920, 960] + + def test_both_set_raises(self): + args = self._make_args(emb_per_list=[1.0], emb_int_list=[1920]) + with pytest.raises(AssertionError, match="not both"): + validate_flextron_per_int_lists(args) + + def test_per_list_out_of_range_raises(self): + args = self._make_args(emb_per_list=[1.5]) + with pytest.raises(AssertionError, match=r"\[0, 1\]"): + validate_flextron_per_int_lists(args) + + def test_per_list_negative_raises(self): + args = self._make_args(emb_per_list=[-0.1]) + with pytest.raises(AssertionError, match=r"\[0, 1\]"): + validate_flextron_per_int_lists(args) + + +class TestSortBudgetListDescending: + def test_ascending_input_gets_reversed(self): + args = Namespace(budget_list=[0.5, 0.7, 1.0], budget_probs=[0.1, 0.4, 0.5]) + sort_budget_list_descending(args) + assert args.budget_list == [1.0, 0.7, 0.5] + assert args.budget_probs == [0.5, 0.4, 0.1] + + def test_descending_input_unchanged(self): + args = Namespace(budget_list=[1.0, 0.5], budget_probs=[0.7, 0.3]) + sort_budget_list_descending(args) + assert args.budget_list == [1.0, 0.5] + assert args.budget_probs == [0.7, 0.3] + + def test_unsorted_input_paired_correctly(self): + # Verify probs follow the same permutation as budgets. + args = Namespace(budget_list=[0.5, 1.0, 0.7], budget_probs=[0.1, 0.5, 0.4]) + sort_budget_list_descending(args) + assert args.budget_list == [1.0, 0.7, 0.5] + assert args.budget_probs == [0.5, 0.4, 0.1] + + def test_no_probs_only_sorts_budgets(self): + args = Namespace(budget_list=[0.5, 1.0], budget_probs=None) + sort_budget_list_descending(args) + assert args.budget_list == [1.0, 0.5] + assert args.budget_probs is None + + def test_single_element_unchanged(self): + args = Namespace(budget_list=[1.0], budget_probs=[1.0]) + sort_budget_list_descending(args) + assert args.budget_list == [1.0] + assert args.budget_probs == [1.0] + + def test_none_budget_list_skipped(self): + args = Namespace(budget_list=None, budget_probs=None) + sort_budget_list_descending(args) # must not raise + assert args.budget_list is None + + def test_length_mismatch_raises(self): + args = Namespace(budget_list=[1.0, 0.5], budget_probs=[1.0]) + with pytest.raises(AssertionError, match="length"): + sort_budget_list_descending(args) + + def test_idempotent(self): + args = Namespace(budget_list=[0.5, 0.7, 1.0], budget_probs=[0.1, 0.4, 0.5]) + sort_budget_list_descending(args) + sort_budget_list_descending(args) # second call must be a no-op + assert args.budget_list == [1.0, 0.7, 0.5] + assert args.budget_probs == [0.5, 0.4, 0.1] diff --git a/tests/unit_tests/elastification/test_flex_budget_utils.py b/tests/unit_tests/elastification/test_flex_budget_utils.py new file mode 100644 index 00000000000..662b0176c3f --- /dev/null +++ b/tests/unit_tests/elastification/test_flex_budget_utils.py @@ -0,0 +1,155 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for megatron.elastification.router.flex_budget_utils.""" + +import pytest + +from megatron.elastification.router.flex_budget_utils import get_num_parameters + +# Reference dimensions used by most tests. Small enough to compute by hand. +_DIMS = dict( + mamba_num_heads=4, + mamba_d_head=2, + mamba_d_state=2, + num_attention_heads=2, + num_query_groups=1, + ffn_hidden_size=8, + hidden_size=4, + kv_channels=2, + vocab_size=10, + num_experts=2, + shared_expert_intermediate_size=0, + moe_router_topk=1, +) + +_EMBED_PLUS_LN = (_DIMS["vocab_size"] * _DIMS["hidden_size"]) + _DIMS["hidden_size"] +_OUTPUT_LAYER = _DIMS["vocab_size"] * _DIMS["hidden_size"] + + +def _att_cost(): + h, k, q = (_DIMS["hidden_size"], _DIMS["kv_channels"], _DIMS["num_query_groups"]) + n_heads = _DIMS["num_attention_heads"] + input_ln = h + linear_proj = n_heads * k * h + linear_qkv = (n_heads + 2 * q) * k * h + return input_ln + linear_proj + linear_qkv + + +def _moe_cost_all(): + pre_mlp_ln = _DIMS["hidden_size"] + n_experts = _DIMS["num_experts"] + ffn = _DIMS["ffn_hidden_size"] + shared = _DIMS["shared_expert_intermediate_size"] + h = _DIMS["hidden_size"] + linear_fc1 = ffn * (h * n_experts + shared) + linear_fc2 = ffn * (h * n_experts + shared) + return pre_mlp_ln + linear_fc1 + linear_fc2 + + +def _moe_cost_active(): + pre_mlp_ln = _DIMS["hidden_size"] + topk = _DIMS["moe_router_topk"] + ffn = _DIMS["ffn_hidden_size"] + shared = _DIMS["shared_expert_intermediate_size"] + h = _DIMS["hidden_size"] + linear_fc1 = ffn * (h * topk + shared) + linear_fc2 = ffn * (h * topk + shared) + return pre_mlp_ln + linear_fc1 + linear_fc2 + + +def _mamba_cost(): + h = _DIMS["hidden_size"] + nheads = _DIMS["mamba_num_heads"] + d_head = _DIMS["mamba_d_head"] + d_state = _DIMS["mamba_d_state"] + d_inner = nheads * d_head + ngroups = 8 # hard-coded in the implementation + cdim = d_inner + 2 * ngroups * d_state + mamba_conv = cdim + cdim * 1 * 4 # bias + weight, kernel=4, stride=1 + mamba_input_ln = h + mamba_in_proj = h * (d_inner * 2 + 2 * ngroups * d_state + nheads) + mamba_norm = d_inner + mamba_out_proj = d_inner * h + scalars = nheads + nheads + nheads # dt_bias + A_log + D + return scalars + mamba_input_ln + mamba_in_proj + mamba_conv + mamba_norm + mamba_out_proj + + +class TestGetNumParameters: + def test_single_moe_layer_matches_manual(self): + total, active = get_num_parameters(hybrid_pattern="E", tied_vocab=False, **_DIMS) + expected_total = _EMBED_PLUS_LN + _OUTPUT_LAYER + _moe_cost_all() + expected_active = _EMBED_PLUS_LN + _OUTPUT_LAYER + _moe_cost_active() + assert total == expected_total + assert active == expected_active + + def test_single_attention_layer(self): + total, active = get_num_parameters(hybrid_pattern="*", tied_vocab=False, **_DIMS) + expected = _EMBED_PLUS_LN + _OUTPUT_LAYER + _att_cost() + assert total == expected + # Attention has no active/total split. + assert active == expected + + def test_single_mamba_layer(self): + total, active = get_num_parameters(hybrid_pattern="M", tied_vocab=False, **_DIMS) + expected = _EMBED_PLUS_LN + _OUTPUT_LAYER + _mamba_cost() + assert total == expected + assert active == expected + + def test_hybrid_pattern_is_sum_of_per_layer_costs(self): + pattern = "MEM*E" + total, active = get_num_parameters(hybrid_pattern=pattern, tied_vocab=False, **_DIMS) + expected_total = ( + _EMBED_PLUS_LN + _OUTPUT_LAYER + 2 * _mamba_cost() + 2 * _moe_cost_all() + _att_cost() + ) + expected_active = ( + _EMBED_PLUS_LN + + _OUTPUT_LAYER + + 2 * _mamba_cost() + + 2 * _moe_cost_active() + + _att_cost() + ) + assert total == expected_total + assert active == expected_active + + def test_tied_vocab_zeros_output_layer(self): + total_tied, _ = get_num_parameters(hybrid_pattern="M", tied_vocab=True, **_DIMS) + total_untied, _ = get_num_parameters(hybrid_pattern="M", tied_vocab=False, **_DIMS) + # Untied adds one more vocab*hidden block. + assert total_untied - total_tied == _DIMS["vocab_size"] * _DIMS["hidden_size"] + + def test_pipe_character_ignored(self): + # The '|' marker (pipeline split) should not contribute any params. + base = get_num_parameters(hybrid_pattern="ME", tied_vocab=False, **_DIMS) + with_pipe = get_num_parameters(hybrid_pattern="M|E", tied_vocab=False, **_DIMS) + assert base == with_pipe + + def test_unknown_layer_char_raises(self): + with pytest.raises(RuntimeError, match="Unknown layer type"): + get_num_parameters(hybrid_pattern="Z", tied_vocab=False, **_DIMS) + + def test_moe_active_less_than_or_equal_total(self): + # topk < num_experts, so active < total; topk == num_experts, active == total. + total_tk1, active_tk1 = get_num_parameters( + hybrid_pattern="E", tied_vocab=False, **{**_DIMS, "moe_router_topk": 1} + ) + total_tkN, active_tkN = get_num_parameters( + hybrid_pattern="E", + tied_vocab=False, + **{**_DIMS, "moe_router_topk": _DIMS["num_experts"]}, + ) + assert active_tk1 < total_tk1 + assert active_tkN == total_tkN + + def test_topk_zero_active_excludes_experts(self): + # With topk=0 the active cost per expert's linear_fc1/fc2 contribution + # collapses to 0 (shared_expert_intermediate_size=0 in our fixture). + _, active = get_num_parameters( + hybrid_pattern="E", tied_vocab=False, **{**_DIMS, "moe_router_topk": 0} + ) + # active == embed + output + pre_mlp_ln (no fc1/fc2 contribution) + assert active == _EMBED_PLUS_LN + _OUTPUT_LAYER + _DIMS["hidden_size"] + + def test_empty_pattern_only_embeddings_and_final_norm(self): + total, active = get_num_parameters(hybrid_pattern="", tied_vocab=False, **_DIMS) + assert total == _EMBED_PLUS_LN + _OUTPUT_LAYER + assert active == _EMBED_PLUS_LN + _OUTPUT_LAYER diff --git a/tests/unit_tests/elastification/test_flextron_config.py b/tests/unit_tests/elastification/test_flextron_config.py new file mode 100644 index 00000000000..34bf2e8a7f3 --- /dev/null +++ b/tests/unit_tests/elastification/test_flextron_config.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for megatron.elastification.flextron_config.""" + +import dataclasses +from argparse import Namespace +from types import SimpleNamespace + +from megatron.elastification.flextron_config import FlextronConfig, inject_flextron_config + + +class TestFlextronConfigDefaults: + def test_default_values(self): + cfg = FlextronConfig() + assert cfg.flextron is False + assert cfg.enable_router is False + assert cfg.router_inter_dim == 128 + assert cfg.hard_sample_th == 0.996 + assert cfg.tau_init == 1.0 + assert cfg.tau_decay == 0.9999 + assert cfg.router_std == 0.1 + assert cfg.budget_type == 'param' + assert cfg.original_model_sample_prob == 0.33 + + def test_all_fields_accessible_after_construction(self): + cfg = FlextronConfig() + for f in dataclasses.fields(FlextronConfig): + # Every declared field should be readable. + getattr(cfg, f.name) + + +class TestInjectFlextronConfig: + def test_copies_all_fields_from_args(self): + args = Namespace( + flextron=True, + enable_router=True, + router_inter_dim=256, + hard_sample_th=0.5, + tau_init=2.0, + tau_decay=0.9, + router_std=0.01, + budget_type='mem', + budget_list=[1.0, 0.5], + original_model_sample_prob=0.0, + ) + target = SimpleNamespace() + inject_flextron_config(args, target) + assert target.flextron is True + assert target.enable_router is True + assert target.router_inter_dim == 256 + assert target.hard_sample_th == 0.5 + assert target.tau_init == 2.0 + assert target.tau_decay == 0.9 + assert target.router_std == 0.01 + assert target.budget_type == 'mem' + assert target.budget_list == [1.0, 0.5] + assert target.original_model_sample_prob == 0.0 + + def test_missing_arg_falls_back_to_default(self): + # args has only a subset of FlextronConfig fields. + args = Namespace(flextron=True) + target = SimpleNamespace() + inject_flextron_config(args, target) + # Present-on-args field is copied. + assert target.flextron is True + # Absent-on-args field gets FlextronConfig default. + assert target.router_inter_dim == 128 + assert target.hard_sample_th == 0.996 + assert target.tau_init == 1.0 + + def test_preserves_unrelated_config_attributes(self): + args = Namespace(flextron=True) + target = SimpleNamespace(hidden_size=1920, num_layers=52) + inject_flextron_config(args, target) + # Fields that are not FlextronConfig fields stay untouched. + assert target.hidden_size == 1920 + assert target.num_layers == 52 + + def test_every_flextron_field_is_set_on_target(self): + args = Namespace() # totally empty + target = SimpleNamespace() + inject_flextron_config(args, target) + for f in dataclasses.fields(FlextronConfig): + assert hasattr(target, f.name), f"field {f.name!r} not injected onto target" + + def test_returns_none(self): + # inject_flextron_config mutates in place and should not return a value. + args = Namespace(flextron=True) + target = SimpleNamespace() + result = inject_flextron_config(args, target) + assert result is None diff --git a/tests/unit_tests/elastification/test_flextron_grouped_mlp_elasticity_manager.py b/tests/unit_tests/elastification/test_flextron_grouped_mlp_elasticity_manager.py new file mode 100644 index 00000000000..fb8593f7b56 --- /dev/null +++ b/tests/unit_tests/elastification/test_flextron_grouped_mlp_elasticity_manager.py @@ -0,0 +1,211 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GPU-backed tests for FlextronGroupedMLPElasticityManager. + +Covers the multi-hook MLP masking pipeline: setup-mask init, the +input/fc1-post/output hook trio that applies emb + intermediate masking, +and detach. The fc1_post_hook calls into expert-tensor-parallel state, +so we initialize MPU at world_size=1 (mask split is the whole mask). + +Run with: + torchrun --nproc_per_node=1 -m pytest tests/unit_tests/elastification/test_flextron_grouped_mlp_elasticity_manager.py +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from megatron.elastification.flextron_elasticity_hooks import ( + FlextronGroupedMLPElasticityManager, + add_flextron_grouped_mlp_elasticity, +) +from tests.unit_tests.test_utilities import Utils + + +def _config(hidden_size=64, ffn_hidden_size=128, soft_mask=True): + return SimpleNamespace( + flextron=True, + soft_mask=soft_mask, + flex_hetero_ffn=False, + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + emb_int_list=[hidden_size, hidden_size // 2], + mlp_int_list=[ffn_hidden_size, ffn_hidden_size // 2], + hybrid_layer_pattern="E", + layernorm_epsilon=1e-5, + ) + + +class _StubGroupedMLP(nn.Module): + """Minimal module exposing the surface attach_hooks needs: + - register_forward_*_hook (inherited from nn.Module) + - a ``linear_fc1`` child (for fc1_post_hook) + + Forward chain mimics a real GroupedMLP: hidden -> fc1 -> ffn-sized + "intermediate" -> projected back to hidden-sized output. Both stages + return ``(tensor, None)`` so the hooks see the (out, bias) tuple shape + they expect.""" + + def __init__(self, hidden_size, ffn_hidden_size): + super().__init__() + self.hidden_size = hidden_size + self.ffn_hidden_size = ffn_hidden_size + # Stash the captured intermediate so tests can inspect what fc1_post_hook + # produced before the output projection runs. + self._captured_intermediate = None + + class _FC1(nn.Module): + def forward(_self, x): + inter = x.new_zeros(*x.shape[:-1], ffn_hidden_size) + inter[..., : x.shape[-1]] = x # plant the input into the lower channels + return inter, None + + self.linear_fc1 = _FC1() + + def forward(self, hidden_states): + intermediate, _ = self.linear_fc1(hidden_states) + # fc1_post_hook may have masked the intermediate before we get here. + self._captured_intermediate = intermediate.detach().clone() + # Project back to hidden dim: take the lower hidden_size channels. + out = intermediate[..., : self.hidden_size].contiguous() + return (out, None) + + +@pytest.mark.internal +class TestFlextronGroupedMLPElasticityManager: + def setup_method(self, method): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_tensor_parallel_size=1, + ) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _make_module(self, cfg): + return _StubGroupedMLP(cfg.hidden_size, cfg.ffn_hidden_size).cuda().to(torch.bfloat16) + + def test_attach_registers_expected_hook_count(self): + cfg = _config() + mod = self._make_module(cfg) + mgr = FlextronGroupedMLPElasticityManager(cfg) + mgr.attach_hooks(mod) + # setup + input_mask + fc1_post + output_mask + cleanup = 5 + assert len(mgr.hook_handles) == 5 + mgr.detach_hooks() + + def test_init_emb_masks_match_choice_list(self): + cfg = _config(hidden_size=64) + mgr = FlextronGroupedMLPElasticityManager(cfg) + mgr._init_embedding_masks() + # One mask per emb_int_list entry, each shape == [hidden_size]. + assert mgr.emb_masks.shape == (len(cfg.emb_int_list), cfg.hidden_size) + # Mask 0 (full): all-ones over the full hidden dim. + torch.testing.assert_close( + mgr.emb_masks[0], torch.ones(cfg.hidden_size, dtype=torch.bfloat16, device="cuda") + ) + # Mask 1 (half): ones on lower half, zeros on upper half. + expected = torch.zeros(cfg.hidden_size, dtype=torch.bfloat16, device="cuda") + expected[: cfg.hidden_size // 2] = 1.0 + torch.testing.assert_close(mgr.emb_masks[1], expected) + + def test_init_mlp_masks_dedupe_and_sort(self): + """``_init_mlp_masks`` dedupes via set() and sorts descending. Verify + the lookup maps each unique value to the right index.""" + cfg = _config(ffn_hidden_size=128) + cfg.mlp_int_list = [128, 128, 64] # duplicate to exercise dedupe + mgr = FlextronGroupedMLPElasticityManager(cfg) + mgr._init_mlp_masks() + # Two unique values, sorted descending: [128, 64]. + assert mgr.mlp_intermediate_masks.shape[0] == 2 + assert mgr.mlp_intermediate_masks_lookup == {128: 0, 64: 1} + + def test_no_router_emb_is_passthrough(self): + """With current_router_emb None, no hook should mutate output.""" + cfg = _config() + mod = self._make_module(cfg) + x = torch.randn(2, cfg.hidden_size, dtype=torch.bfloat16, device="cuda") + baseline_out, baseline_bias = mod(x) + + mgr = FlextronGroupedMLPElasticityManager(cfg) + mgr.attach_hooks(mod) + # current_router_emb is None — the input/fc1/output hooks all early-out. + out, bias = mod(x) + torch.testing.assert_close(out, baseline_out) + mgr.detach_hooks() + + def test_soft_mask_zeros_upper_intermediate_at_half_budget(self): + """fc1_post_hook applies the mlp_intermediate_mask. Soft-mask weighted + sum on a one-hot at the half-budget choice should leave the upper + ffn channels zeroed.""" + cfg = _config(hidden_size=64, ffn_hidden_size=128, soft_mask=True) + mod = self._make_module(cfg) + + mgr = FlextronGroupedMLPElasticityManager(cfg) + mgr.attach_hooks(mod) + # One-hot router_emb on full-emb, one-hot router_mlp on half-ffn (index 1). + emb_logits = torch.tensor([1.0, 0.0], dtype=torch.bfloat16, device="cuda") + mlp_logits = torch.tensor([0.0, 1.0], dtype=torch.bfloat16, device="cuda") + mgr.set_elasticity_params( + router_emb=(emb_logits, cfg.hidden_size), + router_mlp=(mlp_logits, cfg.ffn_hidden_size // 2), + ) + + x = torch.ones(2, cfg.hidden_size, dtype=torch.bfloat16, device="cuda") + mod(x) + intermediate = mod._captured_intermediate + # mlp_int_list sorted-desc dedupe = [128, 64]; one-hot on index 1 -> 64. + # Lower 64 channels active, upper 64 zeroed by the intermediate mask. + assert (intermediate[..., 64:] == 0).all() + assert not (intermediate[..., :64] == 0).all() + mgr.detach_hooks() + + def test_set_elasticity_params_only_updates_provided_axes(self): + """Calling set_elasticity_params with only one kwarg must not clear + the other (regression-guard for the ``if x is not None`` pattern).""" + cfg = _config() + mgr = FlextronGroupedMLPElasticityManager(cfg) + sentinel_emb = (torch.tensor([1.0, 0.0]), cfg.hidden_size) + sentinel_mlp = (torch.tensor([0.0, 1.0]), cfg.ffn_hidden_size // 2) + mgr.set_elasticity_params(router_emb=sentinel_emb, router_mlp=sentinel_mlp) + + # Update only emb; mlp must still be the prior value. + new_emb = (torch.tensor([0.0, 1.0]), cfg.hidden_size // 2) + mgr.set_elasticity_params(router_emb=new_emb) + + assert mgr.current_router_emb is new_emb + assert mgr.current_router_mlp is sentinel_mlp + + def test_detach_clears_hook_handles(self): + cfg = _config() + mod = self._make_module(cfg) + mgr = FlextronGroupedMLPElasticityManager(cfg) + mgr.attach_hooks(mod) + assert len(mgr.hook_handles) == 5 + mgr.detach_hooks() + assert mgr.hook_handles == [] + + +@pytest.mark.internal +class TestAddFlextronGroupedMLPElasticity: + def setup_method(self, method): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_tensor_parallel_size=1, + ) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_factory_returns_manager_with_layer_idx(self): + cfg = _config() + mod = _StubGroupedMLP(cfg.hidden_size, cfg.ffn_hidden_size).cuda().to(torch.bfloat16) + mgr = add_flextron_grouped_mlp_elasticity(mod, cfg, layer_idx=0) + assert isinstance(mgr, FlextronGroupedMLPElasticityManager) + assert mgr.layer_idx == 0 + assert len(mgr.hook_handles) == 5 + mgr.detach_hooks() diff --git a/tests/unit_tests/elastification/test_flextron_mamba_elasticity_manager.py b/tests/unit_tests/elastification/test_flextron_mamba_elasticity_manager.py new file mode 100644 index 00000000000..81ec4839dbe --- /dev/null +++ b/tests/unit_tests/elastification/test_flextron_mamba_elasticity_manager.py @@ -0,0 +1,206 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GPU-backed tests for FlextronMambaElasticityManager. + +Builds a real MambaMixer (mirroring ``tests/unit_tests/ssm/test_mamba_mixer.py``) +and verifies that the elasticity hooks attach, behave as no-ops without +elasticity params, and produce different activations once params are set. + +Run with: + + torchrun --nproc_per_node=1 -m pytest tests/unit_tests/elastification/test_flextron_mamba_elasticity_manager.py +""" + +import pytest +import torch + +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.mamba_mixer import MambaMixer +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from megatron.elastification.flextron_elasticity_hooks import ( + FlextronMambaElasticityManager, + add_flextron_mamba_elasticity, +) +from tests.unit_tests.test_utilities import Utils + + +def _flextron_fields(hidden_size, num_heads): + """Return dict of flextron attrs to copy onto a TransformerConfig.""" + return dict( + flextron=True, + soft_mask=True, + flex_hetero_mamba=False, + flex_hetero_ffn=False, + flex_hetero_moe_expert=False, + hybrid_layer_pattern="M", + emb_int_list=[hidden_size, hidden_size // 2], + mamba_int_list=[num_heads, num_heads // 2], + ) + + +@pytest.mark.internal +class TestFlextronMambaElasticityManager: + + def setup_method(self, method): + pass + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _build_mixer_and_config(self, hidden_size=256, num_heads=8): + """Construct a bf16 MambaMixer on CUDA + a flextron-enabled config.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + model_parallel_cuda_manual_seed(123) + config = TransformerConfig( + hidden_size=hidden_size, + num_layers=1, + num_attention_heads=1, + use_cpu_initialization=True, + use_mamba_mem_eff_path=True, + ) + # Inject the flextron fields directly (bypassing inject_flextron_config + # to avoid pulling in the whole args-parser stack). + for k, v in _flextron_fields(hidden_size, num_heads).items(): + setattr(config, k, v) + + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + mixer = MambaMixer( + config, + hybrid_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + config.hidden_size, + layer_number=1, + pg_collection=pg_collection, + ) + mixer.cuda() + return mixer, config + + def test_attach_produces_hook_handles(self): + mixer, config = self._build_mixer_and_config() + mgr = FlextronMambaElasticityManager(config) + mgr.attach_hooks(mixer) + # There are several hooks: setup + input_mask + in_proj(pre,post) + + # conv1d + norm(pre,post) + output + cleanup. Exact count can change; + # we just verify > 1. + assert len(mgr.hook_handles) > 1 + mgr.detach_hooks() + + def test_current_router_none_preserves_output(self): + """Without elasticity params set, a forward with hooks attached must + produce (approximately) the same output as a forward without hooks.""" + mixer, config = self._build_mixer_and_config() + + seq_len, micro_batch = 16, 2 + x = torch.ones((seq_len, micro_batch, config.hidden_size), device="cuda") + + # Baseline (no elasticity). + baseline_out, _ = mixer(x) + + mgr = FlextronMambaElasticityManager(config) + mgr.attach_hooks(mixer) + # current_router_emb/mamba both None -> all hooks except setup/cleanup + # should no-op. + hooked_out, _ = mixer(x) + torch.testing.assert_close(hooked_out, baseline_out, atol=1e-2, rtol=1e-2) + mgr.detach_hooks() + + def test_full_budget_one_hot_approx_matches_baseline(self): + """One-hot on index 0 (full emb + full mamba heads) should approximately + reproduce the baseline output (up to bf16 / eps drift).""" + mixer, config = self._build_mixer_and_config() + + seq_len, micro_batch = 16, 2 + x = torch.ones((seq_len, micro_batch, config.hidden_size), device="cuda") + baseline_out, _ = mixer(x) + + mgr = FlextronMambaElasticityManager(config) + mgr.attach_hooks(mixer) + emb_logits = torch.tensor([1.0, 0.0], dtype=torch.bfloat16, device="cuda") + mamba_logits = torch.tensor([1.0, 0.0], dtype=torch.bfloat16, device="cuda") + mgr.set_elasticity_params( + router_emb=(emb_logits, config.hidden_size), + router_mamba=(mamba_logits, config.mamba_int_list[0]), + ) + full_out, _ = mixer(x) + # Full-budget one-hot should not materially change the output. + torch.testing.assert_close(full_out, baseline_out, atol=5e-2, rtol=5e-2) + mgr.detach_hooks() + + def test_small_budget_one_hot_changes_output(self): + """One-hot on a smaller choice should change the output norm.""" + mixer, config = self._build_mixer_and_config() + + seq_len, micro_batch = 16, 2 + x = torch.randn((seq_len, micro_batch, config.hidden_size), device="cuda") + baseline_out, _ = mixer(x) + + mgr = FlextronMambaElasticityManager(config) + mgr.attach_hooks(mixer) + emb_logits = torch.tensor([0.0, 1.0], dtype=torch.bfloat16, device="cuda") + mamba_logits = torch.tensor([0.0, 1.0], dtype=torch.bfloat16, device="cuda") + mgr.set_elasticity_params( + router_emb=(emb_logits, config.emb_int_list[1]), + router_mamba=(mamba_logits, config.mamba_int_list[1]), + ) + small_out, _ = mixer(x) + # The small-budget output should measurably differ from the baseline. + assert not torch.allclose(small_out, baseline_out, atol=1e-2) + mgr.detach_hooks() + + def test_detach_restores_baseline(self): + mixer, config = self._build_mixer_and_config() + seq_len, micro_batch = 16, 2 + x = torch.ones((seq_len, micro_batch, config.hidden_size), device="cuda") + baseline_out, _ = mixer(x) + + mgr = FlextronMambaElasticityManager(config) + mgr.attach_hooks(mixer) + emb_logits = torch.tensor([0.0, 1.0], dtype=torch.bfloat16, device="cuda") + mamba_logits = torch.tensor([0.0, 1.0], dtype=torch.bfloat16, device="cuda") + mgr.set_elasticity_params( + router_emb=(emb_logits, config.emb_int_list[1]), + router_mamba=(mamba_logits, config.mamba_int_list[1]), + ) + _ = mixer(x) + mgr.detach_hooks() + + detached_out, _ = mixer(x) + torch.testing.assert_close(detached_out, baseline_out, atol=1e-2, rtol=1e-2) + + +@pytest.mark.internal +class TestAddFlextronMambaElasticity: + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_factory_returns_manager(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + model_parallel_cuda_manual_seed(123) + config = TransformerConfig( + hidden_size=256, + num_layers=1, + num_attention_heads=1, + use_cpu_initialization=True, + use_mamba_mem_eff_path=True, + ) + for k, v in _flextron_fields(256, 8).items(): + setattr(config, k, v) + + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + mixer = MambaMixer( + config, + hybrid_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + config.hidden_size, + layer_number=1, + pg_collection=pg_collection, + ).cuda() + + mgr = add_flextron_mamba_elasticity(mixer, config, layer_idx=0) + assert isinstance(mgr, FlextronMambaElasticityManager) + assert mgr.layer_idx == 0 + mgr.detach_hooks() diff --git a/tests/unit_tests/elastification/test_flextron_stack_elasticity_manager.py b/tests/unit_tests/elastification/test_flextron_stack_elasticity_manager.py new file mode 100644 index 00000000000..c276b519a8f --- /dev/null +++ b/tests/unit_tests/elastification/test_flextron_stack_elasticity_manager.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GPU-backed tests for FlextronStackElasticityManager. + +Tests the final-norm hooks that apply eps modification and sqrt(emb_per) +scaling when the router supplies an embedding choice. Run with: + + torchrun --nproc_per_node=1 -m pytest tests/unit_tests/elastification/test_flextron_stack_elasticity_manager.py +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from megatron.elastification.flextron_elasticity_hooks import ( + FlextronStackElasticityManager, + add_flextron_stack_elasticity, +) + + +def _stack_config(emb_int_list=(256, 128), soft_mask=True, layernorm_epsilon=1e-5): + """Minimal SimpleNamespace exposing every attr the stack manager reads.""" + return SimpleNamespace( + flextron=True, + soft_mask=soft_mask, + hidden_size=256, + emb_int_list=list(emb_int_list), + layernorm_epsilon=layernorm_epsilon, + ) + + +def _stack_with_final_norm(hidden_size=256, eps=1e-5): + """A stand-in HybridStack: only `final_norm` is hooked, nothing else required.""" + stack = nn.Module() + stack.final_norm = nn.LayerNorm(hidden_size, eps=eps).cuda().to(torch.bfloat16) + return stack + + +@pytest.mark.internal +class TestFlextronStackElasticityManager: + def teardown_method(self, method): + # No parallel state was initialized; nothing to tear down. + pass + + def test_disabled_manager_is_noop(self): + config = _stack_config() + config.flextron = False + mgr = FlextronStackElasticityManager(config) + stack = _stack_with_final_norm() + mgr.attach_hooks(stack) # Should silently skip. + assert mgr.hook_handles == [] if hasattr(mgr, "hook_handles") else True + + def test_attach_registers_two_hooks(self): + config = _stack_config() + mgr = FlextronStackElasticityManager(config) + stack = _stack_with_final_norm() + mgr.attach_hooks(stack) + # One pre-hook + one post-hook on final_norm. + assert len(mgr.hook_handles) == 2 + + def test_current_router_emb_none_is_noop(self): + """Without elasticity params set, hooks must pass through unchanged.""" + config = _stack_config() + mgr = FlextronStackElasticityManager(config) + stack = _stack_with_final_norm() + mgr.attach_hooks(stack) + x = torch.randn(4, 2, 256, dtype=torch.bfloat16, device="cuda") + + expected = stack.final_norm(x) # direct call — hooks do run but should no-op + # Hooks were attached in-place, so call again to capture the hooked output. + out = stack.final_norm(x) + torch.testing.assert_close(out, expected) + + def test_soft_mask_scales_output_by_sqrt_emb_per(self): + """With soft_mask and a one-hot router distribution, output should scale by + sqrt(emb_per) of the selected choice.""" + config = _stack_config(emb_int_list=[256, 128], soft_mask=True) + mgr = FlextronStackElasticityManager(config) + stack = _stack_with_final_norm() + mgr.attach_hooks(stack) + + # One-hot on index 1 (emb_int=128 -> per=0.5) + per_logits = torch.tensor([0.0, 1.0], dtype=torch.bfloat16, device="cuda") + mgr.set_elasticity_params(router_emb=(per_logits, 128)) + + x = torch.randn(4, 2, 256, dtype=torch.bfloat16, device="cuda") + # Baseline without elasticity: detach hooks first. + mgr.detach_hooks() + baseline = stack.final_norm(x) + + # Re-attach and run with elasticity. + mgr.attach_hooks(stack) + mgr.set_elasticity_params(router_emb=(per_logits, 128)) + scaled = stack.final_norm(x) + + # Expected: baseline * sqrt(0.5) (since per_logit is 1.0 on idx 1) + expected_scale = (128 / 256) ** 0.5 + torch.testing.assert_close(scaled, baseline * expected_scale, atol=1e-2, rtol=1e-2) + + def test_detach_removes_all_hooks(self): + config = _stack_config() + mgr = FlextronStackElasticityManager(config) + stack = _stack_with_final_norm() + mgr.attach_hooks(stack) + assert len(mgr.hook_handles) == 2 + mgr.detach_hooks() + assert mgr.hook_handles == [] + + +@pytest.mark.internal +class TestAddFlextronStackElasticity: + def test_factory_returns_manager_with_hooks_attached(self): + config = _stack_config() + stack = _stack_with_final_norm() + mgr = add_flextron_stack_elasticity(stack, config) + assert isinstance(mgr, FlextronStackElasticityManager) + assert len(mgr.hook_handles) == 2 + mgr.detach_hooks() diff --git a/tests/unit_tests/elastification/test_flextron_topk_router_elasticity_manager.py b/tests/unit_tests/elastification/test_flextron_topk_router_elasticity_manager.py new file mode 100644 index 00000000000..f21a78a2b98 --- /dev/null +++ b/tests/unit_tests/elastification/test_flextron_topk_router_elasticity_manager.py @@ -0,0 +1,212 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for FlextronTopKRouterElasticityManager. + +Focuses on the hard-mask path: it replaces the router's ``routing`` method +with a wrapper that masks the upper expert indices before delegating, and +must save/restore ``router.expert_bias`` around the call (regression +guard for the prior permanent-mutation bug). + +The soft-mask path goes through ``topk_softmax_with_capacity``, which +requires real MoE plumbing — covered by integration tests, not here. + +Run with: + torchrun --nproc_per_node=1 -m pytest tests/unit_tests/elastification/test_flextron_topk_router_elasticity_manager.py +""" + +from types import SimpleNamespace + +import pytest +import torch + +from megatron.elastification.flextron_elasticity_hooks import ( + FlextronTopKRouterElasticityManager, + add_flextron_topk_router_elasticity, +) + + +def _config(num_moe_experts=8, soft_mask=False, flex_hetero_moe_expert=False): + return SimpleNamespace( + flextron=True, + soft_mask=soft_mask, + flex_hetero_moe_expert=flex_hetero_moe_expert, + num_moe_experts=num_moe_experts, + moe_expert_int_list=[num_moe_experts, num_moe_experts // 2], + hybrid_layer_pattern="E", + ) + + +class _StubRouter: + """Minimal router: holds an ``expert_bias`` tensor and an ``original_routing`` + method that records its inputs so we can assert what was passed.""" + + def __init__(self, expert_bias): + self.expert_bias = expert_bias + # Record (logits_clone, expert_bias_clone) at call-time so we can verify + # what the inner call observed. + self.calls = [] + + def routing(logits, **kwargs): + self.calls.append( + { + "logits": logits.detach().clone(), + "expert_bias": self.expert_bias.detach().clone(), + "kwargs": kwargs, + } + ) + return logits, kwargs + + self.routing = routing + + +@pytest.mark.internal +class TestFlextronTopKRouterElasticityManager: + def test_attach_replaces_routing_method(self): + cfg = _config() + router = _StubRouter(expert_bias=torch.zeros(8)) + original = router.routing + mgr = FlextronTopKRouterElasticityManager(cfg) + mgr.attach_hooks(router) + assert router.routing is not original + # The handle list records the method-replacement entry for detach. + assert len(mgr.hook_handles) == 1 + assert mgr.hook_handles[0][0] == "method_replacement" + + def test_no_elasticity_params_delegates_to_original(self): + """When current_router_moe_expert is None, wrapped_routing must + forward (logits, kwargs) unchanged to the original method.""" + cfg = _config() + router = _StubRouter(expert_bias=torch.zeros(8)) + mgr = FlextronTopKRouterElasticityManager(cfg) + mgr.attach_hooks(router) + + logits = torch.randn(4, 8) + out_logits, out_kwargs = router.routing(logits, foo="bar") + + assert len(router.calls) == 1 + # Logits passed through untouched. + torch.testing.assert_close(router.calls[0]["logits"], logits) + assert router.calls[0]["kwargs"] == {"foo": "bar"} + torch.testing.assert_close(out_logits, logits) + + def test_hard_mask_truncates_upper_logits(self): + """With expert_int=4 (half), logits[:, 4:] should be -inf when the + original routing sees them, and logits[:, :4] should equal the input + scaled by the router_moe_expert logit (max of one-hot).""" + cfg = _config(num_moe_experts=8, soft_mask=False) + router = _StubRouter(expert_bias=torch.zeros(8)) + mgr = FlextronTopKRouterElasticityManager(cfg) + mgr.attach_hooks(router) + + # One-hot on the half-experts choice (index 1 of moe_expert_int_list = 4 experts). + per_logits = torch.tensor([0.0, 1.0]) + mgr.set_elasticity_params(router_moe_expert=(per_logits, 4)) + + logits = torch.ones(2, 8) + router.routing(logits) + + seen = router.calls[0]["logits"] + # Lower 4 columns: scaled by router_moe_expert_logits = max(per_logits) = 1.0. + torch.testing.assert_close(seen[:, :4], torch.ones(2, 4)) + # Upper 4 columns: -inf. + assert torch.isinf(seen[:, 4:]).all() and (seen[:, 4:] < 0).all() + + def test_hard_mask_preserves_expert_bias_after_call(self): + """Regression: the wrapper must save and restore router.expert_bias. + Previously it left a truncated clone bound to router.expert_bias, + leaking into subsequent forwards.""" + cfg = _config(num_moe_experts=8, soft_mask=False) + original_bias = torch.arange(8, dtype=torch.float32) + 1.0 + router = _StubRouter(expert_bias=original_bias.clone()) + mgr = FlextronTopKRouterElasticityManager(cfg) + mgr.attach_hooks(router) + mgr.set_elasticity_params(router_moe_expert=(torch.tensor([0.0, 1.0]), 4)) + + bias_before = router.expert_bias.clone() + router.routing(torch.zeros(2, 8)) + bias_after = router.expert_bias + + # The bias seen *during* the call should have indices 4: zeroed. + seen_bias = router.calls[0]["expert_bias"] + assert (seen_bias[:4] == bias_before[:4]).all() + assert (seen_bias[4:] == 0).all() + # But the bias on the router after the call must be the original. + torch.testing.assert_close(bias_after, original_bias) + # Same Python object, not just equal values. + assert bias_after is not seen_bias + + def test_hard_mask_bias_restored_even_if_inner_raises(self): + """``try/finally`` must restore expert_bias when original_routing raises.""" + cfg = _config(num_moe_experts=8, soft_mask=False) + original_bias = torch.arange(8, dtype=torch.float32) + 1.0 + router = _StubRouter(expert_bias=original_bias.clone()) + + def boom(logits, **kwargs): + raise RuntimeError("simulated downstream failure") + + router.routing = boom + + mgr = FlextronTopKRouterElasticityManager(cfg) + mgr.attach_hooks(router) + mgr.set_elasticity_params(router_moe_expert=(torch.tensor([0.0, 1.0]), 4)) + + with pytest.raises(RuntimeError, match="simulated"): + router.routing(torch.zeros(2, 8)) + + torch.testing.assert_close(router.expert_bias, original_bias) + + def test_hard_mask_with_no_expert_bias(self): + """When router has no expert_bias, the save/restore branch must skip + cleanly and the inner call must still see the masked logits.""" + cfg = _config(num_moe_experts=8, soft_mask=False) + + # Minimal router: no expert_bias attribute, original_routing records + # only the logits it saw (avoids the StubRouter's bias.detach()). + class _RouterNoBias: + def __init__(self): + self.expert_bias = None + self.seen = None + + def routing(logits, **kwargs): + self.seen = logits.detach().clone() + return logits, kwargs + + self.routing = routing + + router = _RouterNoBias() + mgr = FlextronTopKRouterElasticityManager(cfg) + mgr.attach_hooks(router) + mgr.set_elasticity_params(router_moe_expert=(torch.tensor([0.0, 1.0]), 4)) + + router.routing(torch.ones(2, 8)) + assert torch.isinf(router.seen[:, 4:]).all() + # The bias attribute must remain None — no accidental clone-binding. + assert router.expert_bias is None + + def test_detach_restores_original_routing(self): + cfg = _config() + router = _StubRouter(expert_bias=torch.zeros(8)) + original_callable = router.routing # capture pre-attach reference + mgr = FlextronTopKRouterElasticityManager(cfg) + mgr.attach_hooks(router) + wrapped = router.routing + assert wrapped is not original_callable + + mgr.detach_hooks() + + # After detach: routing is back, the helper attribute is gone, the + # handle list is empty. + assert router.routing is original_callable + assert not hasattr(router, "_original_routing") + assert mgr.hook_handles == [] + + +@pytest.mark.internal +class TestAddFlextronTopKRouterElasticity: + def test_factory_returns_manager(self): + cfg = _config() + router = _StubRouter(expert_bias=torch.zeros(8)) + mgr = add_flextron_topk_router_elasticity(router, cfg, layer_idx=0) + assert isinstance(mgr, FlextronTopKRouterElasticityManager) + assert len(mgr.hook_handles) == 1 + mgr.detach_hooks() diff --git a/tests/unit_tests/elastification/test_flextron_transformer_layer_elasticity_manager.py b/tests/unit_tests/elastification/test_flextron_transformer_layer_elasticity_manager.py new file mode 100644 index 00000000000..c14e187cfb1 --- /dev/null +++ b/tests/unit_tests/elastification/test_flextron_transformer_layer_elasticity_manager.py @@ -0,0 +1,141 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GPU-backed tests for FlextronTransformerLayerElasticityManager. + +Tests the pre_mlp_layernorm pre/post hooks for E-layers. Run with: + + torchrun --nproc_per_node=1 -m pytest tests/unit_tests/elastification/test_flextron_transformer_layer_elasticity_manager.py +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from megatron.elastification.flextron_elasticity_hooks import ( + FlextronTransformerLayerElasticityManager, + add_flextron_transformer_layer_elasticity, +) + + +def _tl_config(emb_int_list=(256, 128), soft_mask=True, layernorm_epsilon=1e-5): + return SimpleNamespace( + flextron=True, + soft_mask=soft_mask, + hidden_size=256, + emb_int_list=list(emb_int_list), + layernorm_epsilon=layernorm_epsilon, + ) + + +def _fake_transformer_layer(hidden_size=256, eps=1e-5): + """Minimal module exposing .pre_mlp_layernorm (the only submodule hooked).""" + layer = nn.Module() + layer.pre_mlp_layernorm = nn.LayerNorm(hidden_size, eps=eps).cuda().to(torch.bfloat16) + return layer + + +@pytest.mark.internal +class TestFlextronTransformerLayerElasticityManager: + def teardown_method(self, method): + pass + + def test_attach_registers_two_hooks(self): + config = _tl_config() + mgr = FlextronTransformerLayerElasticityManager(config) + layer = _fake_transformer_layer() + mgr.attach_hooks(layer) + assert len(mgr.hook_handles) == 2 + + def test_current_router_emb_none_is_noop(self): + """With current_router_emb unset, hook behavior must match no-hook forward.""" + config = _tl_config() + layer = _fake_transformer_layer() + x = torch.randn(4, 2, 256, dtype=torch.bfloat16, device="cuda") + expected = layer.pre_mlp_layernorm(x) + + mgr = FlextronTransformerLayerElasticityManager(config) + mgr.attach_hooks(layer) + # current_router_emb is None — no masking, no scaling. + out = layer.pre_mlp_layernorm(x) + torch.testing.assert_close(out, expected) + + def test_soft_mask_scales_output(self): + """With soft_mask one-hot on a smaller choice, the hook should: + (1) zero input channels beyond the chosen emb_int, then + (2) scale the LN output by sqrt(emb_per). + Reproduce the expected output by applying that mask+scale manually + against plain LN and assert equality (within bf16 tolerance).""" + config = _tl_config(emb_int_list=[256, 128], soft_mask=True) + layer = _fake_transformer_layer() + + x = torch.randn(4, 2, 256, dtype=torch.bfloat16, device="cuda") + + # Build expected output: mask upper half, LN, scale by sqrt(emb_per). + mask = torch.zeros(256, dtype=torch.bfloat16, device="cuda") + mask[:128] = 1.0 + expected = layer.pre_mlp_layernorm(x * mask[None, None, :]) * (128 / 256) ** 0.5 + + mgr = FlextronTransformerLayerElasticityManager(config) + mgr.attach_hooks(layer) + # One-hot on index 1 (emb_int=128 -> per=0.5) + per_logits = torch.tensor([0.0, 1.0], dtype=torch.bfloat16, device="cuda") + mgr.set_elasticity_params(router_emb=(per_logits, 128)) + out = layer.pre_mlp_layernorm(x) + + # Tolerance accommodates the tiny eps drift (5e-6 vs 1e-5) inside LN. + torch.testing.assert_close(out, expected, atol=1e-2, rtol=1e-2) + + def test_full_budget_one_hot_preserves_magnitude_order(self): + """When router is one-hot on full budget (index 0 = 100% emb), the + pre-hook masks nothing and post-hook scales by sqrt(1.0)=1.0.""" + config = _tl_config(emb_int_list=[256, 128], soft_mask=True) + layer = _fake_transformer_layer() + x = torch.randn(4, 2, 256, dtype=torch.bfloat16, device="cuda") + baseline = layer.pre_mlp_layernorm(x) + + mgr = FlextronTransformerLayerElasticityManager(config) + mgr.attach_hooks(layer) + per_logits = torch.tensor([1.0, 0.0], dtype=torch.bfloat16, device="cuda") + mgr.set_elasticity_params(router_emb=(per_logits, 256)) + + out = layer.pre_mlp_layernorm(x) + # Full-budget path: input mask is all-ones, scale is sqrt(1.0). Output + # should equal baseline within bf16 tolerance (eps adjustment may add + # tiny drift). + torch.testing.assert_close(out, baseline, atol=5e-2, rtol=5e-2) + + def test_detach_restores_forward(self): + config = _tl_config() + layer = _fake_transformer_layer() + x = torch.randn(4, 2, 256, dtype=torch.bfloat16, device="cuda") + + mgr = FlextronTransformerLayerElasticityManager(config) + mgr.attach_hooks(layer) + per_logits = torch.tensor([0.0, 1.0], dtype=torch.bfloat16, device="cuda") + mgr.set_elasticity_params(router_emb=(per_logits, 128)) + + masked_out = layer.pre_mlp_layernorm(x) + mgr.detach_hooks() + detached_out = layer.pre_mlp_layernorm(x) + + # After detach, the output should match the un-hooked LN output. + expected = nn.LayerNorm(256, eps=layer.pre_mlp_layernorm.eps).cuda().to(torch.bfloat16) + expected.weight.data.copy_(layer.pre_mlp_layernorm.weight.data) + expected.bias.data.copy_(layer.pre_mlp_layernorm.bias.data) + torch.testing.assert_close(detached_out, expected(x), atol=1e-2, rtol=1e-2) + # The masked output from before detach should differ from the detached one. + assert not torch.allclose(masked_out, detached_out, atol=1e-2) + + +@pytest.mark.internal +class TestAddFlextronTransformerLayerElasticity: + def test_factory_returns_manager(self): + config = _tl_config() + layer = _fake_transformer_layer() + mgr = add_flextron_transformer_layer_elasticity(layer, config, layer_idx=3) + assert isinstance(mgr, FlextronTransformerLayerElasticityManager) + assert mgr.layer_idx == 3 + assert len(mgr.hook_handles) == 2 + mgr.detach_hooks() diff --git a/tests/unit_tests/elastification/test_hybrid_flex_router.py b/tests/unit_tests/elastification/test_hybrid_flex_router.py new file mode 100644 index 00000000000..a5720e92dc4 --- /dev/null +++ b/tests/unit_tests/elastification/test_hybrid_flex_router.py @@ -0,0 +1,212 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GPU-backed tests for FlextronRouter. + +Covers construction, forward-pass shape/structure for each axis, and +DP-aware Gumbel determinism (same seed + iteration => identical output). + +Run with: + + torchrun --nproc_per_node=1 -m pytest tests/unit_tests/elastification/test_hybrid_flex_router.py +""" + +from argparse import Namespace + +import pytest +import torch + +import megatron.elastification.router.hybrid_flex_router as _router_module +import megatron.training as _megatron_training +from megatron.core.transformer import TransformerConfig +from megatron.elastification.router.hybrid_flex_router import FlextronRouter +from tests.unit_tests.test_utilities import Utils + + +def _router_config( + hidden_size=256, ffn_hidden_size=128, num_heads=8, mamba_num_heads=8, num_moe_experts=8 +): + """Build a TransformerConfig with every attr FlextronRouter reads.""" + config = TransformerConfig( + hidden_size=hidden_size, + num_layers=2, + num_attention_heads=num_heads, + ffn_hidden_size=ffn_hidden_size, + num_moe_experts=num_moe_experts, + use_cpu_initialization=True, + ) + flex_fields = dict( + flextron=True, + soft_mask=True, + add_skipping=False, + flex_hetero_ffn=False, + flex_hetero_mamba=False, + flex_hetero_moe_expert=False, + hybrid_layer_pattern="ME", + normalize_router_logits=False, + router_inter_dim=32, + router_std=0.1, + router_gbs=2, + router_beta=1.0, + loss_alpha=1.0, + tau_init=1.0, + tau_decay=0.9999, + hard_sample_th=0.996, + # Enable the scaler with a constant 1.0 so `scale` is defined inside + # the axis forwards (they use it unconditionally) but its value is a + # no-op. The get_args stub in setup_method supplies train_iters so + # add_scaler_schedule can construct the linspace. + linear_scaler_start=1.0, + linear_scaler_end=1.0, + budget_list=[1.0, 0.5], + budget_probs=[1.0, 1.0], + budget_type="param", + original_model_sample_prob=0.0, + curr_iteration=0, + mamba_num_heads=mamba_num_heads, + emb_int_list=[hidden_size, hidden_size // 2], + mlp_int_list=[ffn_hidden_size, ffn_hidden_size // 2], + mamba_int_list=[mamba_num_heads, mamba_num_heads // 2], + moe_expert_int_list=[num_moe_experts, num_moe_experts // 2], + override_selected_budget=None, + ) + for k, v in flex_fields.items(): + setattr(config, k, v) + return config + + +@pytest.mark.internal +class TestFlextronRouter: + def setup_method(self, method): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + # FlextronRouter calls _sync_router_weights in __init__, which does an + # NCCL broadcast on CPU params before we get a chance to call .cuda(). + # NCCL has no CPU backend, so we stub the sync out — it's a no-op at + # world_size=1 anyway. Restored in teardown. + self._orig_sync = FlextronRouter._sync_router_weights + FlextronRouter._sync_router_weights = lambda self: None + # forward() pulls `args.curr_iteration` via megatron.training.get_args() + # which fails outside a full Megatron initialize. Install a minimal + # stub that returns the attrs the router needs. + self._orig_get_args = _megatron_training.get_args + # train_iters needs to be >= max curr_iteration used in any test: the + # scaler is a linspace of length train_iters and axis forwards index + # it with curr_iteration. Since start=end=1.0 the values are all 1.0 + # regardless of length, so overshooting is free. + _megatron_training.get_args = lambda: Namespace( + curr_iteration=0, train_iters=1000, train_samples=1000, global_batch_size=1 + ) + # The router's __init__ also reads global microbatch state via + # get_current_global_batch_size / get_micro_batch_size. These return + # None outside a full Megatron initialize — stub them at the module + # level where the router imported them. + self._orig_gbs = _router_module.get_current_global_batch_size + self._orig_mbs = _router_module.get_micro_batch_size + _router_module.get_current_global_batch_size = lambda: 1 + _router_module.get_micro_batch_size = lambda: 1 + + def teardown_method(self, method): + FlextronRouter._sync_router_weights = self._orig_sync + _megatron_training.get_args = self._orig_get_args + _router_module.get_current_global_batch_size = self._orig_gbs + _router_module.get_micro_batch_size = self._orig_mbs + Utils.destroy_model_parallel() + + def test_construction(self): + config = _router_config() + router = FlextronRouter(config).cuda() + # Each gate is a Sequential of two linear layers + activation. + assert hasattr(router, "gate_mlp") + assert hasattr(router, "gate_emb") + assert hasattr(router, "gate_mamba") + assert hasattr(router, "gate_moe_expert") + # Attention head elasticity is not supported. + assert not hasattr(router, "gate_head") + # Skipping was disabled in the config. + assert not hasattr(router, "gate_skip_layer") + + def test_router_params_marked_for_pp_sync(self): + config = _router_config() + router = FlextronRouter(config).cuda() + for p in router.parameters(): + # _mark_router_params_for_pp_sync adds this attribute to every + # trainable parameter so the PP gradient sync picks them up. + assert getattr(p, "flextron_router_pp_sync", False) is True + + def test_forward_returns_five_axis_outputs(self): + config = _router_config() + router = FlextronRouter(config).cuda() + out = router(1.0) + assert len(out) == 5 + # Order (per hybrid_flex_router.forward): + # (mlp, skipping, emb, mamba, moe_expert) + mlp, skipping, emb, mamba, moe_expert = out + # Skipping is None when add_skipping=False. + assert skipping is None + # Each axis output is a (logits, choice) tuple. + for axis in (mlp, emb, mamba, moe_expert): + assert isinstance(axis, tuple) and len(axis) == 2 + + def test_emb_output_shape_matches_choice_count(self): + config = _router_config() + router = FlextronRouter(config).cuda() + _, _, emb, _, _ = router(1.0) + logits, choice = emb + # Logits have one entry per emb_int_list choice. + assert logits.numel() == len(config.emb_int_list) + assert choice in config.emb_int_list + + def test_gumbel_determinism(self): + """Two routers at the same config + iteration + fwd_pass_count should + produce identical Gumbel-softmax samples.""" + config = _router_config() + config.curr_iteration = 0 + + router_a = FlextronRouter(config).cuda() + router_b = FlextronRouter(config).cuda() + # Copy weights so both routers are in the same parameter state; the + # determinism check is about the Gumbel RNG, not init noise. + router_b.load_state_dict(router_a.state_dict()) + + out_a = router_a(1.0) + out_b = router_b(1.0) + for axis_a, axis_b in zip(out_a, out_b): + if axis_a is None: + assert axis_b is None + continue + logits_a, choice_a = axis_a + logits_b, choice_b = axis_b + torch.testing.assert_close(logits_a, logits_b, atol=0, rtol=0) + assert choice_a == choice_b + + def test_fwd_pass_count_increments(self): + config = _router_config() + router = FlextronRouter(config).cuda() + assert router.fwd_pass_count == 0 + router(1.0) + assert router.fwd_pass_count == 1 + router(1.0) + assert router.fwd_pass_count == 2 + + def test_different_iterations_give_different_samples(self): + """Bumping curr_iteration changes the Gumbel seed; logits should differ.""" + config = _router_config() + router = FlextronRouter(config).cuda() + + # Iteration 0 via the default setup-method stub. + out_iter_0 = router(1.0) + + # Swap the stub to return iteration 100, reset fwd_pass_count so + # that is the only thing that varies. train_iters must stay >= + # curr_iteration (matches setup-method stub length). + _megatron_training.get_args = lambda: Namespace( + curr_iteration=100, train_iters=1000, train_samples=1000, global_batch_size=1 + ) + router.fwd_pass_count = 0 + out_iter_100 = router(1.0) + + # Emb-axis logits should differ between iterations. + _, _, emb_0, _, _ = out_iter_0 + _, _, emb_100, _, _ = out_iter_100 + assert not torch.allclose(emb_0[0], emb_100[0]) diff --git a/tests/unit_tests/elastification/test_inject_flextron_forward_logic.py b/tests/unit_tests/elastification/test_inject_flextron_forward_logic.py new file mode 100644 index 00000000000..c802cdb1b9a --- /dev/null +++ b/tests/unit_tests/elastification/test_inject_flextron_forward_logic.py @@ -0,0 +1,217 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for ``inject_flextron_forward_logic``. + +These tests pin two invariants the dev-branch divergence hunt uncovered: + +1. When a Flextron manager is attached and a budget is passed in kwargs, + ``update_hook_elasticity_params`` must be called *before* the original + forward runs — otherwise all hooks fire with ``current_router_emb=None`` + and silently no-op. +2. When no Flextron manager is present, ``flextron_kwargs`` is cleared + before the original forward so no unexpected keyword args leak through. +""" + +from types import SimpleNamespace + +import pytest + +from megatron.elastification.flextron_utils import inject_flextron_forward_logic + + +class _CallLog: + """Record the order in which stub methods are invoked.""" + + def __init__(self): + self.events = [] + + def note(self, name, **payload): + self.events.append((name, payload)) + + +def _make_original_forward(log): + def _fwd( + input_ids=None, + position_ids=None, + attention_mask=None, + decoder_input=None, + labels=None, + inference_context=None, + runtime_gather_output=None, + inference_params=None, + ): + log.note( + "original_forward", + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + ) + return "forward-result" + + return _fwd + + +def _make_manager(log, *, router_present=True, router_output_kwargs=None, loss_func=None): + router_output_kwargs = ( + router_output_kwargs + if router_output_kwargs is not None + else {"router_emb": (object(), object())} + ) + + def process_router_output(budget_item): + log.note("process_router_output", budget_item=budget_item) + return router_output_kwargs, loss_func + + def update_hook_elasticity_params(flextron_kwargs): + log.note("update_hook_elasticity_params", flextron_kwargs=flextron_kwargs) + + return SimpleNamespace( + router=object() if router_present else None, + process_router_output=process_router_output, + update_hook_elasticity_params=update_hook_elasticity_params, + ) + + +def _attach_forward(model_cls, original): + """Return a model whose .forward is `original` and install the wrapper. + + `original` is stored as an instance attribute so Python does not auto-bind + it as a method — we want the stub called as a plain function by + ``flextron_forward``. + """ + model = model_cls() + model.forward = original + inject_flextron_forward_logic(model) + return model + + +class _StubModelNoManager: + """Minimal model with no ._flextron_manager attribute at all.""" + + def __init__(self): + self.config = SimpleNamespace() + + +class _StubModelWithManager: + def __init__(self): + self.config = SimpleNamespace(flextron=True, is_flex_eval=False) + self._flextron_manager = None # filled in by test + + +class TestForwardReplacement: + def test_forward_is_replaced(self): + log = _CallLog() + original = _make_original_forward(log) + model = _StubModelNoManager() + model.forward = original + before = model.forward + inject_flextron_forward_logic(model) + # model.forward is now a bound method wrapper, not the raw stub. + assert model.forward is not before + + +class TestNoManager: + def test_without_manager_original_is_called_directly(self): + log = _CallLog() + original = _make_original_forward(log) + model = _attach_forward(_StubModelNoManager, original) + + result = model.forward( + input_ids="ids", + position_ids="pos", + attention_mask="mask", + budget=0.697, # should be swallowed, not leaked through + ) + + assert result == "forward-result" + # Only original_forward was called; no router / hook-update step. + names = [e[0] for e in log.events] + assert names == ["original_forward"] + # budget kwarg was cleared before reaching original_forward. + assert "budget" not in log.events[0][1] + + def test_manager_with_no_router_skips_router_logic(self): + log = _CallLog() + original = _make_original_forward(log) + model = _attach_forward(_StubModelWithManager, original) + model._flextron_manager = _make_manager(log, router_present=False) + + model.forward(input_ids="ids", position_ids="pos", attention_mask="mask", budget=0.5) + + names = [e[0] for e in log.events] + assert names == ["original_forward"] + + +class TestManagerOrdering: + def test_budget_kwarg_triggers_router_and_hooks_before_forward(self): + """Core invariant: update_hook_elasticity_params runs *before* original_forward.""" + log = _CallLog() + original = _make_original_forward(log) + model = _attach_forward(_StubModelWithManager, original) + model._flextron_manager = _make_manager(log) + + model.forward(input_ids="ids", position_ids="pos", attention_mask="mask", budget=0.697) + + names = [e[0] for e in log.events] + # Exact expected sequence. + assert names == [ + "process_router_output", + "update_hook_elasticity_params", + "original_forward", + ] + # Sanity: the budget actually used is the one passed in kwargs. + assert log.events[0][1]["budget_item"] == 0.697 + + def test_loss_func_invoked_when_returned(self): + log = _CallLog() + original = _make_original_forward(log) + model = _attach_forward(_StubModelWithManager, original) + + def _loss_func(kwargs, budget_item): + log.note("loss_func", budget_item=budget_item) + return "budget-loss" + + model._flextron_manager = _make_manager(log, loss_func=_loss_func) + + model.forward(input_ids="ids", position_ids="pos", attention_mask="mask", budget=0.697) + + names = [e[0] for e in log.events] + # loss_func must run after router output and before the hook update. + assert names == [ + "process_router_output", + "loss_func", + "update_hook_elasticity_params", + "original_forward", + ] + + +class TestOverrideSelectedBudget: + def test_override_non_one_sets_budget_from_override(self): + log = _CallLog() + original = _make_original_forward(log) + model = _attach_forward(_StubModelWithManager, original) + model._flextron_manager = _make_manager(log) + model.config.is_flex_eval = True + model.config.override_selected_budget = [0.577] + + # No budget kwarg on the caller side — override should supply it. + model.forward(input_ids="ids", position_ids="pos", attention_mask="mask") + + names = [e[0] for e in log.events] + assert names == [ + "process_router_output", + "update_hook_elasticity_params", + "original_forward", + ] + assert log.events[0][1]["budget_item"] == 0.577 + + def test_override_without_flex_eval_raises(self): + log = _CallLog() + original = _make_original_forward(log) + model = _attach_forward(_StubModelWithManager, original) + model._flextron_manager = _make_manager(log) + model.config.is_flex_eval = False + model.config.override_selected_budget = [0.577] + + with pytest.raises(AssertionError): + model.forward(input_ids="ids", position_ids="pos", attention_mask="mask") diff --git a/tests/unit_tests/elastification/test_loss_func.py b/tests/unit_tests/elastification/test_loss_func.py new file mode 100644 index 00000000000..d5f27aa1ba8 --- /dev/null +++ b/tests/unit_tests/elastification/test_loss_func.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for megatron.elastification.loss_func. + +Covers the non-distributed paths of ``_mask_loss`` and ``loss_func``: +no tensor-parallel reductions, no sequence parallel, no KD. The +distributed/KD branches require a multi-rank torch.distributed init and +are left for a follow-up. +""" + +from argparse import Namespace +from unittest.mock import MagicMock + +import pytest +import torch + +from megatron.elastification import loss_func as loss_func_module +from megatron.elastification.loss_func import _mask_loss, loss_func + + +def _stub_args(**overrides): + defaults = dict( + router_beta=1.0, + loss_alpha=1.0, + freeze_router=False, + tensor_model_parallel_size=1, + export_kd_teacher_load=False, + budget_list=[1.0, 0.5], + ) + defaults.update(overrides) + return Namespace(**defaults) + + +@pytest.fixture +def patch_get_args(monkeypatch): + """Replace get_args in the loss_func module with a stub returning the args we set.""" + holder = {"args": _stub_args()} + + def _set_args(**overrides): + holder["args"] = _stub_args(**overrides) + + monkeypatch.setattr(loss_func_module, "get_args", lambda: holder["args"]) + return _set_args + + +def _flat_loss_tensor(values): + """Build the (B, S) loss tensor that ``_mask_loss`` expects.""" + return torch.tensor(values, dtype=torch.float32).reshape(1, -1) + + +class TestMaskLossPlainTensor: + """When output_tensor is a plain Tensor, no param_loss is reported.""" + + def test_returns_scalar_loss_tensor(self, patch_get_args): + patch_get_args() + out = torch.tensor([[1.0, 2.0, 4.0]]) + mask = torch.tensor([[1.0, 1.0, 0.0]]) + result = _mask_loss(out, mask) + assert isinstance(result, torch.Tensor) + assert result.item() == pytest.approx(3.0) # 1*1 + 2*1 + 4*0 + + +class TestMaskLossWithParamLossTuple: + """output_tensor is (output, (param_loss, extra_dict)).""" + + def test_positive_param_loss_added_to_lm(self, patch_get_args): + patch_get_args(loss_alpha=1.0) + out = torch.tensor([[1.0, 2.0]]) + mask = torch.tensor([[1.0, 1.0]]) + param_loss = torch.tensor([0.5]) + loss, param_item = _mask_loss((out, (param_loss, {})), mask) + # lm = 1+2 = 3; param contribution = 0.5 * num_tokens(2) * alpha(1) = 1.0 + assert loss.item() == pytest.approx(4.0) + assert param_item.item() == pytest.approx(1.0) + + def test_negative_param_loss_scaled_by_router_beta(self, patch_get_args): + # router_beta flips and scales negative param losses. + patch_get_args(router_beta=2.0, loss_alpha=1.0) + out = torch.tensor([[1.0, 1.0]]) + mask = torch.tensor([[1.0, 1.0]]) + param_loss = torch.tensor([-0.25]) + loss, param_item = _mask_loss((out, (param_loss, {})), mask) + # param_loss negated and scaled: -2.0 * (-0.25) = 0.5 + # param_item = 0.5 * num_tokens(2) * alpha(1) = 1.0 + # lm contribution = 2; total = 3 + assert param_item.item() == pytest.approx(1.0) + assert loss.item() == pytest.approx(3.0) + + def test_freeze_router_drops_param_contribution(self, patch_get_args): + patch_get_args(freeze_router=True) + out = torch.tensor([[1.0, 1.0]]) + mask = torch.tensor([[1.0, 1.0]]) + param_loss = torch.tensor([0.5]) + result = _mask_loss((out, (param_loss, {})), mask) + # When router is frozen, param_loss isn't added — bare scalar returned. + assert isinstance(result, torch.Tensor) + assert result.item() == pytest.approx(2.0) + + def test_loss_alpha_scales_param_contribution(self, patch_get_args): + patch_get_args(loss_alpha=10.0) + out = torch.tensor([[1.0, 1.0]]) + mask = torch.tensor([[1.0, 1.0]]) + param_loss = torch.tensor([0.5]) + loss, param_item = _mask_loss((out, (param_loss, {})), mask) + # param_item = 0.5 * 2 tokens * 10 = 10.0 + assert param_item.item() == pytest.approx(10.0) + assert loss.item() == pytest.approx(12.0) # 2 (lm) + 10 (param) + + +class TestLossFuncReportingNoKD: + """Top-level loss_func paths that don't enter the KD branch.""" + + def _model(self, training=True): + m = MagicMock() + m.training = training + return m + + def test_full_model_step_routes_to_lm_loss_full(self, patch_get_args): + patch_get_args() + out = torch.tensor([[1.0, 2.0]]) + mask = torch.tensor([[1.0, 1.0]]) + # param_loss = 0 → recognized as full-model step + zero_param = torch.tensor([0.0]) + loss, num_tokens, report = loss_func( + mask, (out, (zero_param, {})), self._model(training=True) + ) + # The report dict must contain both keys, but only "(full)" carries data. + assert "lm loss (full)" in report and "lm loss (budget)" in report + full_val, full_den = report["lm loss (full)"][0], report["lm loss (full)"][1] + budget_val, budget_den = report["lm loss (budget)"][0], report["lm loss (budget)"][1] + assert budget_val.item() == 0.0 and budget_den.item() == 0.0 + # full_val gets lm loss minus param contribution. param_loss=0 → just lm. + assert full_val.item() == pytest.approx(3.0) + assert num_tokens.item() == 2 + + def test_sub_budget_step_routes_to_lm_loss_budget(self, patch_get_args): + patch_get_args() + out = torch.tensor([[1.0, 2.0]]) + mask = torch.tensor([[1.0, 1.0]]) + nonzero_param = torch.tensor([0.5]) # signals sub-budget step + loss, num_tokens, report = loss_func( + mask, (out, (nonzero_param, {})), self._model(training=True) + ) + full_val = report["lm loss (full)"][0] + budget_val = report["lm loss (budget)"][0] + assert full_val.item() == 0.0 + # budget side carries the lm loss (3.0 = 1+2) + assert budget_val.item() == pytest.approx(3.0) + + def test_num_tokens_clamped_when_all_masked(self, patch_get_args): + patch_get_args() + out = torch.tensor([[5.0, 5.0]]) + mask = torch.tensor([[0.0, 0.0]]) + zero_param = torch.tensor([0.0]) + _, num_tokens, _ = loss_func(mask, (out, (zero_param, {})), self._model(training=True)) + # Guard at line 94 clamps to min=1 to avoid divide-by-zero downstream. + assert num_tokens.item() == 1 + + def test_report_values_are_packed_pairs(self, patch_get_args): + """Every report entry is converted to a (value, num_tokens) tensor pair.""" + patch_get_args() + out = torch.tensor([[1.0, 1.0]]) + mask = torch.tensor([[1.0, 1.0]]) + zero_param = torch.tensor([0.0]) + _, _, report = loss_func(mask, (out, (zero_param, {})), self._model(training=True)) + for key, val in report.items(): + assert isinstance(val, torch.Tensor), f"report[{key}] not packed" + assert val.shape == (2,), f"report[{key}] not (value, num_tokens)" diff --git a/tests/unit_tests/elastification/test_memory_config.py b/tests/unit_tests/elastification/test_memory_config.py new file mode 100644 index 00000000000..70ea488108c --- /dev/null +++ b/tests/unit_tests/elastification/test_memory_config.py @@ -0,0 +1,128 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for megatron.elastification.memory_config.""" + +from argparse import Namespace + +import pytest +import yaml + +from megatron.elastification.memory_config import MemoryConfig, load_memory_config + + +@pytest.fixture +def profiles_file(tmp_path): + """Write a minimal memory-profiles YAML and return its path.""" + data = { + "presets": { + "bf16": { + "params": 2, + "kv_cache": 2, + "ssm_cache": 2, + "max_buffer": 2, + "param_budget_target": "active", + }, + "fp8_kv": { + "params": 2, + "kv_cache": 1, + "ssm_cache": 2, + "max_buffer": 2, + "param_budget_target": "active", + }, + "total_target": { + "params": 2, + "kv_cache": 2, + "ssm_cache": 2, + "max_buffer": 2, + "param_budget_target": "total", + }, + } + } + path = tmp_path / "memory_profiles.yaml" + path.write_text(yaml.safe_dump(data)) + return str(path) + + +def _make_args(profile="bf16", profiles_path=None, **overrides): + defaults = dict( + memory_profile=profile, + memory_profile_path=profiles_path, + bpe_params=None, + bpe_kv_cache=None, + bpe_ssm_cache=None, + bpe_max_buffer=None, + param_budget_target=None, + ) + defaults.update(overrides) + return Namespace(**defaults) + + +class TestMemoryConfigDataclass: + def test_default_values(self): + cfg = MemoryConfig() + assert cfg.bpe_params == 2.0 + assert cfg.bpe_kv_cache == 2.0 + assert cfg.bpe_ssm_cache == 2.0 + assert cfg.bpe_max_buffer == 2.0 + assert cfg.param_budget_target == "active" + + def test_invalid_param_budget_target_rejected(self): + with pytest.raises(ValueError, match="param_budget_target"): + MemoryConfig(param_budget_target="bogus") + + def test_valid_param_budget_target_accepted(self): + MemoryConfig(param_budget_target="active") + MemoryConfig(param_budget_target="total") + + +class TestLoadMemoryConfig: + def test_preset_applied(self, profiles_file): + args = _make_args(profile="fp8_kv", profiles_path=profiles_file) + cfg = load_memory_config(args) + assert cfg.bpe_params == 2.0 + assert cfg.bpe_kv_cache == 1.0 # FP8 + assert cfg.bpe_ssm_cache == 2.0 + assert cfg.bpe_max_buffer == 2.0 + + def test_preset_param_budget_target(self, profiles_file): + args = _make_args(profile="total_target", profiles_path=profiles_file) + cfg = load_memory_config(args) + assert cfg.param_budget_target == "total" + + def test_cli_override_takes_priority_over_preset(self, profiles_file): + args = _make_args( + profile="bf16", profiles_path=profiles_file, bpe_kv_cache=0.5625 # override + ) + cfg = load_memory_config(args) + assert cfg.bpe_kv_cache == 0.5625 # override wins + assert cfg.bpe_params == 2.0 # preset preserved + + def test_param_budget_target_override(self, profiles_file): + args = _make_args(profile="bf16", profiles_path=profiles_file, param_budget_target="total") + cfg = load_memory_config(args) + assert cfg.param_budget_target == "total" + + def test_unknown_profile_raises(self, profiles_file): + args = _make_args(profile="nonexistent", profiles_path=profiles_file) + with pytest.raises(ValueError, match="not found"): + load_memory_config(args) + + def test_missing_profile_file_raises(self, tmp_path): + args = _make_args(profile="bf16", profiles_path=str(tmp_path / "missing.yaml")) + with pytest.raises(FileNotFoundError): + load_memory_config(args) + + def test_none_profile_name_defaults_to_bf16(self, profiles_file): + args = _make_args(profile=None, profiles_path=profiles_file) + cfg = load_memory_config(args) + # bf16 defaults in the fixture. + assert cfg.bpe_params == 2.0 + assert cfg.bpe_kv_cache == 2.0 + + def test_default_profiles_path_loads_bundled_yaml(self): + # When profiles_path is None, the loader falls back to the bundled + # megatron/elastification/memory_profiles.yaml. + args = _make_args(profile="bf16", profiles_path=None) + cfg = load_memory_config(args) + assert cfg.bpe_params == 2.0 + assert cfg.bpe_kv_cache == 2.0